usr.sbin/pkg/dns_utils.c

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