Implement recvfrom_into and sendto for UDP

4 files changed, 45 insertions, 40 deletions

diff --git a/ports/esp32s2/Makefile b/ports/esp32s2/Makefile
index 55d6e91d4..a2e6cb73d 100644
--- a/ports/esp32s2/Makefile
+++ b/ports/esp32s2/Makefile
@@ -188,6 +188,7 @@ SRC_C += \
 	lib/utils/pyexec.c \
 	lib/utils/stdout_helpers.c \
 	lib/utils/sys_stdio_mphal.c \
+	lib/netutils/netutils.c \
 	peripherals/pcnt.c \
 	peripherals/pins.c \
 	peripherals/rmt.c \
diff --git a/ports/esp32s2/common-hal/socketpool/Socket.c b/ports/esp32s2/common-hal/socketpool/Socket.c
index de9e6ad31..937d199f2 100644
--- a/ports/esp32s2/common-hal/socketpool/Socket.c
+++ b/ports/esp32s2/common-hal/socketpool/Socket.c
@@ -32,6 +32,11 @@
 #include "py/runtime.h"
 #include "supervisor/shared/tick.h"
 
+#include "components/lwip/lwip/src/include/lwip/err.h"
+#include "components/lwip/lwip/src/include/lwip/sockets.h"
+#include "components/lwip/lwip/src/include/lwip/sys.h"
+#include "components/lwip/lwip/src/include/lwip/netdb.h"
+
 void common_hal_socketpool_socket_settimeout(socketpool_socket_obj_t* self, mp_uint_t timeout_ms) {
     self->timeout_ms = timeout_ms;
 }
@@ -122,58 +127,55 @@ mp_uint_t common_hal_socketpool_socket_recv_into(socketpool_socket_obj_t* self,
 mp_uint_t common_hal_socketpool_socket_sendto(socketpool_socket_obj_t* self,
     const char* host, size_t hostlen, uint8_t port, const uint8_t* buf, mp_uint_t len) {
 
-    struct sockaddr_in dest_addr;
-    dest_addr.sin_addr.s_addr = inet_addr(HOST_IP_ADDR);
-    dest_addr.sin_family = AF_INET;
-    dest_addr.sin_port = htons(port);
-
-
+    // Get the IP address string
     const struct addrinfo hints = {
         .ai_family = AF_INET,
         .ai_socktype = SOCK_STREAM,
     };
-    struct addrinfo *res;
-    int err = getaddrinfo(host, NULL, &hints, &res);
-    if (err != 0 || res == NULL) {
-        return mp_const_none;
+    struct addrinfo *result;
+    int error = lwip_getaddrinfo(host, NULL, &hints, &result);
+    if (error != 0 || result == NULL) {
+        return 0;
     }
-
     #pragma GCC diagnostic push
     #pragma GCC diagnostic ignored "-Wcast-align"
-    struct in_addr *addr = &((struct sockaddr_in *)res->ai_addr)->sin_addr;
+    struct in_addr *addr = &((struct sockaddr_in *)result->ai_addr)->sin_addr;
     #pragma GCC diagnostic pop
     char ip_str[IP4ADDR_STRLEN_MAX];
     inet_ntoa_r(*addr, ip_str, IP4ADDR_STRLEN_MAX);
-    mp_obj_t ip_obj = mp_obj_new_str(ip_str, strlen(ip_str));
-    freeaddrinfo(res);
-
+    freeaddrinfo(result);
 
+    // Set parameters
+    struct sockaddr_in dest_addr;
+    dest_addr.sin_addr.s_addr = inet_addr((const char *)ip_str);
+    dest_addr.sin_family = AF_INET;
+    dest_addr.sin_port = htons(port);
 
-    int err = lwip_sendto(self->num, buf, len, 0 /* flags */,
-        (struct sockaddr *)&dest_addr, sizeof(dest_addr));
+    int bytes_sent = lwip_sendto(self->num, buf, len, 0, (struct sockaddr *)&dest_addr, sizeof(dest_addr));
+    if (bytes_sent < 0) {
+        mp_raise_BrokenPipeError();
+        return 0;
+    }
+    return bytes_sent;
 }
 
 mp_uint_t common_hal_socketpool_socket_recvfrom_into(socketpool_socket_obj_t* self,
-    const uint8_t* buf, mp_uint_t len, uint8_t* ip, uint8_t port) {
+    uint8_t* buf, mp_uint_t len, uint8_t* ip, uint *port) {
 
-    const struct addrinfo hints = {
-        .ai_family = AF_INET,
-        .ai_socktype = SOCK_STREAM,
-    };
-    struct addrinfo *res;
-    int err = getaddrinfo(host, NULL, &hints, &res);
-    if (err != 0 || res == NULL) {
-        return mp_const_none;
-    }
+    struct sockaddr_in source_addr;
+    socklen_t socklen = sizeof(source_addr);
+    int bytes_received = lwip_recvfrom(self->num, buf, len - 1, 0, (struct sockaddr *)&source_addr, &socklen);
 
-    #pragma GCC diagnostic push
-    #pragma GCC diagnostic ignored "-Wcast-align"
-    struct in_addr *addr = &((struct sockaddr_in *)res->ai_addr)->sin_addr;
-    #pragma GCC diagnostic pop
-    char ip_str[IP4ADDR_STRLEN_MAX];
-    inet_ntoa_r(*addr, ip_str, IP4ADDR_STRLEN_MAX);
-    mp_obj_t ip_obj = mp_obj_new_str(ip_str, strlen(ip_str));
-    freeaddrinfo(res);
+    memcpy((void *)ip, (void*)&source_addr.sin_addr.s_addr, sizeof source_addr.sin_addr.s_addr);
+    *port = source_addr.sin_port;
+
+    if (bytes_received < 0) {
+        mp_raise_BrokenPipeError();
+        return 0;
+    } else {
+        buf[bytes_received] = 0; // Null-terminate whatever we received
+        return bytes_received;
+    }
 }
 
 void common_hal_socketpool_socket_close(socketpool_socket_obj_t* self) {
diff --git a/shared-bindings/socketpool/Socket.c b/shared-bindings/socketpool/Socket.c
index 591b7d8bb..99e2c8fcc 100644
--- a/shared-bindings/socketpool/Socket.c
+++ b/shared-bindings/socketpool/Socket.c
@@ -36,6 +36,8 @@
 #include "py/runtime.h"
 #include "py/mperrno.h"
 
+#include "lib/netutils/netutils.h"
+
 //| class Socket:
 //|     """TCP, UDP and RAW socket. Cannot be created directly. Instead, call
 //|        `SocketPool.socket()`.
@@ -298,7 +300,7 @@ STATIC mp_obj_t socketpool_socket_sendto(mp_obj_t self_in, mp_obj_t data_in, mp_
     mp_int_t port = mp_obj_get_int(addr_items[1]);
 
     mp_int_t ret = common_hal_socketpool_socket_sendto(self, host, hostlen, port, bufinfo.buf, bufinfo.len);
-    if (!ok) {
+    if (!ret) {
         mp_raise_OSError(0);
     }
 
@@ -324,11 +326,11 @@ STATIC mp_obj_t socketpool_socket_recvfrom_into(mp_obj_t self_in, mp_obj_t data_
     byte ip[4];
     mp_uint_t port;
     mp_int_t ret = common_hal_socketpool_socket_recvfrom_into(self,
-        (byte*)bufinfo.buf, len, ip, &port);
+        (byte*)bufinfo.buf, bufinfo.len, ip, &port);
     mp_obj_t tuple_contents[2];
     tuple_contents[0] = mp_obj_new_int_from_uint(ret);
     tuple_contents[1] = netutils_format_inet_addr(ip, port, NETUTILS_BIG);
-    return mp_obj_new_tuple(2, tuple);
+    return mp_obj_new_tuple(2, tuple_contents);
 }
 STATIC MP_DEFINE_CONST_FUN_OBJ_2(socketpool_socket_recvfrom_into_obj, socketpool_socket_recvfrom_into);
 
diff --git a/shared-bindings/socketpool/Socket.h b/shared-bindings/socketpool/Socket.h
index 72a4c9e3c..54fbe59b2 100644
--- a/shared-bindings/socketpool/Socket.h
+++ b/shared-bindings/socketpool/Socket.h
@@ -36,9 +36,9 @@ bool common_hal_socketpool_socket_connect(socketpool_socket_obj_t* self, const c
 mp_uint_t common_hal_socketpool_socket_send(socketpool_socket_obj_t* self, const uint8_t* buf, mp_uint_t len);
 mp_uint_t common_hal_socketpool_socket_recv_into(socketpool_socket_obj_t* self, const uint8_t* buf, mp_uint_t len);
 mp_uint_t common_hal_socketpool_socket_sendto(socketpool_socket_obj_t* self,
-    const uint8_t* host, size_t hostlen, uint8_t port, const uint8_t* buf, mp_uint_t len);
+    const char* host, size_t hostlen, uint8_t port, const uint8_t* buf, mp_uint_t len);
 mp_uint_t common_hal_socketpool_socket_recvfrom_into(socketpool_socket_obj_t* self,
-    const uint8_t* buf, mp_uint_t len, uint8_t* ip, uint8_t port);
+    uint8_t* buf, mp_uint_t len, uint8_t* ip, uint *port);
 void common_hal_socketpool_socket_close(socketpool_socket_obj_t* self);
 bool common_hal_socketpool_socket_get_closed(socketpool_socket_obj_t* self);
 bool common_hal_socketpool_socket_get_connected(socketpool_socket_obj_t* self);