From 438b4f87d3ca5fbea78f60d0b5ccb60356b20e90 Mon Sep 17 00:00:00 2001
From: Cameron Gutman <aicommander@gmail.com>
Date: Tue, 13 Aug 2019 17:51:11 -0700
Subject: [PATCH] Rewrite STUN code to avoid requiring TCP and use multiple
 servers in parallel

---
 src/SimpleStun.c | 62 +++++++++++++++++++++++++++++-------------------
 1 file changed, 37 insertions(+), 25 deletions(-)

diff --git a/src/SimpleStun.c b/src/SimpleStun.c
index f012448..8367f14 100644
--- a/src/SimpleStun.c
+++ b/src/SimpleStun.c
@@ -38,9 +38,9 @@ typedef struct _STUN_MESSAGE {
 int LiFindExternalAddressIP4(const char* stunServer, unsigned short stunPort, unsigned int* wanAddr)
 {
     SOCKET sock;
-    struct sockaddr_storage stunAddr;
-    struct sockaddr_in* stunAddrIn;
-    SOCKADDR_LEN stunAddrLen;
+    struct addrinfo* stunAddrs;
+    struct addrinfo hints;
+    char stunPortStr[6];
     int err;
     STUN_MESSAGE reqMsg;
     int i;
@@ -58,13 +58,21 @@ int LiFindExternalAddressIP4(const char* stunServer, unsigned short stunPort, un
         return err;
     }
 
-    err = resolveHostName(stunServer, AF_INET, 3478, &stunAddr, &stunAddrLen);
-    if (err != 0) {
+    memset(&hints, 0, sizeof(hints));
+    hints.ai_family = AF_INET;
+    hints.ai_socktype = SOCK_DGRAM;
+    hints.ai_protocol = IPPROTO_UDP;
+    hints.ai_flags = AI_ADDRCONFIG;
+
+    sprintf(stunPortStr, "%u", stunPort);
+    err = getaddrinfo(stunServer, stunPortStr, &hints, &stunAddrs);
+    if (err != 0 || stunAddrs == NULL) {
         Limelog("Failed to resolve STUN server: %d\n", err);
+        stunAddrs = NULL;
         goto Exit;
     }
 
-    sock = bindUdpSocket(AF_INET, 2048);
+    sock = bindUdpSocket(hints.ai_family, 2048);
     if (sock == INVALID_SOCKET) {
         err = LastSocketFail();
         Limelog("Failed to connect to STUN server: %d\n", err);
@@ -79,29 +87,29 @@ int LiFindExternalAddressIP4(const char* stunServer, unsigned short stunPort, un
     }
 
     bytesRead = SOCKET_ERROR;
-    for (i = 0; i < STUN_RECV_TIMEOUT_SEC * 1000 / UDP_RECV_POLL_TIMEOUT_MS; i++) {
-        // Retransmit the request every second until the timeout elapses
+    for (i = 0; i < STUN_RECV_TIMEOUT_SEC * 1000 / UDP_RECV_POLL_TIMEOUT_MS && bytesRead <= 0; i++) {
+        // Retransmit the request every second until the timeout elapses or we get a response
         if (i % (1000 / UDP_RECV_POLL_TIMEOUT_MS) == 0) {
-            stunAddrIn = (struct sockaddr_in*)&stunAddr;
-            stunAddrIn->sin_port = htons(stunPort);
-            err = (int)sendto(sock, (char *)&reqMsg, sizeof(reqMsg), 0,
-                               (struct sockaddr*)stunAddrIn, stunAddrLen);
-            if (err == SOCKET_ERROR) {
-                err = LastSocketFail();
-                Limelog("Failed to send STUN binding request: %d\n", err);
-                closeSocket(sock);
-                goto Exit;
+            struct addrinfo *current;
+
+            // Send a request to each resolved address but stop if we get a response
+            for (current = stunAddrs; current != NULL && bytesRead <= 0; current = current->ai_next) {
+                err = (int)sendto(sock, (char *)&reqMsg, sizeof(reqMsg), 0, current->ai_addr, current->ai_addrlen);
+                if (err == SOCKET_ERROR) {
+                    err = LastSocketFail();
+                    Limelog("Failed to send STUN binding request: %d\n", err);
+                    continue;
+                }
+
+                // Wait UDP_RECV_POLL_TIMEOUT_MS before moving on to the next server to
+                // avoid having to spam the other STUN servers if we find a working one.
+                bytesRead = recvUdpSocket(sock, resp.buf, sizeof(resp.buf), 1);
             }
         }
-
-        // This waits in UDP_RECV_POLL_TIMEOUT_MS increments
-        bytesRead = recvUdpSocket(sock, resp.buf, sizeof(resp.buf), 1);
-        if (bytesRead == 0) {
-            // Timeout - continue looping
-            continue;
+        else {
+            // This waits in UDP_RECV_POLL_TIMEOUT_MS increments
+            bytesRead = recvUdpSocket(sock, resp.buf, sizeof(resp.buf), 1);
         }
-
-        break;
     }
 
     closeSocket(sock);
@@ -176,6 +184,10 @@ int LiFindExternalAddressIP4(const char* stunServer, unsigned short stunPort, un
     err = -6;
 
 Exit:
+    if (stunAddrs != NULL) {
+        freeaddrinfo(stunAddrs);
+    }
+
     cleanupPlatformSockets();
     return err;
 }