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 } 226