Make return value of recvfrom_into & accept right for v6 addresses

jepler · jepler · commit bee32ea82656 · 2024-07-17T17:15:22.000-05:00
only tested recvfrom_into as can't bind() v6 addresses yet

wifi workflow may be broken by this or maybe it was broken before
diff --git a/ports/espressif/common-hal/socketpool/Socket.c b/ports/espressif/common-hal/socketpool/Socket.c
@@ -12,6 +12,7 @@
 #include "py/runtime.h"
 #include "shared-bindings/socketpool/SocketPool.h"
 #include "common-hal/socketpool/__init__.h"
+#include "common-hal/wifi/__init__.h"
 #if CIRCUITPY_SSL
 #include "shared-bindings/ssl/SSLSocket.h"
 #include "shared-module/ssl/SSLSocket.h"
@@ -269,9 +270,9 @@ socketpool_socket_obj_t *common_hal_socketpool_socket(socketpool_socketpool_obj_
     return sock;
 }
 
-int socketpool_socket_accept(socketpool_socket_obj_t *self, uint8_t *ip, uint32_t *port, socketpool_socket_obj_t *accepted) {
-    struct sockaddr_in accept_addr;
-    socklen_t socklen = sizeof(accept_addr);
+int socketpool_socket_accept(socketpool_socket_obj_t *self, mp_obj_t *peer_out, socketpool_socket_obj_t *accepted) {
+    struct sockaddr_storage peer_addr;
+    socklen_t socklen = sizeof(peer_addr);
     int newsoc = -1;
     bool timed_out = false;
     uint64_t start_ticks = supervisor_ticks_ms64();
@@ -282,20 +283,17 @@ int socketpool_socket_accept(socketpool_socket_obj_t *self, uint8_t *ip, uint32_
             timed_out = supervisor_ticks_ms64() - start_ticks >= self->timeout_ms;
         }
         RUN_BACKGROUND_TASKS;
-        newsoc = lwip_accept(self->num, (struct sockaddr *)&accept_addr, &socklen);
+        newsoc = lwip_accept(self->num, (struct sockaddr *)&peer_addr, &socklen);
         // In non-blocking mode, fail instead of timing out
         if (newsoc == -1 && (self->timeout_ms == 0 || mp_hal_is_interrupted())) {
             return -MP_EAGAIN;
         }
     }
 
-    if (!timed_out) {
-        // harmless on failure but avoiding memcpy is faster
-        memcpy((void *)ip, (void *)&accept_addr.sin_addr.s_addr, sizeof(accept_addr.sin_addr.s_addr));
-        *port = accept_addr.sin_port;
-    } else {
+    if (timed_out) {
         return -ETIMEDOUT;
     }
+
     if (newsoc < 0) {
         return -MP_EBADF;
     }
@@ -320,13 +318,16 @@ int socketpool_socket_accept(socketpool_socket_obj_t *self, uint8_t *ip, uint32_
         accepted->type = self->type;
     }
 
+    if (peer_out) {
+        *peer_out = sockaddr_to_tuple(&peer_addr);
+    }
+
     return newsoc;
 }
 
-socketpool_socket_obj_t *common_hal_socketpool_socket_accept(socketpool_socket_obj_t *self,
-    uint8_t *ip, uint32_t *port) {
+socketpool_socket_obj_t *common_hal_socketpool_socket_accept(socketpool_socket_obj_t *self, mp_obj_t *peer_out) {
     socketpool_socket_obj_t *sock = m_new_obj_with_finaliser(socketpool_socket_obj_t);
-    int newsoc = socketpool_socket_accept(self, ip, port, NULL);
+    int newsoc = socketpool_socket_accept(self, peer_out, NULL);
 
     if (newsoc > 0) {
         // Create the socket
@@ -344,6 +345,7 @@ socketpool_socket_obj_t *common_hal_socketpool_socket_accept(socketpool_socket_o
     }
 }
 
+// TODO: ipv6
 size_t common_hal_socketpool_socket_bind(socketpool_socket_obj_t *self,
     const char *host, size_t hostlen, uint32_t port) {
     struct sockaddr_in bind_addr;
@@ -490,9 +492,9 @@ bool common_hal_socketpool_socket_listen(socketpool_socket_obj_t *self, int back
 }
 
 mp_uint_t common_hal_socketpool_socket_recvfrom_into(socketpool_socket_obj_t *self,
-    uint8_t *buf, uint32_t len, uint8_t *ip, uint32_t *port) {
+    uint8_t *buf, uint32_t len, mp_obj_t *source_out) {
 
-    struct sockaddr_in source_addr;
+    struct sockaddr_storage source_addr;
     socklen_t socklen = sizeof(source_addr);
 
     // LWIP Socket
@@ -514,10 +516,7 @@ mp_uint_t common_hal_socketpool_socket_recvfrom_into(socketpool_socket_obj_t *se
         }
     }
 
-    if (!timed_out) {
-        memcpy((void *)ip, (void *)&source_addr.sin_addr.s_addr, sizeof(source_addr.sin_addr.s_addr));
-        *port = htons(source_addr.sin_port);
-    } else {
+    if (timed_out) {
         mp_raise_OSError(ETIMEDOUT);
     }
 
@@ -526,6 +525,10 @@ mp_uint_t common_hal_socketpool_socket_recvfrom_into(socketpool_socket_obj_t *se
         return 0;
     }
 
+    if (source_out) {
+        *source_out = sockaddr_to_tuple(&source_addr);
+    }
+
     return received;
 }
 
diff --git a/ports/espressif/common-hal/wifi/__init__.c b/ports/espressif/common-hal/wifi/__init__.c
@@ -289,6 +289,23 @@ mp_obj_t sockaddr_to_str(const struct sockaddr_storage *sockaddr) {
     return mp_obj_new_str(buf, strlen(buf));
 }
 
+mp_obj_t sockaddr_to_tuple(const struct sockaddr_storage *sockaddr) {
+    mp_obj_t args[4] = {
+        sockaddr_to_str(sockaddr),
+    };
+    int n = 2;
+    if (sockaddr->ss_family == AF_INET6) {
+        const struct sockaddr_in6 *addr6 = (const void *)sockaddr;
+        args[1] = MP_OBJ_NEW_SMALL_INT(htons(addr6->sin6_port));
+        args[2] = MP_OBJ_NEW_SMALL_INT(addr6->sin6_flowinfo);
+        args[3] = MP_OBJ_NEW_SMALL_INT(addr6->sin6_scope_id);
+        n = 4;
+    } else {
+        const struct sockaddr_in *addr = (const void *)sockaddr;
+        args[1] = MP_OBJ_NEW_SMALL_INT(htons(addr->sin_port));
+    }
+    return mp_obj_new_tuple(n, args);
+}
 
 void sockaddr_to_espaddr(const struct sockaddr_storage *sockaddr, esp_ip_addr_t *espaddr) {
     MP_STATIC_ASSERT(IPADDR_TYPE_V4 == ESP_IPADDR_TYPE_V4);
diff --git a/ports/espressif/common-hal/wifi/__init__.h b/ports/espressif/common-hal/wifi/__init__.h
@@ -19,6 +19,7 @@ void ipaddress_ipaddress_to_esp_idf(mp_obj_t ip_address, ip_addr_t *esp_ip_addre
 void ipaddress_ipaddress_to_esp_idf_ip4(mp_obj_t ip_address, esp_ip4_addr_t *esp_ip_address);
 
 mp_obj_t sockaddr_to_str(const struct sockaddr_storage *addr);
+mp_obj_t sockaddr_to_tuple(const struct sockaddr_storage *addr);
 mp_obj_t espaddr_to_str(const esp_ip_addr_t *espaddr);
 mp_obj_t espaddr4_to_str(const esp_ip4_addr_t *espaddr);
 mp_obj_t espaddr6_to_str(const esp_ip6_addr_t *espaddr);
diff --git a/shared-bindings/socketpool/Socket.c b/shared-bindings/socketpool/Socket.c
@@ -57,14 +57,10 @@ static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(socketpool_socket___exit___obj, 4, 4,
 //|         Returns a tuple of (new_socket, remote_address)"""
 static mp_obj_t _socketpool_socket_accept(mp_obj_t self_in) {
     socketpool_socket_obj_t *self = MP_OBJ_TO_PTR(self_in);
-    uint8_t ip[4];
-    uint32_t port;
-
-    socketpool_socket_obj_t *sock = common_hal_socketpool_socket_accept(self, ip, &port);
 
     mp_obj_t tuple_contents[2];
-    tuple_contents[0] = MP_OBJ_FROM_PTR(sock);
-    tuple_contents[1] = netutils_format_inet_addr(ip, port, NETUTILS_BIG);
+    tuple_contents[0] = MP_OBJ_FROM_PTR(common_hal_socketpool_socket_accept(self, &tuple_contents[1]));
+
     return mp_obj_new_tuple(2, tuple_contents);
 }
 static MP_DEFINE_CONST_FUN_OBJ_1(socketpool_socket_accept_obj, _socketpool_socket_accept);
@@ -158,13 +154,9 @@ static mp_obj_t socketpool_socket_recvfrom_into(mp_obj_t self_in, mp_obj_t data_
     mp_buffer_info_t bufinfo;
     mp_get_buffer_raise(data_in, &bufinfo, MP_BUFFER_WRITE);
 
-    byte ip[4];
-    uint32_t port;
-    mp_int_t ret = common_hal_socketpool_socket_recvfrom_into(self,
-        (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);
+    tuple_contents[0] = mp_obj_new_int_from_uint(common_hal_socketpool_socket_recvfrom_into(self,
+        (byte *)bufinfo.buf, bufinfo.len, &tuple_contents[1]));
     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
@@ -10,7 +10,7 @@
 
 extern const mp_obj_type_t socketpool_socket_type;
 
-socketpool_socket_obj_t *common_hal_socketpool_socket_accept(socketpool_socket_obj_t *self, uint8_t *ip, uint32_t *port);
+socketpool_socket_obj_t *common_hal_socketpool_socket_accept(socketpool_socket_obj_t *self, mp_obj_t *peer_out);
 size_t common_hal_socketpool_socket_bind(socketpool_socket_obj_t *self, const char *host, size_t hostlen, uint32_t port);
 void common_hal_socketpool_socket_close(socketpool_socket_obj_t *self);
 void common_hal_socketpool_socket_connect(socketpool_socket_obj_t *self, const char *host, size_t hostlen, uint32_t port);
@@ -20,7 +20,7 @@ mp_uint_t common_hal_socketpool_socket_get_timeout(socketpool_socket_obj_t *self
 mp_int_t common_hal_socketpool_socket_get_type(socketpool_socket_obj_t *self);
 bool common_hal_socketpool_socket_listen(socketpool_socket_obj_t *self, int backlog);
 mp_uint_t common_hal_socketpool_socket_recvfrom_into(socketpool_socket_obj_t *self,
-    uint8_t *buf, uint32_t len, uint8_t *ip, uint32_t *port);
+    uint8_t *buf, uint32_t len, mp_obj_t *peer_out);
 mp_uint_t common_hal_socketpool_socket_recv_into(socketpool_socket_obj_t *self, const uint8_t *buf, uint32_t len);
 mp_uint_t common_hal_socketpool_socket_send(socketpool_socket_obj_t *self, const uint8_t *buf, uint32_t len);
 mp_uint_t common_hal_socketpool_socket_sendto(socketpool_socket_obj_t *self,
@@ -31,7 +31,7 @@ bool common_hal_socketpool_readable(socketpool_socket_obj_t *self);
 bool common_hal_socketpool_writable(socketpool_socket_obj_t *self);
 
 // Non-allocating versions for internal use.
-int socketpool_socket_accept(socketpool_socket_obj_t *self, uint8_t *ip, uint32_t *port, socketpool_socket_obj_t *accepted);
+int socketpool_socket_accept(socketpool_socket_obj_t *self, mp_obj_t *peer_out, socketpool_socket_obj_t *accepted);
 void socketpool_socket_close(socketpool_socket_obj_t *self);
 int socketpool_socket_send(socketpool_socket_obj_t *self, const uint8_t *buf, uint32_t len);
 int socketpool_socket_recv_into(socketpool_socket_obj_t *self,
diff --git a/supervisor/shared/web_workflow/web_workflow.c b/supervisor/shared/web_workflow/web_workflow.c
@@ -1592,12 +1592,10 @@ void supervisor_web_workflow_background(void *data) {
         if ((!common_hal_socketpool_socket_get_connected(&active) ||
              (!active_request.in_progress && !active_request.new_socket)) &&
             !common_hal_socketpool_socket_get_closed(&listening)) {
-            uint32_t ip;
-            uint32_t port;
             if (!common_hal_socketpool_socket_get_closed(&active)) {
                 common_hal_socketpool_socket_close(&active);
             }
-            int newsoc = socketpool_socket_accept(&listening, (uint8_t *)&ip, &port, &active);
+            int newsoc = socketpool_socket_accept(&listening, NULL, &active);
             if (newsoc == -EBADF) {
                 common_hal_socketpool_socket_close(&listening);
                 break;
