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