From 665309986e781cf734557be7bca78687a262ad9e Mon Sep 17 00:00:00 2001
From: Christopher Haster <chaster@utexas.edu>
Date: Tue, 9 Aug 2016 14:53:00 -0500
Subject: [PATCH] [nsapi] Added support for multiple results in a dns-query

Internal API changes:
+ nsapi_dns_query_multiple

Note, dns_query_multiple takes a buffer of nsapi_addr_t to fill. This is
less dynamic than the linked-list method used in linux's gethostbyname,
but is easier to manage and more constrained.
---
 features/net/network-socket/nsapi_dns.cpp | 57 ++++++++++++++++++-----
 1 file changed, 45 insertions(+), 12 deletions(-)

diff --git a/features/net/network-socket/nsapi_dns.cpp b/features/net/network-socket/nsapi_dns.cpp
index 6c781f36b0..cf2744ea07 100644
--- a/features/net/network-socket/nsapi_dns.cpp
+++ b/features/net/network-socket/nsapi_dns.cpp
@@ -17,6 +17,7 @@
 #include "nsapi_dns.h"
 #include "network-socket/UDPSocket.h"
 #include <string.h>
+#include <stdlib.h>
 
 
 #define DNS_BUFFER_SIZE 256 
@@ -88,7 +89,7 @@ static void dns_append_question(uint8_t **p, const char *host)
     dns_append_word(p, 1); // qclass = 1
 }
 
-static bool dns_scan_response(const uint8_t **p, nsapi_addr_t *addr)
+static int dns_scan_response(const uint8_t **p, nsapi_addr_t *addr, unsigned addr_count)
 {
     // scan header
     uint16_t id    = dns_scan_word(p);
@@ -104,7 +105,7 @@ static bool dns_scan_response(const uint8_t **p, nsapi_addr_t *addr)
 
     // verify header is response to query
     if (!(id == 1 && qr && opcode == 0 && rcode == 0)) {
-        return false;
+        return 0;
     }
 
     // skip questions
@@ -123,7 +124,9 @@ static bool dns_scan_response(const uint8_t **p, nsapi_addr_t *addr)
     }
 
     // scan each response
-    for (int i = 0; i < ancount; i++) {
+    unsigned count = 0;
+
+    for (int i = 0; i < ancount && count < addr_count; i++) {
         while (true) {
             uint8_t len = dns_scan_byte(p);
             if (len == 0) {
@@ -152,14 +155,15 @@ static bool dns_scan_response(const uint8_t **p, nsapi_addr_t *addr)
             addr->bytes[i] = dns_scan_byte(p);
         }
 
-        return true;
+        addr += 1;
+        count += 1;
     }
 
-    return false;
+    return count;
 }
 
-
-int nsapi_dns_query(NetworkStack *stack, nsapi_addr_t *addr, const char *host)
+// core query function
+int nsapi_dns_query_multiple(NetworkStack *stack, nsapi_addr_t *addr, unsigned addr_count, const char *host)
 {
     // check for valid host name
     int host_len = host ? strlen(host) : 0;
@@ -208,9 +212,9 @@ int nsapi_dns_query(NetworkStack *stack, nsapi_addr_t *addr, const char *host)
         }
 
         p = packet;
-        bool found = dns_scan_response((const uint8_t **)&p, addr);
+        int found = dns_scan_response((const uint8_t **)&p, addr, addr_count);
         if (found) {
-            result = 0;
+            result = found;
             break;
         }
     }
@@ -228,16 +232,45 @@ int nsapi_dns_query(NetworkStack *stack, nsapi_addr_t *addr, const char *host)
     return result;
 }
 
+// convenience functions for other forms of queries
+NSAPI_C_LINKAGE
+int nsapi_dns_query_multiple(nsapi_stack_t *stack, nsapi_addr_t *addr, unsigned addr_count, const char *host)
+{
+    return nsapi_dns_query_multiple(nsapi_create_stack(stack), addr, addr_count, host);
+}
+
+int nsapi_dns_query_multiple(NetworkStack *stack, SocketAddress *addresses, unsigned addr_count, const char *host)
+{
+    nsapi_addr_t *addrs = new nsapi_addr_t[addr_count];
+    int result = nsapi_dns_query_multiple(stack, addrs, addr_count, host);
+
+    if (result > 0) {
+        for (int i = 0; i < result; i++) {
+            addresses[i].set_addr(addrs[i]);
+        }
+    }
+
+    delete[] addrs;
+    return result;
+}
+
 NSAPI_C_LINKAGE
 int nsapi_dns_query(nsapi_stack_t *stack, nsapi_addr_t *addr, const char *host)
 {
-    return nsapi_dns_query(nsapi_create_stack(stack), addr, host);
+    int result = nsapi_dns_query_multiple(nsapi_create_stack(stack), addr, 1, host);
+    return (result > 0) ? 0 : result;
+}
+
+int nsapi_dns_query(NetworkStack *stack, nsapi_addr_t *addr, const char *host)
+{
+    int result = nsapi_dns_query_multiple(stack, addr, 1, host);
+    return (result > 0) ? 0 : result;
 }
 
 int nsapi_dns_query(NetworkStack *stack, SocketAddress *address, const char *host)
 {
     nsapi_addr_t addr;
-    int err = nsapi_dns_query(stack, &addr, host);
+    int result = nsapi_dns_query_multiple(stack, &addr, 1, host);
     address->set_addr(addr);
-    return err;
+    return (result > 0) ? 0 : result;
 }