Merge pull request #8752 from carson-coder/add-proto-to-sockets

Add proto to sockets

Author: jepler

ports/espressif/common-hal/socketpool/Socket.c

@@ -182,9 +182,11 @@ STATIC void mark_user_socket(int fd, socketpool_socket_obj_t *obj) {
 
 STATIC bool _socketpool_socket(socketpool_socketpool_obj_t *self,
     socketpool_socketpool_addressfamily_t family, socketpool_socketpool_sock_t type,
+    int proto,
     socketpool_socket_obj_t *sock) {
     int addr_family;
     int ipproto;
+
     if (family == SOCKETPOOL_AF_INET) {
         addr_family = AF_INET;
         ipproto = IPPROTO_IP;
@@ -202,6 +204,7 @@ STATIC bool _socketpool_socket(socketpool_socketpool_obj_t *self,
         socket_type = SOCK_DGRAM;
     } else { // SOCKETPOOL_SOCK_RAW
         socket_type = SOCK_RAW;
+        ipproto = proto;
     }
     sock->type = socket_type;
     sock->family = addr_family;
@@ -226,9 +229,9 @@ STATIC bool _socketpool_socket(socketpool_socketpool_obj_t *self,
 // special entry for workflow listener (register system socket)
 bool socketpool_socket(socketpool_socketpool_obj_t *self,
     socketpool_socketpool_addressfamily_t family, socketpool_socketpool_sock_t type,
-    socketpool_socket_obj_t *sock) {
+    int proto, socketpool_socket_obj_t *sock) {
 
-    if (!_socketpool_socket(self, family, type, sock)) {
+    if (!_socketpool_socket(self, family, type, proto, sock)) {
         return false;
     }
 
@@ -241,15 +244,15 @@ bool socketpool_socket(socketpool_socketpool_obj_t *self,
 }
 
 socketpool_socket_obj_t *common_hal_socketpool_socket(socketpool_socketpool_obj_t *self,
-    socketpool_socketpool_addressfamily_t family, socketpool_socketpool_sock_t type) {
+    socketpool_socketpool_addressfamily_t family, socketpool_socketpool_sock_t type, int proto) {
     if (family != SOCKETPOOL_AF_INET) {
         mp_raise_NotImplementedError(translate("Only IPv4 sockets supported"));
     }
 
     socketpool_socket_obj_t *sock = m_new_obj_with_finaliser(socketpool_socket_obj_t);
     sock->base.type = &socketpool_socket_type;
 
-    if (!_socketpool_socket(self, family, type, sock)) {
+    if (!_socketpool_socket(self, family, type, proto, sock)) {
         mp_raise_RuntimeError(translate("Out of sockets"));
     }
     mark_user_socket(sock->num, sock);

ports/raspberrypi/common-hal/socketpool/Socket.c

@@ -661,6 +661,7 @@ STATIC void mark_user_socket(socketpool_socket_obj_t *obj) {
 
 bool socketpool_socket(socketpool_socketpool_obj_t *self,
     socketpool_socketpool_addressfamily_t family, socketpool_socketpool_sock_t type,
+    int proto,
     socketpool_socket_obj_t *socket) {
 
     if (!register_open_socket(socket)) {
@@ -690,7 +691,7 @@ bool socketpool_socket(socketpool_socketpool_obj_t *self,
             break;
         #if MICROPY_PY_LWIP_SOCK_RAW
         case SOCKETPOOL_SOCK_RAW: {
-            socket->pcb.raw = raw_new(0);
+            socket->pcb.raw = raw_new(proto);
             break;
         }
         #endif
@@ -728,7 +729,7 @@ bool socketpool_socket(socketpool_socketpool_obj_t *self,
 }
 
 socketpool_socket_obj_t *common_hal_socketpool_socket(socketpool_socketpool_obj_t *self,
-    socketpool_socketpool_addressfamily_t family, socketpool_socketpool_sock_t type) {
+    socketpool_socketpool_addressfamily_t family, socketpool_socketpool_sock_t type, int proto) {
     if (family != SOCKETPOOL_AF_INET) {
         mp_raise_NotImplementedError(translate("Only IPv4 sockets supported"));
     }
@@ -737,7 +738,7 @@ socketpool_socket_obj_t *common_hal_socketpool_socket(socketpool_socketpool_obj_
     socketpool_socket_obj_t *socket = m_new_ll_obj_with_finaliser(socketpool_socket_obj_t);
     socket->base.type = &socketpool_socket_type;
 
-    if (!socketpool_socket(self, family, type, socket)) {
+    if (!socketpool_socket(self, family, type, proto, socket)) {
         mp_raise_RuntimeError(translate("Out of sockets"));
     }
     mark_user_socket(socket);

shared-bindings/socketpool/SocketPool.c

@@ -81,26 +81,34 @@ MP_DEFINE_EXCEPTION(gaierror, OSError)
 //|
 //|     TCP_NODELAY: int
 //|
-//|     IPPROTO_TCP: int
 //|     IPPROTO_IP: int
+//|     IPPROTO_ICMP: int
+//|     IPPROTO_TCP: int
+//|     IPPROTO_UDP: int
+//|     IPPROTO_IPV6: int
+//|     IPPROTO_RAW: int
 //|
 //|     IP_MULTICAST_TTL: int
 //|
-//|     def socket(self, family: int = AF_INET, type: int = SOCK_STREAM) -> socketpool.Socket:
+//|     def socket(
+//|         self, family: int = AF_INET, type: int = SOCK_STREAM, proto: int = IPPROTO_IP
+//|     ) -> socketpool.Socket:
 //|         """Create a new socket
 //|
 //|         :param ~int family: AF_INET or AF_INET6
 //|         :param ~int type: SOCK_STREAM, SOCK_DGRAM or SOCK_RAW
+//|         :param ~int proto: IPPROTO_IP, IPPROTO_ICMP, IPPROTO_TCP, IPPROTO_UDP, IPPROTO_IPV6or IPPROTO_RAW. Only works with SOCK_RAW
 //|
-//|         The ``proto`` (protocol) and ``fileno`` arguments available in ``socket.socket()``
-//|         in CPython are not supported.
+//|         The ``fileno`` argument available in ``socket.socket()``
+//|         in CPython is not supported.
 //|         """
 //|         ...
 STATIC mp_obj_t socketpool_socketpool_socket(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args) {
-    enum { ARG_family, ARG_type };
+    enum { ARG_family, ARG_type, ARG_proto };
     static const mp_arg_t allowed_args[] = {
         { MP_QSTR_family, MP_ARG_INT, {.u_int = SOCKETPOOL_AF_INET} },
         { MP_QSTR_type, MP_ARG_INT, {.u_int = SOCKETPOOL_SOCK_STREAM} },
+        { MP_QSTR_proto, MP_ARG_INT, {.u_int = SOCKETPOOL_IPPROTO_IP} },
     };
     socketpool_socketpool_obj_t *self = MP_OBJ_TO_PTR(pos_args[0]);
 
@@ -109,8 +117,13 @@ STATIC mp_obj_t socketpool_socketpool_socket(size_t n_args, const mp_obj_t *pos_
 
     socketpool_socketpool_addressfamily_t family = args[ARG_family].u_int;
     socketpool_socketpool_sock_t type = args[ARG_type].u_int;
+    socketpool_socketpool_ipproto_t proto = args[ARG_proto].u_int;
+
+    if (proto < 0) {
+        proto = 0;
+    }
 
-    return common_hal_socketpool_socket(self, family, type);
+    return common_hal_socketpool_socket(self, family, type, proto);
 }
 MP_DEFINE_CONST_FUN_OBJ_KW(socketpool_socketpool_socket_obj, 1, socketpool_socketpool_socket);
 
@@ -185,8 +198,12 @@ STATIC const mp_rom_map_elem_t socketpool_socketpool_locals_dict_table[] = {
 
     { MP_ROM_QSTR(MP_QSTR_TCP_NODELAY), MP_ROM_INT(SOCKETPOOL_TCP_NODELAY) },
 
-    { MP_ROM_QSTR(MP_QSTR_IPPROTO_TCP), MP_ROM_INT(SOCKETPOOL_IPPROTO_TCP) },
     { MP_ROM_QSTR(MP_QSTR_IPPROTO_IP), MP_ROM_INT(SOCKETPOOL_IPPROTO_IP) },
+    { MP_ROM_QSTR(MP_QSTR_IPPROTO_ICMP), MP_ROM_INT(SOCKETPOOL_IPPROTO_ICMP) },
+    { MP_ROM_QSTR(MP_QSTR_IPPROTO_TCP), MP_ROM_INT(SOCKETPOOL_IPPROTO_TCP) },
+    { MP_ROM_QSTR(MP_QSTR_IPPROTO_UDP), MP_ROM_INT(SOCKETPOOL_IPPROTO_UDP) },
+    { MP_ROM_QSTR(MP_QSTR_IPPROTO_IPV6), MP_ROM_INT(SOCKETPOOL_IPPROTO_IPV6) },
+    { MP_ROM_QSTR(MP_QSTR_IPPROTO_RAW), MP_ROM_INT(SOCKETPOOL_IPPROTO_RAW) },
     { MP_ROM_QSTR(MP_QSTR_IP_MULTICAST_TTL), MP_ROM_INT(SOCKETPOOL_IP_MULTICAST_TTL) },
 
     { MP_ROM_QSTR(MP_QSTR_EAI_NONAME), MP_ROM_INT(SOCKETPOOL_EAI_NONAME) },

shared-bindings/socketpool/SocketPool.h

@@ -46,7 +46,11 @@ typedef enum {
 
 typedef enum {
     SOCKETPOOL_IPPROTO_IP = 0,
+    SOCKETPOOL_IPPROTO_ICMP = 1,
     SOCKETPOOL_IPPROTO_TCP = 6,
+    SOCKETPOOL_IPPROTO_UDP = 17,
+    SOCKETPOOL_IPPROTO_IPV6 = 41,
+    SOCKETPOOL_IPPROTO_RAW = 255,
 } socketpool_socketpool_ipproto_t;
 
 typedef enum {
@@ -64,7 +68,7 @@ typedef enum {
 void common_hal_socketpool_socketpool_construct(socketpool_socketpool_obj_t *self, mp_obj_t radio);
 
 socketpool_socket_obj_t *common_hal_socketpool_socket(socketpool_socketpool_obj_t *self,
-    socketpool_socketpool_addressfamily_t family, socketpool_socketpool_sock_t type);
+    socketpool_socketpool_addressfamily_t family, socketpool_socketpool_sock_t type, int proto);
 
 mp_obj_t common_hal_socketpool_socketpool_gethostbyname(socketpool_socketpool_obj_t *self,
     const char *host);
@@ -76,7 +80,7 @@ mp_obj_t common_hal_socketpool_socketpool_gethostbyname_raise(socketpool_socketp
 // closed automatically.
 bool socketpool_socket(socketpool_socketpool_obj_t *self,
     socketpool_socketpool_addressfamily_t family, socketpool_socketpool_sock_t type,
-    socketpool_socket_obj_t *sock);
+    int proto, socketpool_socket_obj_t *sock);
 
 NORETURN void common_hal_socketpool_socketpool_raise_gaierror_noname(void);
 

supervisor/shared/web_workflow/web_workflow.c

@@ -354,7 +354,7 @@ bool supervisor_start_web_workflow(void) {
     }
 
     if (common_hal_socketpool_socket_get_closed(&listening)) {
-        socketpool_socket(&pool, SOCKETPOOL_AF_INET, SOCKETPOOL_SOCK_STREAM, &listening);
+        socketpool_socket(&pool, SOCKETPOOL_AF_INET, SOCKETPOOL_SOCK_STREAM, 0, &listening);
         common_hal_socketpool_socket_settimeout(&listening, 0);
         // Bind to any ip. (Not checking for failures)
         common_hal_socketpool_socket_bind(&listening, "", 0, web_api_port);