From 7d57037906cf0274af08bd2eccbfffabe0ea66e3 Mon Sep 17 00:00:00 2001
From: Paul Sokolovsky <pfalcon@users.sourceforge.net>
Date: Sun, 17 Apr 2016 18:06:45 +0300
Subject: [PATCH] extmod/modlwip: Add ability to run callback on "recv" and
 "accept" events.

To use: .setsockopt(SOL_SOCKET, 20, lambda sock: print(sock)). There's a
single underlying callback slot. For normal sockets, it serves as data
received callback, for listening sockets - connection arrived callback.
---
 extmod/modlwip.c | 28 +++++++++++++++++++++++++++-
 1 file changed, 27 insertions(+), 1 deletion(-)

diff --git a/extmod/modlwip.c b/extmod/modlwip.c
index d007b4273d..778a741b3d 100644
--- a/extmod/modlwip.c
+++ b/extmod/modlwip.c
@@ -234,6 +234,7 @@ typedef struct _lwip_socket_obj_t {
         struct pbuf *pbuf;
         struct tcp_pcb *connection;
     } incoming;
+    mp_obj_t callback;
     byte peer[4];
     mp_uint_t peer_port;
     mp_uint_t timeout;
@@ -261,6 +262,12 @@ static inline void poll_sockets(void) {
 /*******************************************************************************/
 // Callback functions for the lwIP raw API.
 
+static inline void exec_user_callback(lwip_socket_obj_t *socket) {
+    if (socket->callback != MP_OBJ_NULL) {
+        mp_call_function_1(socket->callback, socket);
+    }
+}
+
 // Callback for incoming UDP packets. We simply stash the packet and the source address,
 // in case we need it for recvfrom.
 STATIC void _lwip_udp_incoming(void *arg, struct udp_pcb *upcb, struct pbuf *p, ip_addr_t *addr, u16_t port) {
@@ -315,6 +322,7 @@ STATIC err_t _lwip_tcp_accept(void *arg, struct tcp_pcb *newpcb, err_t err) {
         return ERR_BUF;
     } else {
         socket->incoming.connection = newpcb;
+        exec_user_callback(socket);
         return ERR_OK;
     }
 }
@@ -327,12 +335,16 @@ STATIC err_t _lwip_tcp_recv(void *arg, struct tcp_pcb *tcpb, struct pbuf *p, err
         // Other side has closed connection.
         DEBUG_printf("_lwip_tcp_recv[%p]: other side closed connection\n", socket);
         socket->state = STATE_PEER_CLOSED;
+        exec_user_callback(socket);
         return ERR_OK;
     } else if (socket->incoming.pbuf != NULL) {
         // No room in the inn, let LWIP know it's still responsible for delivery later
         return ERR_BUF;
     }
     socket->incoming.pbuf = p;
+
+    exec_user_callback(socket);
+
     return ERR_OK;
 }
 
@@ -542,6 +554,7 @@ STATIC mp_obj_t lwip_socket_make_new(const mp_obj_type_t *type, mp_uint_t n_args
     socket->base.type = (mp_obj_t)&lwip_socket_type;
     socket->domain = MOD_NETWORK_AF_INET;
     socket->type = MOD_NETWORK_SOCK_STREAM;
+    socket->callback = MP_OBJ_NULL;
     if (n_args >= 1) {
         socket->domain = mp_obj_get_int(args[0]);
         if (n_args >= 2) {
@@ -717,6 +730,7 @@ STATIC mp_obj_t lwip_socket_accept(mp_obj_t self_in) {
     socket2->timeout = socket->timeout;
     socket2->state = STATE_CONNECTED;
     socket2->leftover_count = 0;
+    socket2->callback = MP_OBJ_NULL;
     tcp_arg(socket2->pcb.tcp, (void*)socket2);
     tcp_err(socket2->pcb.tcp, _lwip_tcp_error);
     tcp_recv(socket2->pcb.tcp, _lwip_tcp_recv);
@@ -978,8 +992,20 @@ STATIC MP_DEFINE_CONST_FUN_OBJ_2(lwip_socket_setblocking_obj, lwip_socket_setblo
 STATIC mp_obj_t lwip_socket_setsockopt(mp_uint_t n_args, const mp_obj_t *args) {
     (void)n_args; // always 4
     lwip_socket_obj_t *socket = args[0];
+
+    int opt = mp_obj_get_int(args[2]);
+    if (opt == 20) {
+        if (args[3] == mp_const_none) {
+            socket->callback = MP_OBJ_NULL;
+        } else {
+            socket->callback = args[3];
+        }
+        return mp_const_none;
+    }
+
+    // Integer options
     mp_int_t val = mp_obj_get_int(args[3]);
-    switch (mp_obj_get_int(args[2])) {
+    switch (opt) {
         case SOF_REUSEADDR:
             // Options are common for UDP and TCP pcb's.
             if (val) {