Add v6 address support to SocketPool.getaddrinfo (on espressif)

jepler · jepler · commit eaa6f99d0e04 · 2024-07-17T12:18:29.000-05:00
diff --git a/ports/espressif/common-hal/socketpool/SocketPool.c b/ports/espressif/common-hal/socketpool/SocketPool.c
@@ -22,37 +22,122 @@ void common_hal_socketpool_socketpool_construct(socketpool_socketpool_obj_t *sel
 
 // common_hal_socketpool_socket is in socketpool/Socket.c to centralize open socket tracking.
 
-mp_obj_t common_hal_socketpool_socketpool_gethostbyname(socketpool_socketpool_obj_t *self,
-    const char *host) {
-
+static int getaddrinfo_common(const char *host, int service, const struct addrinfo *hints, struct addrinfo **res) {
     // As of 2022, the version of lwip in esp-idf does not handle the
     // trailing-dot syntax of domain names, so emulate it.
     // Remove this once https://github.com/espressif/esp-idf/issues/10013 has
     // been implemented
-    size_t strlen_host = strlen(host);
-    if (strlen_host && host[strlen_host - 1] == '.') {
-        mp_obj_t nodot = mp_obj_new_str(host, strlen_host - 1);
-        host = mp_obj_str_get_str(nodot);
+    if (host) {
+        size_t strlen_host = strlen(host);
+        if (strlen_host && host[strlen_host - 1] == '.') {
+            mp_obj_t nodot = mp_obj_new_str(host, strlen_host - 1);
+            host = mp_obj_str_get_str(nodot);
+        }
+    }
+
+    char service_buf[6];
+    snprintf(service_buf, sizeof(service_buf), "%d", service);
+
+    return lwip_getaddrinfo(host, service_buf, hints, res);
+}
+
+static mp_obj_t format_address(const struct sockaddr *addr, int family) {
+    char ip_str[IPADDR_STRLEN_MAX]; // big enough for any supported address type
+    const struct sockaddr_in *a = (void *)addr;
+
+    switch (family) {
+        #if CIRCUITPY_SOCKETPOOL_IPV6
+        case AF_INET6:
+            inet_ntop(family, &((const struct sockaddr_in6 *)a)->sin6_addr, ip_str, sizeof(ip_str));
+            break;
+        #endif
+        default:
+        case AF_INET:
+            inet_ntop(family, &((const struct sockaddr_in *)a)->sin_addr, ip_str, sizeof(ip_str));
+            break;
     }
+    return mp_obj_new_str(ip_str, strlen(ip_str));
+}
+
+mp_obj_t common_hal_socketpool_socketpool_gethostbyname(socketpool_socketpool_obj_t *self,
+    const char *host) {
 
     const struct addrinfo hints = {
         .ai_family = AF_INET,
         .ai_socktype = SOCK_STREAM,
     };
-    struct addrinfo *res;
-    int err = lwip_getaddrinfo(host, NULL, &hints, &res);
+    struct addrinfo *res = NULL;
+    int err = getaddrinfo_common(host, 0, &hints, &res);
     if (err != 0 || res == NULL) {
         return mp_const_none;
     }
 
-    #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));
-    lwip_freeaddrinfo(res);
+    nlr_buf_t nlr;
+    if (nlr_push(&nlr) == 0) {
+        mp_obj_t result = format_address(res->ai_addr, res->ai_family);
+        nlr_pop();
+        lwip_freeaddrinfo(res);
+        return result;
+    } else {
+        lwip_freeaddrinfo(res);
+        nlr_raise(MP_OBJ_FROM_PTR(nlr.ret_val));
+    };
+}
+
+#if CIRCUITPY_SOCKETPOOL_IPV6
+static mp_obj_t convert_sockaddr(const struct addrinfo *ai, int port) {
+    mp_int_t n_tuple = ai->ai_family == AF_INET6 ? 4 : 2;
+    mp_obj_tuple_t *result = MP_OBJ_TO_PTR(mp_obj_new_tuple(n_tuple, NULL));
+    result->items[0] = format_address(ai->ai_addr, ai->ai_family);
+    result->items[1] = MP_OBJ_NEW_SMALL_INT(port);
+    if (ai->ai_family == AF_INET6) {
+        const struct sockaddr_in6 *ai6 = (void *)ai->ai_addr;
+        result->items[2] = MP_OBJ_NEW_SMALL_INT(ai6->sin6_flowinfo);
+        result->items[3] = MP_OBJ_NEW_SMALL_INT(ai6->sin6_scope_id);
+    }
+    return result;
+}
 
-    return ip_obj;
+static mp_obj_t convert_addrinfo(const struct addrinfo *ai, int port) {
+    MP_STATIC_ASSERT(AF_INET == SOCKETPOOL_AF_INET);
+    MP_STATIC_ASSERT(AF_INET6 == SOCKETPOOL_AF_INET6);
+    // MP_STATIC_ASSERT(AF_UNSPEC == SOCKETPOOL_AF_UNSPEC);
+    mp_obj_tuple_t *result = MP_OBJ_TO_PTR(mp_obj_new_tuple(5, NULL));
+    result->items[0] = MP_OBJ_NEW_SMALL_INT(ai->ai_family);
+    result->items[1] = MP_OBJ_NEW_SMALL_INT(ai->ai_socktype);
+    result->items[2] = MP_OBJ_NEW_SMALL_INT(ai->ai_protocol);
+    result->items[3] = ai->ai_canonname ? mp_obj_new_str(ai->ai_canonname, strlen(ai->ai_canonname)) : MP_OBJ_NEW_QSTR(MP_QSTR_);
+    result->items[4] = convert_sockaddr(ai, port);
+    return result;
+}
+
+mp_obj_t common_hal_socketpool_getaddrinfo_raise(socketpool_socketpool_obj_t *self, const char *host, int port, int family, int type, int proto, int flags) {
+    MP_STATIC_ASSERT(LWIP_IPV6);
+    const struct addrinfo hints = {
+        .ai_flags = flags,
+        .ai_family = family,
+        .ai_protocol = proto,
+        .ai_socktype = type,
+    };
+
+    struct addrinfo *res = NULL;
+    int err = getaddrinfo_common(host, port, &hints, &res);
+    if (err != 0 || res == NULL) {
+        common_hal_socketpool_socketpool_raise_gaierror_noname();
+    }
+
+    nlr_buf_t nlr;
+    if (nlr_push(&nlr) == 0) {
+        mp_obj_t result = mp_obj_new_list(0, NULL);
+        for (struct addrinfo *ai = res; ai; ai = ai->ai_next) {
+            mp_obj_list_append(result, convert_addrinfo(ai, port));
+        }
+        nlr_pop();
+        lwip_freeaddrinfo(res);
+        return result;
+    } else {
+        lwip_freeaddrinfo(res);
+        nlr_raise(MP_OBJ_FROM_PTR(nlr.ret_val));
+    }
 }
+#endif
diff --git a/ports/espressif/mpconfigport.mk b/ports/espressif/mpconfigport.mk
@@ -49,6 +49,7 @@ CIRCUITPY_SYNTHIO_MAX_CHANNELS ?= 12
 CIRCUITPY_TOUCHIO_USE_NATIVE ?= 1
 CIRCUITPY_WATCHDOG ?= 1
 CIRCUITPY_WIFI ?= 1
+CIRCUITPY_SOCKETPOOL_IPV6 ?= 1
 
 # Enable _eve module
 CIRCUITPY__EVE ?= 1
diff --git a/py/circuitpy_mpconfig.mk b/py/circuitpy_mpconfig.mk
@@ -496,6 +496,9 @@ CFLAGS += -DCIRCUITPY_SKIP_SAFE_MODE_WAIT=$(CIRCUITPY_SKIP_SAFE_MODE_WAIT)
 CIRCUITPY_SOCKETPOOL ?= $(CIRCUITPY_WIFI)
 CFLAGS += -DCIRCUITPY_SOCKETPOOL=$(CIRCUITPY_SOCKETPOOL)
 
+CIRCUITPY_SOCKETPOOL_IPV6 ?= 0
+CFLAGS += -DCIRCUITPY_SOCKETPOOL_IPV6=$(CIRCUITPY_SOCKETPOOL_IPV6)
+
 CIRCUITPY_SSL ?= $(CIRCUITPY_WIFI)
 CFLAGS += -DCIRCUITPY_SSL=$(CIRCUITPY_SSL)
 
diff --git a/shared-bindings/socketpool/SocketPool.c b/shared-bindings/socketpool/SocketPool.c
@@ -142,6 +142,16 @@ static mp_obj_t socketpool_socketpool_getaddrinfo(size_t n_args, const mp_obj_t
     mp_arg_val_t args[MP_ARRAY_SIZE(allowed_args)];
     mp_arg_parse_all(n_args - 1, pos_args + 1, kw_args, MP_ARRAY_SIZE(allowed_args), allowed_args, args);
 
+    #if CIRCUITPY_SOCKETPOOL_IPV6
+    return common_hal_socketpool_getaddrinfo_raise(
+        self,
+        mp_obj_str_get_str(args[ARG_host].u_obj),
+        args[ARG_port].u_int,
+        args[ARG_family].u_int,
+        args[ARG_type].u_int,
+        args[ARG_proto].u_int,
+        args[ARG_flags].u_int);
+    #else
     const char *host = mp_obj_str_get_str(args[ARG_host].u_obj);
     mp_int_t port = args[ARG_port].u_int;
     mp_obj_t ip_str = mp_const_none;
@@ -164,6 +174,7 @@ static mp_obj_t socketpool_socketpool_getaddrinfo(size_t n_args, const mp_obj_t
     sockaddr->items[1] = MP_OBJ_NEW_SMALL_INT(port);
     tuple->items[4] = MP_OBJ_FROM_PTR(sockaddr);
     return mp_obj_new_list(1, (mp_obj_t *)&tuple);
+    #endif
 }
 static MP_DEFINE_CONST_FUN_OBJ_KW(socketpool_socketpool_getaddrinfo_obj, 1, socketpool_socketpool_getaddrinfo);
 
diff --git a/shared-bindings/socketpool/SocketPool.h b/shared-bindings/socketpool/SocketPool.h
@@ -32,3 +32,7 @@ bool socketpool_socket(socketpool_socketpool_obj_t *self,
     int proto, socketpool_socket_obj_t *sock);
 
 NORETURN void common_hal_socketpool_socketpool_raise_gaierror_noname(void);
+
+#if CIRCUITPY_SOCKETPOOL_IPV6
+mp_obj_t common_hal_socketpool_getaddrinfo_raise(socketpool_socketpool_obj_t *self, const char *host, int port, int family, int type, int proto, int flags);
+#endif
