From e2dfac946dd4e4fa2312fde0aae90b18e738ad12 Mon Sep 17 00:00:00 2001
From: Deomid Ryabkov <rojer@cesanta.com>
Date: Mon, 10 Dec 2018 19:08:12 +0000
Subject: [PATCH] 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 |  1 +
 mongoose.c                                  | 26 +++++++++++++++++----
 mongoose.h                                  |  1 +
 src/common/platforms/lwip/mg_lwip_ev_mgr.c  |  1 -
 src/mg_net.c                                | 25 +++++++++++++++++---
 src/mg_net.h                                |  1 +
 6 files changed, 47 insertions(+), 8 deletions(-)

diff --git a/docs/c-api/mg_net.h/struct_mg_connection.md b/docs/c-api/mg_net.h/struct_mg_connection.md
index be69d16fd6..cf6bd7ff9e 100644
--- a/docs/c-api/mg_net.h/struct_mg_connection.md
+++ b/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  */
diff --git a/mongoose.c b/mongoose.c
index 3196c065e9..457b23b9c6 100644
--- a/mongoose.c
+++ b/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;
       }
diff --git a/mongoose.h b/mongoose.h
index 2b7c083051..f9dd87d4e7 100644
--- a/mongoose.h
+++ b/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  */
diff --git a/src/common/platforms/lwip/mg_lwip_ev_mgr.c b/src/common/platforms/lwip/mg_lwip_ev_mgr.c
index b495464e91..ab8f19f783 100644
--- a/src/common/platforms/lwip/mg_lwip_ev_mgr.c
+++ b/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;
       }
diff --git a/src/mg_net.c b/src/mg_net.c
index ebf0513f91..3b1b7038fc 100644
--- a/src/mg_net.c
+++ b/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);
     }
diff --git a/src/mg_net.h b/src/mg_net.h
index aa8ed17158..cbe75b4c3f 100644
--- a/src/mg_net.h
+++ b/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  */