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         chosen = malloc(sizeof(int) * (last - first + 1));
  88         totalweight = 0;
  89
  90         for (i = 0; i <= last; i++)
  91                 totalweight += d[i]->weight;
  92
  93         if (totalweight == 0)
  94                 return;
  95
  96         for (i = 0; i <= last; i++) {
  97                 for (;;) {
  98                         chosen[i] = random() % (d[i]->weight * 100 / totalweight);
  99                         for (j = 0; j < i; j++) {
 100                                 if (chosen[i] == chosen[j])
 101                                         break;
 102                         }
 103                         if (j == i) {
 104                                 d[i]->finalweight = chosen[i];
 105                                 break;
 106                         }
 107                 }
 108         }
 109
 110         free(chosen);
 111 }
 112
 113 struct dns_srvinfo *
 114 dns_getsrvinfo(const char *zone)
 115 {
 116         struct dns_srvinfo **res, *first;
 117         unsigned char *end, *p;
 118         char host[MAXHOSTNAMELEN];
 119         dns_query q;
 120         int len, qdcount, ancount, n, i, f, l;
 121         unsigned int type, class, ttl, priority, weight, port;
 122
 123         if ((len = res_query(zone, C_IN, T_SRV, q.buf, sizeof(q.buf))) == -1 ||
 124             len < (int)sizeof(HEADER))
 125                 return (NULL);
 126
 127         qdcount = ntohs(q.hdr.qdcount);
 128         ancount = ntohs(q.hdr.ancount);
 129
 130         end = q.buf + len;
 131         p = q.buf + sizeof(HEADER);
 132
 133         while(qdcount > 0 && p < end) {
 134                 qdcount--;
 135                 if((len = dn_expand(q.buf, end, p, host, MAXHOSTNAMELEN)) < 0)
 136                         return (NULL);
 137                 p += len + NS_QFIXEDSZ;
 138         }
 139
 140         res = calloc(ancount, sizeof(struct dns_srvinfo *));
 141         if (res == NULL)
 142                 return (NULL);
 143
 144         n = 0;
 145         while (ancount > 0 && p < end) {
 146                 ancount--;
 147                 len = dn_expand(q.buf, end, p, host, MAXHOSTNAMELEN);
 148                 if (len < 0) {
 149                         for (i = 0; i < n; i++)
 150                                 free(res[i]);
 151                         free(res);
 152                         return NULL;
 153                 }
 154
 155                 p += len;
 156
 157                 NS_GET16(type, p);
 158                 NS_GET16(class, p);
 159                 NS_GET32(ttl, p);
 160                 NS_GET16(len, p);
 161
 162                 if (type != T_SRV) {
 163                         p += len;
 164                         continue;
 165                 }
 166
 167                 NS_GET16(priority, p);
 168                 NS_GET16(weight, p);
 169                 NS_GET16(port, p);
 170
 171                 len = dn_expand(q.buf, end, p, host, MAXHOSTNAMELEN);
 172                 if (len < 0) {
 173                         for (i = 0; i < n; i++)
 174                                 free(res[i]);
 175                         free(res);
 176                         return (NULL);
 177                 }
 178
 179                 res[n] = malloc(sizeof(struct dns_srvinfo));
 180                 if (res[n] == NULL) {
 181                         for (i = 0; i < n; i++)
 182                                 free(res[i]);
 183                         free(res);
 184                         return (NULL);
 185                 }
 186                 res[n]->type = type;
 187                 res[n]->class = class;
 188                 res[n]->ttl = ttl;
 189                 res[n]->priority = priority;
 190                 res[n]->weight = weight;
 191                 res[n]->port = port;
 192                 res[n]->next = NULL;
 193                 strlcpy(res[n]->host, host, MAXHOSTNAMELEN);
 194
 195                 p += len;
 196                 n++;
 197         }
 198
 199         qsort(res, n, sizeof(res[0]), srv_priority_cmp);
 200
 201         priority = f = l = 0;
 202         for (i = 0; i < n; i++) {
 203                 if (res[i]->priority != priority) {
 204                         if (f != l)
 205                                 compute_weight(res, f, l);
 206                         f = i;
 207                         priority = res[i]->priority;
 208                 }
 209                 l = i;
 210         }
 211
 212         qsort(res, n, sizeof(res[0]), srv_final_cmp);
 213
 214         for (i = 0; i < n - 1; i++)
 215                 res[i]->next = res[i + 1];
 216
 217         first = res[0];
 218         free(res);
 219
 220         return (first);
 221 }