extmod/modlwip: Implement raw sockets for lwIP.

Configurable via MICROPY_PY_LWIP_SOCK_RAW.

1 files changed, 101 insertions, 23 deletions

diff --git a/extmod/modlwip.c b/extmod/modlwip.c
index be932b6a9..050937750 100644
--- a/extmod/modlwip.c
+++ b/extmod/modlwip.c
@@ -40,7 +40,7 @@
 #include "lwip/init.h"
 #include "lwip/tcp.h"
 #include "lwip/udp.h"
-//#include "lwip/raw.h"
+#include "lwip/raw.h"
 #include "lwip/dns.h"
 #include "lwip/igmp.h"
 #if LWIP_VERSION_MAJOR < 2
@@ -280,6 +280,7 @@ typedef struct _lwip_socket_obj_t {
     volatile union {
         struct tcp_pcb *tcp;
         struct udp_pcb *udp;
+        struct raw_pcb *raw;
     } pcb;
     volatile union {
         struct pbuf *pbuf;
@@ -361,6 +362,26 @@ static inline void exec_user_callback(lwip_socket_obj_t *socket) {
     }
 }
 
+#if MICROPY_PY_LWIP_SOCK_RAW
+// Callback for incoming raw packets.
+#if LWIP_VERSION_MAJOR < 2
+STATIC u8_t _lwip_raw_incoming(void *arg, struct raw_pcb *pcb, struct pbuf *p, ip_addr_t *addr)
+#else
+STATIC u8_t _lwip_raw_incoming(void *arg, struct raw_pcb *pcb, struct pbuf *p, const ip_addr_t *addr)
+#endif
+{
+    lwip_socket_obj_t *socket = (lwip_socket_obj_t*)arg;
+
+    if (socket->incoming.pbuf != NULL) {
+        pbuf_free(p);
+    } else {
+        socket->incoming.pbuf = p;
+        memcpy(&socket->peer, addr, sizeof(socket->peer));
+    }
+    return 1; // we ate the packet
+}
+#endif
+
 // Callback for incoming UDP packets. We simply stash the packet and the source address,
 // in case we need it for recvfrom.
 #if LWIP_VERSION_MAJOR < 2
@@ -515,8 +536,8 @@ STATIC err_t _lwip_tcp_recv(void *arg, struct tcp_pcb *tcpb, struct pbuf *p, err
 // Functions for socket send/receive operations. Socket send/recv and friends call
 // these to do the work.
 
-// Helper function for send/sendto to handle UDP packets.
-STATIC mp_uint_t lwip_udp_send(lwip_socket_obj_t *socket, const byte *buf, mp_uint_t len, byte *ip, mp_uint_t port, int *_errno) {
+// Helper function for send/sendto to handle raw/UDP packets.
+STATIC mp_uint_t lwip_raw_udp_send(lwip_socket_obj_t *socket, const byte *buf, mp_uint_t len, byte *ip, mp_uint_t port, int *_errno) {
     if (len > 0xffff) {
         // Any packet that big is probably going to fail the pbuf_alloc anyway, but may as well try
         len = 0xffff;
@@ -536,11 +557,25 @@ STATIC mp_uint_t lwip_udp_send(lwip_socket_obj_t *socket, const byte *buf, mp_ui
 
     err_t err;
     if (ip == NULL) {
-        err = udp_send(socket->pcb.udp, p);
+        #if MICROPY_PY_LWIP_SOCK_RAW
+        if (socket->type == MOD_NETWORK_SOCK_RAW) {
+            err = raw_send(socket->pcb.raw, p);
+        } else
+        #endif
+        {
+            err = udp_send(socket->pcb.udp, p);
+        }
     } else {
         ip_addr_t dest;
         IP4_ADDR(&dest, ip[0], ip[1], ip[2], ip[3]);
-        err = udp_sendto(socket->pcb.udp, p, &dest, port);
+        #if MICROPY_PY_LWIP_SOCK_RAW
+        if (socket->type == MOD_NETWORK_SOCK_RAW) {
+            err = raw_sendto(socket->pcb.raw, p, &dest);
+        } else
+        #endif
+        {
+            err = udp_sendto(socket->pcb.udp, p, &dest, port);
+        }
     }
 
     pbuf_free(p);
@@ -558,8 +593,8 @@ STATIC mp_uint_t lwip_udp_send(lwip_socket_obj_t *socket, const byte *buf, mp_ui
     return len;
 }
 
-// Helper function for recv/recvfrom to handle UDP packets
-STATIC mp_uint_t lwip_udp_receive(lwip_socket_obj_t *socket, byte *buf, mp_uint_t len, byte *ip, mp_uint_t *port, int *_errno) {
+// Helper function for recv/recvfrom to handle raw/UDP packets
+STATIC mp_uint_t lwip_raw_udp_receive(lwip_socket_obj_t *socket, byte *buf, mp_uint_t len, byte *ip, mp_uint_t *port, int *_errno) {
 
     if (socket->incoming.pbuf == NULL) {
         if (socket->timeout != -1) {
@@ -795,7 +830,13 @@ STATIC mp_obj_t lwip_socket_make_new(const mp_obj_type_t *type, size_t n_args, s
             socket->pcb.udp = udp_new();
             socket->incoming.pbuf = NULL;
             break;
-        //case MOD_NETWORK_SOCK_RAW: socket->pcb.raw = raw_new(); break;
+        #if MICROPY_PY_LWIP_SOCK_RAW
+        case MOD_NETWORK_SOCK_RAW: {
+            mp_int_t proto = n_args <= 2 ? 0 : mp_obj_get_int(args[2]);
+            socket->pcb.raw = raw_new(proto);
+            break;
+        }
+        #endif
         default: mp_raise_OSError(MP_EINVAL);
     }
 
@@ -817,6 +858,14 @@ STATIC mp_obj_t lwip_socket_make_new(const mp_obj_type_t *type, size_t n_args, s
             udp_recv(socket->pcb.udp, _lwip_udp_incoming, (void*)socket);
             break;
         }
+        #if MICROPY_PY_LWIP_SOCK_RAW
+        case MOD_NETWORK_SOCK_RAW: {
+            // Register our receive callback now. Since raw sockets don't require binding or connection
+            // before use, there's no other good time to do it.
+            raw_recv(socket->pcb.raw, _lwip_raw_incoming, (void*)socket);
+            break;
+        }
+        #endif
     }
 
     socket->timeout = -1;
@@ -1045,6 +1094,12 @@ STATIC mp_obj_t lwip_socket_connect(mp_obj_t self_in, mp_obj_t addr_in) {
             err = udp_connect(socket->pcb.udp, &dest, port);
             break;
         }
+        #if MICROPY_PY_LWIP_SOCK_RAW
+        case MOD_NETWORK_SOCK_RAW: {
+            err = raw_connect(socket->pcb.raw, &dest);
+            break;
+        }
+        #endif
     }
 
     if (err != ERR_OK) {
@@ -1079,10 +1134,12 @@ STATIC mp_obj_t lwip_socket_send(mp_obj_t self_in, mp_obj_t buf_in) {
             ret = lwip_tcp_send(socket, bufinfo.buf, bufinfo.len, &_errno);
             break;
         }
-        case MOD_NETWORK_SOCK_DGRAM: {
-            ret = lwip_udp_send(socket, bufinfo.buf, bufinfo.len, NULL, 0, &_errno);
+        case MOD_NETWORK_SOCK_DGRAM:
+        #if MICROPY_PY_LWIP_SOCK_RAW
+        case MOD_NETWORK_SOCK_RAW:
+        #endif
+            ret = lwip_raw_udp_send(socket, bufinfo.buf, bufinfo.len, NULL, 0, &_errno);
             break;
-        }
     }
     if (ret == -1) {
         mp_raise_OSError(_errno);
@@ -1108,10 +1165,12 @@ STATIC mp_obj_t lwip_socket_recv(mp_obj_t self_in, mp_obj_t len_in) {
             ret = lwip_tcp_receive(socket, (byte*)vstr.buf, len, &_errno);
             break;
         }
-        case MOD_NETWORK_SOCK_DGRAM: {
-            ret = lwip_udp_receive(socket, (byte*)vstr.buf, len, NULL, NULL, &_errno);
+        case MOD_NETWORK_SOCK_DGRAM:
+        #if MICROPY_PY_LWIP_SOCK_RAW
+        case MOD_NETWORK_SOCK_RAW:
+        #endif
+            ret = lwip_raw_udp_receive(socket, (byte*)vstr.buf, len, NULL, NULL, &_errno);
             break;
-        }
     }
     if (ret == -1) {
         mp_raise_OSError(_errno);
@@ -1143,10 +1202,12 @@ STATIC mp_obj_t lwip_socket_sendto(mp_obj_t self_in, mp_obj_t data_in, mp_obj_t
             ret = lwip_tcp_send(socket, bufinfo.buf, bufinfo.len, &_errno);
             break;
         }
-        case MOD_NETWORK_SOCK_DGRAM: {
-            ret = lwip_udp_send(socket, bufinfo.buf, bufinfo.len, ip, port, &_errno);
+        case MOD_NETWORK_SOCK_DGRAM:
+        #if MICROPY_PY_LWIP_SOCK_RAW
+        case MOD_NETWORK_SOCK_RAW:
+        #endif
+            ret = lwip_raw_udp_send(socket, bufinfo.buf, bufinfo.len, ip, port, &_errno);
             break;
-        }
     }
     if (ret == -1) {
         mp_raise_OSError(_errno);
@@ -1176,10 +1237,12 @@ STATIC mp_obj_t lwip_socket_recvfrom(mp_obj_t self_in, mp_obj_t len_in) {
             ret = lwip_tcp_receive(socket, (byte*)vstr.buf, len, &_errno);
             break;
         }
-        case MOD_NETWORK_SOCK_DGRAM: {
-            ret = lwip_udp_receive(socket, (byte*)vstr.buf, len, ip, &port, &_errno);
+        case MOD_NETWORK_SOCK_DGRAM:
+        #if MICROPY_PY_LWIP_SOCK_RAW
+        case MOD_NETWORK_SOCK_RAW:
+        #endif
+            ret = lwip_raw_udp_receive(socket, (byte*)vstr.buf, len, ip, &port, &_errno);
             break;
-        }
     }
     if (ret == -1) {
         mp_raise_OSError(_errno);
@@ -1331,7 +1394,10 @@ STATIC mp_uint_t lwip_socket_read(mp_obj_t self_in, void *buf, mp_uint_t size, i
         case MOD_NETWORK_SOCK_STREAM:
             return lwip_tcp_receive(socket, buf, size, errcode);
         case MOD_NETWORK_SOCK_DGRAM:
-            return lwip_udp_receive(socket, buf, size, NULL, NULL, errcode);
+        #if MICROPY_PY_LWIP_SOCK_RAW
+        case MOD_NETWORK_SOCK_RAW:
+        #endif
+            return lwip_raw_udp_receive(socket, buf, size, NULL, NULL, errcode);
     }
     // Unreachable
     return MP_STREAM_ERROR;
@@ -1344,7 +1410,10 @@ STATIC mp_uint_t lwip_socket_write(mp_obj_t self_in, const void *buf, mp_uint_t
         case MOD_NETWORK_SOCK_STREAM:
             return lwip_tcp_send(socket, buf, size, errcode);
         case MOD_NETWORK_SOCK_DGRAM:
-            return lwip_udp_send(socket, buf, size, NULL, 0, errcode);
+        #if MICROPY_PY_LWIP_SOCK_RAW
+        case MOD_NETWORK_SOCK_RAW:
+        #endif
+            return lwip_raw_udp_send(socket, buf, size, NULL, 0, errcode);
     }
     // Unreachable
     return MP_STREAM_ERROR;
@@ -1385,6 +1454,11 @@ STATIC mp_uint_t lwip_socket_ioctl(mp_obj_t self_in, mp_uint_t request, uintptr_
             if (socket->type == MOD_NETWORK_SOCK_DGRAM && socket->pcb.udp != NULL) {
                 // UDP socket is writable
                 ret |= MP_STREAM_POLL_WR;
+            #if MICROPY_PY_LWIP_SOCK_RAW
+            } else if (socket->type == MOD_NETWORK_SOCK_RAW && socket->pcb.raw != NULL) {
+                // raw socket is writable
+                ret |= MP_STREAM_POLL_WR;
+            #endif
             } else if (socket->pcb.tcp != NULL && tcp_sndbuf(socket->pcb.tcp) > 0) {
                 // TCP socket is writable
                 // Note: pcb.tcp==NULL if state<0, and in this case we can't call tcp_sndbuf
@@ -1438,7 +1512,9 @@ STATIC mp_uint_t lwip_socket_ioctl(mp_obj_t self_in, mp_uint_t request, uintptr_
                 break;
             }
             case MOD_NETWORK_SOCK_DGRAM: udp_remove(socket->pcb.udp); break;
-            //case MOD_NETWORK_SOCK_RAW: raw_remove(socket->pcb.raw); break;
+            #if MICROPY_PY_LWIP_SOCK_RAW
+            case MOD_NETWORK_SOCK_RAW: raw_remove(socket->pcb.raw); break;
+            #endif
         }
 
         socket->pcb.tcp = NULL;
@@ -1660,7 +1736,9 @@ STATIC const mp_rom_map_elem_t mp_module_lwip_globals_table[] = {
 
     { MP_ROM_QSTR(MP_QSTR_SOCK_STREAM), MP_ROM_INT(MOD_NETWORK_SOCK_STREAM) },
     { MP_ROM_QSTR(MP_QSTR_SOCK_DGRAM), MP_ROM_INT(MOD_NETWORK_SOCK_DGRAM) },
+    #if MICROPY_PY_LWIP_SOCK_RAW
     { MP_ROM_QSTR(MP_QSTR_SOCK_RAW), MP_ROM_INT(MOD_NETWORK_SOCK_RAW) },
+    #endif
 
     { MP_ROM_QSTR(MP_QSTR_SOL_SOCKET), MP_ROM_INT(1) },
     { MP_ROM_QSTR(MP_QSTR_SO_REUSEADDR), MP_ROM_INT(SOF_REUSEADDR) },