Ensure that user sees all the data before connection is closed

If user throttles receive by setting recv_mbuf_limit,
after the net interface reports connection as closed we must wait
for data to trickle through before disposing of it.
There can still b data in the buffers (e.g. SSL).

CL: mg: Ensure that user sees all the data before connection is closed

PUBLISHED_FROM=22be0fa368950a9fdb03cfb00febc7c0a1674b01

docs/c-api/mg_net.h/struct_mg_connection.md

@@ -46,6 +46,7 @@ signature: |
   #define MG_F_WANT_READ (1 << 6)          /* SSL specific */
   #define MG_F_WANT_WRITE (1 << 7)         /* SSL specific */
   #define MG_F_IS_WEBSOCKET (1 << 8)       /* Websocket specific */
+  #define MG_F_RECV_AND_CLOSE (1 << 9) /* Drain rx and close the connection. */
   
   /* Flags that are settable by user */
   #define MG_F_SEND_AND_CLOSE (1 << 10)      /* Push remaining data and close  */

mongoose.c

@@ -2474,8 +2474,16 @@ MG_INTERNAL size_t recv_avail_size(struct mg_connection *conn, size_t max) {
 static int mg_do_recv(struct mg_connection *nc);
 
 int mg_if_poll(struct mg_connection *nc, double now) {
-  if ((nc->flags & MG_F_CLOSE_IMMEDIATELY) ||
-      (nc->send_mbuf.len == 0 && (nc->flags & MG_F_SEND_AND_CLOSE))) {
+  if (nc->flags & MG_F_CLOSE_IMMEDIATELY) {
+    mg_close_conn(nc);
+    return 0;
+  } else if (nc->flags & MG_F_SEND_AND_CLOSE) {
+    if (nc->send_mbuf.len == 0) {
+      nc->flags |= MG_F_CLOSE_IMMEDIATELY;
+      mg_close_conn(nc);
+      return 0;
+    }
+  } else if (nc->flags & MG_F_RECV_AND_CLOSE) {
     mg_close_conn(nc);
     return 0;
   }
@@ -2518,6 +2526,13 @@ void mg_destroy_conn(struct mg_connection *conn, int destroy_if) {
 }
 
 void mg_close_conn(struct mg_connection *conn) {
+  /* See if there's any remaining data to deliver. Skip if user completely
+   * throttled the connection there will be no progress anyway. */
+  if (conn->sock != INVALID_SOCKET && mg_do_recv(conn) == -2) {
+    /* Receive is throttled, wait. */
+    conn->flags |= MG_F_RECV_AND_CLOSE;
+    return;
+  }
 #if MG_ENABLE_SSL
   if (conn->flags & MG_F_SSL_HANDSHAKE_DONE) {
     mg_ssl_if_conn_close_notify(conn);
@@ -2608,6 +2623,7 @@ void mg_mgr_free(struct mg_mgr *m) {
 
   for (conn = m->active_connections; conn != NULL; conn = tmp_conn) {
     tmp_conn = conn->next;
+    conn->flags |= MG_F_CLOSE_IMMEDIATELY;
     mg_close_conn(conn);
   }
 
@@ -2913,7 +2929,10 @@ static int mg_do_recv(struct mg_connection *nc) {
   }
   do {
     len = recv_avail_size(nc, len);
-    if (len == 0) return -2;
+    if (len == 0) {
+      res = -2;
+      break;
+    }
     if (nc->recv_mbuf.size < nc->recv_mbuf.len + len) {
       mbuf_resize(&nc->recv_mbuf, nc->recv_mbuf.len + len);
     }
@@ -15866,7 +15885,6 @@ void mg_ev_mgr_lwip_process_signals(struct mg_mgr *mgr) {
         break;
       }
       case MG_SIG_CLOSE_CONN: {
-        nc->flags |= MG_F_SEND_AND_CLOSE;
         mg_close_conn(nc);
         break;
       }

mongoose.h

@@ -3949,6 +3949,7 @@ struct mg_connection {
 #define MG_F_WANT_READ (1 << 6)          /* SSL specific */
 #define MG_F_WANT_WRITE (1 << 7)         /* SSL specific */
 #define MG_F_IS_WEBSOCKET (1 << 8)       /* Websocket specific */
+#define MG_F_RECV_AND_CLOSE (1 << 9) /* Drain rx and close the connection. */
 
 /* Flags that are settable by user */
 #define MG_F_SEND_AND_CLOSE (1 << 10)      /* Push remaining data and close  */

src/common/platforms/lwip/mg_lwip_ev_mgr.c

@@ -67,7 +67,6 @@ void mg_ev_mgr_lwip_process_signals(struct mg_mgr *mgr) {
         break;
       }
       case MG_SIG_CLOSE_CONN: {
-        nc->flags |= MG_F_SEND_AND_CLOSE;
         mg_close_conn(nc);
         break;
       }

src/mg_net.c

@@ -128,8 +128,16 @@ MG_INTERNAL size_t recv_avail_size(struct mg_connection *conn, size_t max) {
 static int mg_do_recv(struct mg_connection *nc);
 
 int mg_if_poll(struct mg_connection *nc, double now) {
-  if ((nc->flags & MG_F_CLOSE_IMMEDIATELY) ||
-      (nc->send_mbuf.len == 0 && (nc->flags & MG_F_SEND_AND_CLOSE))) {
+  if (nc->flags & MG_F_CLOSE_IMMEDIATELY) {
+    mg_close_conn(nc);
+    return 0;
+  } else if (nc->flags & MG_F_SEND_AND_CLOSE) {
+    if (nc->send_mbuf.len == 0) {
+      nc->flags |= MG_F_CLOSE_IMMEDIATELY;
+      mg_close_conn(nc);
+      return 0;
+    }
+  } else if (nc->flags & MG_F_RECV_AND_CLOSE) {
     mg_close_conn(nc);
     return 0;
   }
@@ -172,6 +180,13 @@ void mg_destroy_conn(struct mg_connection *conn, int destroy_if) {
 }
 
 void mg_close_conn(struct mg_connection *conn) {
+  /* See if there's any remaining data to deliver. Skip if user completely
+   * throttled the connection there will be no progress anyway. */
+  if (conn->sock != INVALID_SOCKET && mg_do_recv(conn) == -2) {
+    /* Receive is throttled, wait. */
+    conn->flags |= MG_F_RECV_AND_CLOSE;
+    return;
+  }
 #if MG_ENABLE_SSL
   if (conn->flags & MG_F_SSL_HANDSHAKE_DONE) {
     mg_ssl_if_conn_close_notify(conn);
@@ -262,6 +277,7 @@ void mg_mgr_free(struct mg_mgr *m) {
 
   for (conn = m->active_connections; conn != NULL; conn = tmp_conn) {
     tmp_conn = conn->next;
+    conn->flags |= MG_F_CLOSE_IMMEDIATELY;
     mg_close_conn(conn);
   }
 
@@ -567,7 +583,10 @@ static int mg_do_recv(struct mg_connection *nc) {
   }
   do {
     len = recv_avail_size(nc, len);
-    if (len == 0) return -2;
+    if (len == 0) {
+      res = -2;
+      break;
+    }
     if (nc->recv_mbuf.size < nc->recv_mbuf.len + len) {
       mbuf_resize(&nc->recv_mbuf, nc->recv_mbuf.len + len);
     }

src/mg_net.h

@@ -140,6 +140,7 @@ struct mg_connection {
 #define MG_F_WANT_READ (1 << 6)          /* SSL specific */
 #define MG_F_WANT_WRITE (1 << 7)         /* SSL specific */
 #define MG_F_IS_WEBSOCKET (1 << 8)       /* Websocket specific */
+#define MG_F_RECV_AND_CLOSE (1 << 9) /* Drain rx and close the connection. */
 
 /* Flags that are settable by user */
 #define MG_F_SEND_AND_CLOSE (1 << 10)      /* Push remaining data and close  */