extmod/modusocket: Add socket state to track new/listening/conn/closed.

iabdalkader · dpgeorge · commit 70bf6ab6fb72 · 2022-06-08T14:13:59.000+10:00
diff --git a/extmod/modnetwork.h b/extmod/modnetwork.h
@@ -47,6 +47,11 @@
 #define MOD_NETWORK_SO_SNDTIMEO     (0x1005)
 #define MOD_NETWORK_SO_RCVTIMEO     (0x1006)
 
+#define MOD_NETWORK_SS_NEW          (0)
+#define MOD_NETWORK_SS_LISTENING    (1)
+#define MOD_NETWORK_SS_CONNECTED    (2)
+#define MOD_NETWORK_SS_CLOSED       (3)
+
 #if MICROPY_PY_LWIP
 struct netif;
 void mod_network_lwip_init(void);
@@ -89,9 +94,10 @@ typedef struct _mod_network_socket_obj_t {
     int32_t fileno  : 16;
     int32_t timeout;
     mp_obj_t callback;
+    int32_t state   : 8;
     #if MICROPY_PY_USOCKET_EXTENDED_STATE
     // Extended socket state for NICs/ports that need it.
-    void *state;
+    void *_private;
     #endif
 } mod_network_socket_obj_t;
 
diff --git a/extmod/modusocket.c b/extmod/modusocket.c
@@ -73,8 +73,9 @@ STATIC mp_obj_t socket_make_new(const mp_obj_type_t *type, size_t n_args, size_t
     }
     s->timeout = -1;
     s->callback = MP_OBJ_NULL;
+    s->state = MOD_NETWORK_SS_NEW;
     #if MICROPY_PY_USOCKET_EXTENDED_STATE
-    s->state = NULL;
+    s->_private = NULL;
     #endif
 
     return MP_OBJ_FROM_PTR(s);
@@ -143,6 +144,9 @@ STATIC mp_obj_t socket_listen(size_t n_args, const mp_obj_t *args) {
         mp_raise_OSError(_errno);
     }
 
+    // set socket state
+    self->state = MOD_NETWORK_SS_LISTENING;
+
     return mp_const_none;
 }
 STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(socket_listen_obj, 1, 2, socket_listen);
@@ -171,8 +175,9 @@ STATIC mp_obj_t socket_accept(mp_obj_t self_in) {
     socket2->fileno = -1;
     socket2->timeout = -1;
     socket2->callback = MP_OBJ_NULL;
+    socket2->state = MOD_NETWORK_SS_NEW;
     #if MICROPY_PY_USOCKET_EXTENDED_STATE
-    socket2->state = NULL;
+    socket2->_private = NULL;
     #endif
 
     // accept incoming connection
@@ -213,6 +218,9 @@ STATIC mp_obj_t socket_connect(mp_obj_t self_in, mp_obj_t addr_in) {
         mp_raise_OSError(_errno);
     }
 
+    // set socket state
+    self->state = MOD_NETWORK_SS_CONNECTED;
+
     return mp_const_none;
 }
 STATIC MP_DEFINE_CONST_FUN_OBJ_2(socket_connect_obj, socket_connect);
@@ -493,6 +501,7 @@ mp_uint_t socket_ioctl(mp_obj_t self_in, mp_uint_t request, uintptr_t arg, int *
         if (self->nic != MP_OBJ_NULL) {
             self->nic_type->close(self);
             self->nic = MP_OBJ_NULL;
+            self->state = MOD_NETWORK_SS_CLOSED;
         }
         return 0;
     }
