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