Merge pull request #2135 from wh201906/udp_win

Add UDP support on Windows

CHANGELOG.md

@@ -3,6 +3,7 @@ All notable changes to this project will be documented in this file.
 This project uses the changelog in accordance with [keepchangelog](http://keepachangelog.com/). Please use this to write notable changes, which is not the same as git commit log...
 
 ## [unreleased][unreleased]
+ - Added UDP support on Windows (@wh201906)
  - Added client communication timeout to preferences (@iceman1001)
  - Added IPv6 support (@wh201906)
  - Fixed `lf hid clone --bin` - now correctly handles sentinel bits (@iceman1001)

client/src/uart/uart_win32.c

@@ -17,6 +17,7 @@
 //-----------------------------------------------------------------------------
 
 #include "uart.h"
+#include "ringbuffer.h"
 
 #include <stdio.h>
 #include <string.h>
@@ -37,6 +38,7 @@ typedef struct {
     DCB dcb;               // Device control settings
     COMMTIMEOUTS ct;       // Serial port time-out configuration
     SOCKET hSocket;        // Socket handle
+    RingBuffer* udpBuffer; // Buffer for UDP
 } serial_port_windows_t;
 
 // this is for TCP connection
@@ -63,8 +65,7 @@ static int uart_reconfigure_timeouts_polling(serial_port sp) {
         return PM3_SUCCESS;
     newtimeout_pending = false;
 
-    serial_port_windows_t *spw;
+    serial_port_windows_t *spw = (serial_port_windows_t *)sp;
-    spw = (serial_port_windows_t *)sp;
     spw->ct.ReadIntervalTimeout         = newtimeout_value;
     spw->ct.ReadTotalTimeoutMultiplier  = 0;
     spw->ct.ReadTotalTimeoutConstant    = newtimeout_value;
@@ -90,6 +91,8 @@ serial_port uart_open(const char *pcPortName, uint32_t speed) {
         return INVALID_SERIAL_PORT;
     }
 
+    sp->udpBuffer = NULL;
+
     char *prefix = strdup(pcPortName);
     if (prefix == NULL) {
         PrintAndLogEx(ERR, "error:  string duplication");
@@ -225,6 +228,126 @@ serial_port uart_open(const char *pcPortName, uint32_t speed) {
         return sp;
     }
 
+    if (memcmp(prefix, "udp:", 4) == 0) {
+        free(prefix);
+
+        if (strlen(pcPortName) <= 4) {
+            PrintAndLogEx(ERR, "error: tcp port name length too short");
+            free(sp);
+            return INVALID_SERIAL_PORT;
+        }
+
+        struct addrinfo *addr = NULL, *rp;
+
+        char *addrPortStr = strdup(pcPortName + 4);
+        char *addrstr = addrPortStr;
+        const char *portstr;
+        if (addrPortStr == NULL) {
+            PrintAndLogEx(ERR, "error: string duplication");
+            free(sp);
+            return INVALID_SERIAL_PORT;
+        }
+
+        timeout.tv_usec = UART_TCP_CLIENT_RX_TIMEOUT_MS * 1000;
+
+        // find the start of the address
+        char *endBracket = strrchr(addrPortStr, ']');
+        if (addrPortStr[0] == '[') {
+            addrstr += 1;
+            if (endBracket == NULL) {
+                PrintAndLogEx(ERR, "error: wrong address: [] unmatched");
+                free(addrPortStr);
+                free(sp);
+                return INVALID_SERIAL_PORT;
+            }
+        }
+
+        // find the port
+        char *lColon = strchr(addrPortStr, ':');
+        char *rColon = strrchr(addrPortStr, ':');
+        if (rColon == NULL) {
+            // no colon
+            // "tcp:<ipv4 address>", "tcp:[<ipv4 address>]"
+            portstr = "18888";
+        } else if (lColon == rColon) {
+            // only one colon
+            // "tcp:<ipv4 address>:<port>", "tcp:[<ipv4 address>]:<port>"
+            portstr = rColon + 1;
+        } else {
+            // two or more colon, IPv6 address
+            // tcp:[<ipv6 address>]:<port>
+            // "tcp:<ipv6 address>", "tcp:[<ipv6 address>]"
+            if (endBracket != NULL && rColon == endBracket + 1) {
+                portstr = rColon + 1;
+            } else {
+                portstr = "18888";
+            }
+        }
+
+        // handle the end of the address
+        if (endBracket != NULL) {
+            *endBracket = '\0';
+        } else if (rColon != NULL && lColon == rColon) {
+            *rColon = '\0';
+        }
+
+        WSADATA wsaData;
+        struct addrinfo info;
+        int iResult;
+
+        iResult = WSAStartup(MAKEWORD(2, 2), &wsaData);
+        if (iResult != 0) {
+            PrintAndLogEx(ERR, "error: WSAStartup failed with error: %d", iResult);
+            free(addrPortStr);
+            free(sp);
+            return INVALID_SERIAL_PORT;
+        }
+
+        memset(&info, 0, sizeof(info));
+        info.ai_family = AF_UNSPEC;
+        info.ai_socktype = SOCK_DGRAM;
+        info.ai_protocol = IPPROTO_UDP;
+
+        int s = getaddrinfo(addrstr, portstr, &info, &addr);
+        if (s != 0) {
+            PrintAndLogEx(ERR, "error: getaddrinfo: %d: %s", s, gai_strerror(s));
+            freeaddrinfo(addr);
+            free(addrPortStr);
+            free(sp);
+            WSACleanup();
+            return INVALID_SERIAL_PORT;
+        }
+
+        SOCKET hSocket = INVALID_SOCKET;
+        for (rp = addr; rp != NULL; rp = rp->ai_next) {
+            hSocket = socket(rp->ai_family, rp->ai_socktype, rp->ai_protocol);
+
+            if (hSocket == INVALID_SOCKET)
+                continue;
+
+            if (connect(hSocket, rp->ai_addr, (int)rp->ai_addrlen) != INVALID_SOCKET)
+                break;
+
+            closesocket(hSocket);
+            hSocket = INVALID_SOCKET;
+        }
+
+        freeaddrinfo(addr);
+        free(addrPortStr);
+
+        if (rp == NULL) {               /* No address succeeded */
+            PrintAndLogEx(ERR, "error: Could not connect");
+            WSACleanup();
+            free(sp);
+            return INVALID_SERIAL_PORT;
+        }
+
+        sp->hSocket = hSocket;
+        sp->udpBuffer = RingBuf_create(MAX(sizeof(PacketResponseNGRaw), sizeof(PacketResponseOLD)) * 30);
+
+        return sp;
+    }
+
     // Copy the input "com?" to "\\.\COM?" format
     snprintf(acPortName, sizeof(acPortName), "\\\\.\\%s", pcPortName);
     _strupr(acPortName);
@@ -277,13 +400,14 @@ void uart_close(const serial_port sp) {
         closesocket(spw->hSocket);
         WSACleanup();
     }
+    RingBuf_destroy(spw->udpBuffer);
     if (spw->hPort != INVALID_HANDLE_VALUE)
         CloseHandle(spw->hPort);
     free(sp);
 }
 
 bool uart_set_speed(serial_port sp, const uint32_t uiPortSpeed) {
-    serial_port_windows_t *spw;
+    serial_port_windows_t *spw = (serial_port_windows_t *)sp;
 
     // Set port speed (Input and Output)
     switch (uiPortSpeed) {
@@ -301,7 +425,6 @@ bool uart_set_speed(serial_port sp, const uint32_t uiPortSpeed) {
             return false;
     };
 
-    spw = (serial_port_windows_t *)sp;
     spw->dcb.BaudRate = uiPortSpeed;
     bool result = SetCommState(spw->hPort, &spw->dcb);
     PurgeComm(spw->hPort, PURGE_RXABORT | PURGE_RXCLEAR);
@@ -320,11 +443,12 @@ uint32_t uart_get_speed(const serial_port sp) {
 }
 
 int uart_receive(const serial_port sp, uint8_t *pbtRx, uint32_t pszMaxRxLen, uint32_t *pszRxLen) {
-    serial_port_windows_t *spw = (serial_port_windows_t *)sp;
+    const serial_port_windows_t *spw = (serial_port_windows_t *)sp;
-    if (spw->hSocket == INVALID_SOCKET) { // serial port
+    if (spw->hSocket == INVALID_SOCKET) {
+        // serial port
         uart_reconfigure_timeouts_polling(sp);
 
-        int res = ReadFile(((serial_port_windows_t *)sp)->hPort, pbtRx, pszMaxRxLen, (LPDWORD)pszRxLen, NULL);
+        int res = ReadFile(spw->hPort, pbtRx, pszMaxRxLen, (LPDWORD)pszRxLen, NULL);
         if (res)
             return PM3_SUCCESS;
 
@@ -335,7 +459,8 @@ int uart_receive(const serial_port sp, uint8_t *pbtRx, uint32_t pszMaxRxLen, uin
         }
 
         return PM3_ENOTTY;
-    } else { // TCP
+    } else {
+        // TCP or UDP
         uint32_t byteCount;  // FIONREAD returns size on 32b
         fd_set rfds;
         struct timeval tv;
@@ -347,12 +472,31 @@ int uart_receive(const serial_port sp, uint8_t *pbtRx, uint32_t pszMaxRxLen, uin
         // Reset the output count
         *pszRxLen = 0;
         do {
+            int res;
+            if(spw->udpBuffer != NULL) {
+                // for UDP connection, try to use the data from the buffer
+
+                byteCount = RingBuf_getAvailableSize(spw->udpBuffer);
+                // Cap the number of bytes, so we don't overrun the buffer
+                if (pszMaxRxLen - (*pszRxLen) < byteCount) {
+                    // PrintAndLogEx(ERR, "UART:: RX prevent overrun (have %u, need %u)", pszMaxRxLen - (*pszRxLen), byteCount);
+                    byteCount = pszMaxRxLen - (*pszRxLen);
+                }
+                res = RingBuf_dequeueBatch(spw->udpBuffer, pbtRx + (*pszRxLen), byteCount);
+                *pszRxLen += res;
+
+                if (*pszRxLen == pszMaxRxLen) {
+                    // We have all the data we wanted.
+                    return PM3_SUCCESS;
+                }
+            }
+
             // Reset file descriptor
             FD_ZERO(&rfds);
             FD_SET(spw->hSocket, &rfds);
             tv = timeout;
             // the first argument nfds is ignored in Windows
-            int res = select(0, &rfds, NULL, NULL, &tv);
+            res = select(0, &rfds, NULL, NULL, &tv);
 
             // Read error
             if (res == SOCKET_ERROR) {
@@ -375,6 +519,27 @@ int uart_receive(const serial_port sp, uint8_t *pbtRx, uint32_t pszMaxRxLen, uin
             // PrintAndLogEx(ERR, "UART:: RX ioctl res %d byteCount %u", res, byteCount);
             if (res == SOCKET_ERROR) return PM3_ENOTTY;
 
+            // For UDP connection, put the incoming data into the buffer and handle them in the next round
+            if (spw->udpBuffer != NULL) {
+                if (RingBuf_getContinousAvailableSize(spw->udpBuffer) >= byteCount) {
+                    // write to the buffer directly
+                    res = recv(spw->hSocket, (char *)RingBuf_getRearPtr(spw->udpBuffer), RingBuf_getAvailableSize(spw->udpBuffer), 0);
+                    if (res >= 0) {
+                        RingBuf_postEnqueueBatch(spw->udpBuffer, res);
+                    }
+                } else {
+                    // use transit buffer
+                    uint8_t transitBuf[MAX(sizeof(PacketResponseNGRaw), sizeof(PacketResponseOLD)) * 30];
+                    res = recv(spw->hSocket, (char *)transitBuf, RingBuf_getAvailableSize(spw->udpBuffer), 0);
+                    RingBuf_enqueueBatch(spw->udpBuffer, transitBuf, res);
+                }
+                // Stop if the OS has some troubles reading the data
+                if (res < 0) {
+                    return PM3_EIO;
+                }
+                continue;
+            }
+
             // Cap the number of bytes, so we don't overrun the buffer
             if (pszMaxRxLen - (*pszRxLen) < byteCount) {
                 //            PrintAndLogEx(ERR, "UART:: RX prevent overrun (have %u, need %u)", pszMaxRxLen - (*pszRxLen), byteCount);
@@ -402,10 +567,10 @@ int uart_receive(const serial_port sp, uint8_t *pbtRx, uint32_t pszMaxRxLen, uin
 }
 
 int uart_send(const serial_port sp, const uint8_t *p_tx, const uint32_t len) {
-    serial_port_windows_t *spw = (serial_port_windows_t *)sp;
+    const serial_port_windows_t *spw = (serial_port_windows_t *)sp;
     if (spw->hSocket == INVALID_SOCKET) { // serial port
         DWORD txlen = 0;
-        int res = WriteFile(((serial_port_windows_t *)sp)->hPort, p_tx, len, &txlen, NULL);
+        int res = WriteFile(spw->hPort, p_tx, len, &txlen, NULL);
         if (res)
             return PM3_SUCCESS;