Condense getservinfo and bounds check

1 files changed, 34 insertions, 59 deletions

diff --git a/getservinfo.c b/getservinfo.c
index 0db7bc0..219f72a 100644
--- a/getservinfo.c
+++ b/getservinfo.c
@@ -41,34 +41,12 @@
 #define EAI_PROTOCOL EAI_BADFLAGS
 #endif
 
-static uint16_t u16(uint8_t **ptr) {
-	uint16_t x = *(*ptr)++;
-	x = x << 8 | *(*ptr)++;
-	return x;
-}
-
-static void nameSkip(uint8_t **ptr) {
-	for (uint8_t len; (len = *(*ptr)++); *ptr += len) {
-		if (len & 0xC0) {
-			(*ptr)++;
-			break;
-		}
-	}
-}
-
-static char *nameString(uint8_t **ptr) {
-	char *name = (char *)(*ptr + 1);
-	for (uint8_t len; (len = **ptr); *ptr += len) {
-		*(*ptr)++ = '.';
-	}
-	return name;
-}
-
 /* A wrapper around getaddrinfo(3) which first performs SRV record (RFC 2782)
  * lookup. hints must be provided and hints->ai_protocol must be set. SRV
  * lookup is skipped if servname is numerical. If SRV lookup is successful, the
  * ai_canonname field of the first addrinfo structure returned is set to the
- * target name.
+ * target name. Only the first SRV record is used. Priority and weight are
+ * ignored.
  */
 int getservinfo(
 	const char *hostname, const char *servname,
@@ -76,12 +54,12 @@ int getservinfo(
 ) {
 	if (!hints) return EAI_BADHINTS;
 	if (!hints->ai_protocol) return EAI_PROTOCOL;
-	if (hints->ai_flags & AI_NUMERICHOST) return EAI_BADFLAGS;
-	if (hints->ai_flags & AI_NUMERICSERV) return EAI_BADFLAGS;
 
 	char *rest;
 	strtoul(servname, &rest, 10);
-	if (!*rest) return getaddrinfo(hostname, servname, hints, res);
+	if (!*rest || hints->ai_flags & (AI_NUMERICHOST | AI_NUMERICSERV)) {
+		return getaddrinfo(hostname, servname, hints, res);
+	}
 
 	struct protoent *proto = getprotobynumber(hints->ai_protocol);
 	if (!proto) return EAI_PROTOCOL;
@@ -95,47 +73,44 @@ int getservinfo(
 
 	uint8_t msg[512];
 	len = res_query(dname, 1 /* IN */, 33 /* SRV */, msg, sizeof(msg));
-	if (len < 0) return getaddrinfo(hostname, servname, hints, res);
-
-	uint8_t *ptr = msg;
-	u16(&ptr); // ID
-	uint16_t rcode = u16(&ptr) & 0x000F;
-	uint16_t qdcount = u16(&ptr);
-	uint16_t ancount = u16(&ptr);
-	u16(&ptr); // NSCOUNT
-	u16(&ptr); // ARCOUNT
-
-	if (rcode || !ancount) return getaddrinfo(hostname, servname, hints, res);
+	if (len < 12) return getaddrinfo(hostname, servname, hints, res);
+	uint8_t *ptr = &msg[12];
+
+#define NAME_SKIP \
+	for (uint8_t n; ptr < &msg[len] && (n = *ptr++); ptr += n) { \
+		if (n & 0xC0) { \
+			ptr++; \
+			break; \
+		} \
+	}
 
-	for (uint16_t q = 0; q < qdcount; ++q) {
-		nameSkip(&ptr); // QNAME
-		u16(&ptr); // QTYPE
-		u16(&ptr); // QCLASS
+	uint16_t qdcount = msg[4] << 8 | msg[5];
+	for (uint16_t q = 0; ptr < &msg[len] && q < qdcount; ++q) {
+		NAME_SKIP; // QNAME
+		ptr += 4; // QTYPE, QCLASS
 	}
 
-	// Read only one answer, ignoring Priority and Weight. Does anyone actually
-	// use them?
-	nameSkip(&ptr); // NAME
-	u16(&ptr); // TYPE
-	u16(&ptr); // CLASS
-	u16(&ptr); // TTL
-	u16(&ptr); // TTL
-	u16(&ptr); // RDLENGTH
-	u16(&ptr); // Priority
-	u16(&ptr); // Weight
+	NAME_SKIP; // NAME
+	ptr += 14; // TYPE, CLASS, TTL, RDLENGTH, Priority, Weight
+	if (&msg[len] < ptr + 3) {
+		return getaddrinfo(hostname, servname, hints, res);
+	}
 
-	uint16_t port = u16(&ptr);
-	hostname = nameString(&ptr);
-	if (!hostname[0]) return EAI_NONAME;
+	char port[sizeof("65535")];
+	snprintf(port, sizeof(port), "%d", ptr[0] << 8 | ptr[1]);
+	ptr += 2;
 
-	char myServ[sizeof("65535")];
-	snprintf(myServ, sizeof(myServ), "%hu", port);
+	// Name compression is not used for Target.
+	if (!ptr[0]) return EAI_NONAME;
+	hostname = (const char *)&ptr[1];
+	for (uint8_t n; ptr < &msg[len] && (n = *ptr); ptr += n) {
+		*ptr++ = '.';
+	}
 
 	struct addrinfo myHints = *hints;
 	myHints.ai_flags |= AI_NUMERICSERV;
 	myHints.ai_flags &= ~AI_CANONNAME;
-
-	int error = getaddrinfo(hostname, myServ, &myHints, res);
+	int error = getaddrinfo(hostname, port, &myHints, res);
 	if (error) return error;
 
 	(*res)->ai_canonname = strdup(hostname);