usr.sbin/pkg/dns_utils.c

   1 /*-
   2  * Copyright (c) 2012-2013 Baptiste Daroussin <bapt@FreeBSD.org>
   3  * All rights reserved.
   4  *
   5  * Redistribution and use in source and binary forms, with or without
   6  * modification, are permitted provided that the following conditions
   7  * are met:
   8  * 1. Redistributions of source code must retain the above copyright
   9  *    notice, this list of conditions and the following disclaimer,
  10  *    without modification, immediately at the beginning of the file.
  11  * 2. Redistributions in binary form must reproduce the above copyright
  12  *    notice, this list of conditions and the following disclaimer in the
  13  *    documentation and/or other materials provided with the distribution.
  14  *
  15  * THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
  16  * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
  17  * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
  18  * IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT,
  19  * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
  20  * NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
  21  * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
  22  * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
  23  * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
  24  * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  25  */
  26
  27 #include <sys/cdefs.h>
  28 __FBSDID("$FreeBSD$");
  29
  30 #include <stdlib.h>
  31 #include <string.h>
  32 #include <netinet/in.h>
  33 #include <resolv.h>
  34
  35 #include "dns_utils.h"
  36
  37 typedef union {
  38         HEADER hdr;
  39         unsigned char buf[1024];
  40 } dns_query;
  41
  42 static int
  43 srv_priority_cmp(const void *a, const void *b)
  44 {
  45         const struct dns_srvinfo *da, *db;
  46         unsigned int r, l;
  47
  48         da = *(struct dns_srvinfo * const *)a;
  49         db = *(struct dns_srvinfo * const *)b;
  50
  51         l = da->priority;
  52         r = db->priority;
  53
  54         return ((l > r) - (l < r));
  55 }
  56
  57 static int
  58 srv_final_cmp(const void *a, const void *b)
  59 {
  60         const struct dns_srvinfo *da, *db;
  61         unsigned int r, l, wr, wl;
  62         int res;
  63
  64         da = *(struct dns_srvinfo * const *)a;
  65         db = *(struct dns_srvinfo * const *)b;
  66
  67         l = da->priority;
  68         r = db->priority;
  69
  70         res = ((l > r) - (l < r));
  71
  72         if (res == 0) {
  73                 wl = da->finalweight;
  74                 wr = db->finalweight;
  75                 res = ((wr > wl) - (wr < wl));
  76         }
  77
  78         return (res);
  79 }
  80
  81 static void
  82 compute_weight(struct dns_srvinfo **d, int first, int last)
  83 {
  84         int i, j, totalweight;
  85         int *chosen;
  86
  87         totalweight = 0;
  88
  89         for (i = 0; i <= last; i++)
  90                 totalweight += d[i]->weight;
  91
  92         if (totalweight == 0)
  93                 return;
  94
  95         chosen = malloc(sizeof(int) * (last - first + 1));
  96
  97         for (i = 0; i <= last; i++) {
  98                 for (;;) {
  99                         chosen[i] = random() % (d[i]->weight * 100 / totalweight);
 100                         for (j = 0; j < i; j++) {
 101                                 if (chosen[i] == chosen[j])
 102                                         break;
 103                         }
 104                         if (j == i) {
 105                                 d[i]->finalweight = chosen[i];
 106                                 break;
 107                         }
 108                 }
 109         }
 110
 111         free(chosen);
 112 }
 113
 114 struct dns_srvinfo *
 115 dns_getsrvinfo(const char *zone)
 116 {
 117         struct dns_srvinfo **res, *first;
 118         unsigned char *end, *p;
 119         char host[MAXHOSTNAMELEN];
 120         dns_query q;
 121         int len, qdcount, ancount, n, i, f, l;
 122         unsigned int type, class, ttl, priority, weight, port;
 123
 124         if ((len = res_query(zone, C_IN, T_SRV, q.buf, sizeof(q.buf))) == -1 ||
 125             len < (int)sizeof(HEADER))
 126                 return (NULL);
 127
 128         qdcount = ntohs(q.hdr.qdcount);
 129         ancount = ntohs(q.hdr.ancount);
 130
 131         end = q.buf + len;
 132         p = q.buf + sizeof(HEADER);
 133
 134         while(qdcount > 0 && p < end) {
 135                 qdcount--;
 136                 if((len = dn_expand(q.buf, end, p, host, MAXHOSTNAMELEN)) < 0)
 137                         return (NULL);
 138                 p += len + NS_QFIXEDSZ;
 139         }
 140
 141         res = calloc(ancount, sizeof(struct dns_srvinfo *));
 142         if (res == NULL)
 143                 return (NULL);
 144
 145         n = 0;
 146         while (ancount > 0 && p < end) {
 147                 ancount--;
 148                 len = dn_expand(q.buf, end, p, host, MAXHOSTNAMELEN);
 149                 if (len < 0) {
 150                         for (i = 0; i < n; i++)
 151                                 free(res[i]);
 152                         free(res);
 153                         return NULL;
 154                 }
 155
 156                 p += len;
 157
 158                 NS_GET16(type, p);
 159                 NS_GET16(class, p);
 160                 NS_GET32(ttl, p);
 161                 NS_GET16(len, p);
 162
 163                 if (type != T_SRV) {
 164                         p += len;
 165                         continue;
 166                 }
 167
 168                 NS_GET16(priority, p);
 169                 NS_GET16(weight, p);
 170                 NS_GET16(port, p);
 171
 172                 len = dn_expand(q.buf, end, p, host, MAXHOSTNAMELEN);
 173                 if (len < 0) {
 174                         for (i = 0; i < n; i++)
 175                                 free(res[i]);
 176                         free(res);
 177                         return (NULL);
 178                 }
 179
 180                 res[n] = malloc(sizeof(struct dns_srvinfo));
 181                 if (res[n] == NULL) {
 182                         for (i = 0; i < n; i++)
 183                                 free(res[i]);
 184                         free(res);
 185                         return (NULL);
 186                 }
 187                 res[n]->type = type;
 188                 res[n]->class = class;
 189                 res[n]->ttl = ttl;
 190                 res[n]->priority = priority;
 191                 res[n]->weight = weight;
 192                 res[n]->port = port;
 193                 res[n]->next = NULL;
 194                 strlcpy(res[n]->host, host, MAXHOSTNAMELEN);
 195
 196                 p += len;
 197                 n++;
 198         }
 199
 200         qsort(res, n, sizeof(res[0]), srv_priority_cmp);
 201
 202         priority = f = l = 0;
 203         for (i = 0; i < n; i++) {
 204                 if (res[i]->priority != priority) {
 205                         if (f != l)
 206                                 compute_weight(res, f, l);
 207                         f = i;
 208                         priority = res[i]->priority;
 209                 }
 210                 l = i;
 211         }
 212
 213         qsort(res, n, sizeof(res[0]), srv_final_cmp);
 214
 215         for (i = 0; i < n - 1; i++)
 216                 res[i]->next = res[i + 1];
 217
 218         first = res[0];
 219         free(res);
 220
 221         return (first);
 222 }