make FIN wait for TLS drain

mongoose.c

@@ -10514,6 +10514,7 @@ void mg_mgr_poll(struct mg_mgr *mgr, int ms) {
   mg_tcpip_poll(mgr->ifp, now);
   for (c = mgr->conns; c != NULL; c = tmp) {
     struct connstate *s = (struct connstate *) (c + 1);
+    long flush = 0;
     bool is_tls = c->is_tls && !c->is_resolving && !c->is_arplooking &&
                   !c->is_listening && !c->is_connecting;
     tmp = c->next;
@@ -10526,10 +10527,12 @@ void mg_mgr_poll(struct mg_mgr *mgr, int ms) {
     if (is_tls && (c->rtls.len > 0 || mg_tls_pending(c) > 0))
       c->is_tls_hs ? mg_tls_handshake(c) : handle_tls_recv(c);
     if (can_write(c)) write_conn(c);
-    if (is_tls && c->send.len == 0) mg_tls_flush(c);
+    if (is_tls && c->send.len == 0) flush = mg_tls_flush(c);
+    if (flush == MG_IO_ERR) mg_error(c, "tx err");
     if (c->is_draining && c->send.len == 0) {
       if (c->is_udp) c->is_closing = 1;
-      if (!c->is_udp && s->ttype != MIP_TTYPE_FIN) init_closure(c);
+      if (!c->is_udp && flush == 0 && s->ttype != MIP_TTYPE_FIN)
+        init_closure(c);
     }
     // For non-TLS, close immediately upon completing the 3-way closure
     // For TLS, handle any pending data (above) until MIP_TTYPE_FIN expires
@@ -14998,6 +15001,7 @@ void mg_mgr_poll(struct mg_mgr *mgr, int ms) {
   mg_ota_poll(mgr);
 
   for (c = mgr->conns; c != NULL; c = tmp) {
+    long flush = 0;
     bool is_resp = c->is_resp;
     tmp = c->next;
     mg_call(c, MG_EV_POLL, &now);
@@ -15019,10 +15023,12 @@ void mg_mgr_poll(struct mg_mgr *mgr, int ms) {
     } else {
       if (c->is_readable) read_conn(c);
       if (c->is_writable) write_conn(c);
-      if (c->is_tls && !c->is_tls_hs && c->send.len == 0) mg_tls_flush(c);
+      if (c->is_tls && !c->is_tls_hs && c->send.len == 0)
+        flush = mg_tls_flush(c);
     }
 
-    if (c->is_draining && c->send.len == 0) c->is_closing = 1;
+    if (flush == MG_IO_ERR) mg_error(c, "tx err");
+    if (c->is_draining && c->send.len == 0 && flush == 0) c->is_closing = 1;
     if (c->is_closing) close_conn(c);
   }
 }
@@ -19378,13 +19384,16 @@ size_t mg_tls_pending(struct mg_connection *c) {
   return tls != NULL ? tls->recv_len : 0;
 }
 
-void mg_tls_flush(struct mg_connection *c) {
+long mg_tls_flush(struct mg_connection *c) {
   struct tls_data *tls = (struct tls_data *) c->tls;
-  long n;
+  long n = 0;
   while (tls->send.len > 0 &&
          (n = mg_io_send(c, tls->send.buf, tls->send.len)) > 0) {
     mg_iobuf_del(&tls->send, 0, (size_t) n);
   }
+  if (n == MG_IO_ERR) return n;
+  c->is_tls_throttled = tls->send.len > 0;
+  return c->is_tls_throttled ? MG_IO_WAIT : 0;
 }
 
 void mg_tls_ctx_init(struct mg_mgr *mgr) {
@@ -20780,8 +20789,9 @@ size_t mg_tls_pending(struct mg_connection *c) {
   (void) c;
   return 0;
 }
-void mg_tls_flush(struct mg_connection *c) {
+long mg_tls_flush(struct mg_connection *c) {
   (void) c;
+  return 0;
 }
 void mg_tls_ctx_init(struct mg_mgr *mgr) {
   (void) mgr;
@@ -21045,17 +21055,20 @@ long mg_tls_send(struct mg_connection *c, const void *buf, size_t len) {
   return n;
 }
 
-void mg_tls_flush(struct mg_connection *c) {
+long mg_tls_flush(struct mg_connection *c) {
   struct mg_tls *tls = (struct mg_tls *) c->tls;
+  long n = 0;
   if (c->is_tls_throttled && c->is_draining) {
-    long n =
-        mbedtls_ssl_write(&tls->ssl, tls->throttled_buf, tls->throttled_len);
+    n = mbedtls_ssl_write(&tls->ssl, tls->throttled_buf, tls->throttled_len);
 #if defined(MBEDTLS_ERR_SSL_RECEIVED_NEW_SESSION_TICKET)
-    if (n == MBEDTLS_ERR_SSL_RECEIVED_NEW_SESSION_TICKET) return;
+    if (n == MBEDTLS_ERR_SSL_RECEIVED_NEW_SESSION_TICKET) return MG_IO_WAIT;
 #endif
     c->is_tls_throttled =
         (n == MBEDTLS_ERR_SSL_WANT_READ || n == MBEDTLS_ERR_SSL_WANT_WRITE);
+    if (c->is_tls_throttled) return MG_IO_WAIT;
+    if (n <= 0) return MG_IO_ERR;
   }
+  return 0;
 }
 
 void mg_tls_ctx_init(struct mg_mgr *mgr) {
@@ -21444,8 +21457,9 @@ long mg_tls_send(struct mg_connection *c, const void *buf, size_t len) {
   return n;
 }
 
-void mg_tls_flush(struct mg_connection *c) {
+long mg_tls_flush(struct mg_connection *c) {
   (void) c;
+  return 0;
 }
 
 void mg_tls_ctx_init(struct mg_mgr *mgr) {

mongoose.h

@@ -2916,7 +2916,7 @@ void mg_tls_free(struct mg_connection *);
 long mg_tls_send(struct mg_connection *, const void *buf, size_t len);
 long mg_tls_recv(struct mg_connection *, void *buf, size_t len);
 size_t mg_tls_pending(struct mg_connection *);
-void mg_tls_flush(struct mg_connection *);
+long mg_tls_flush(struct mg_connection *);
 void mg_tls_handshake(struct mg_connection *);
 void mg_tls_ctx_init(struct mg_mgr *);
 void mg_tls_ctx_free(struct mg_mgr *);

src/net_builtin.c

@@ -2273,6 +2273,7 @@ void mg_mgr_poll(struct mg_mgr *mgr, int ms) {
   mg_tcpip_poll(mgr->ifp, now);
   for (c = mgr->conns; c != NULL; c = tmp) {
     struct connstate *s = (struct connstate *) (c + 1);
+    long flush = 0;
     bool is_tls = c->is_tls && !c->is_resolving && !c->is_arplooking &&
                   !c->is_listening && !c->is_connecting;
     tmp = c->next;
@@ -2285,10 +2286,12 @@ void mg_mgr_poll(struct mg_mgr *mgr, int ms) {
     if (is_tls && (c->rtls.len > 0 || mg_tls_pending(c) > 0))
       c->is_tls_hs ? mg_tls_handshake(c) : handle_tls_recv(c);
     if (can_write(c)) write_conn(c);
-    if (is_tls && c->send.len == 0) mg_tls_flush(c);
+    if (is_tls && c->send.len == 0) flush = mg_tls_flush(c);
+    if (flush == MG_IO_ERR) mg_error(c, "tx err");
     if (c->is_draining && c->send.len == 0) {
       if (c->is_udp) c->is_closing = 1;
-      if (!c->is_udp && s->ttype != MIP_TTYPE_FIN) init_closure(c);
+      if (!c->is_udp && flush == 0 && s->ttype != MIP_TTYPE_FIN)
+        init_closure(c);
     }
     // For non-TLS, close immediately upon completing the 3-way closure
     // For TLS, handle any pending data (above) until MIP_TTYPE_FIN expires

src/sock.c

@@ -770,6 +770,7 @@ void mg_mgr_poll(struct mg_mgr *mgr, int ms) {
   mg_ota_poll(mgr);
 
   for (c = mgr->conns; c != NULL; c = tmp) {
+    long flush = 0;
     bool is_resp = c->is_resp;
     tmp = c->next;
     mg_call(c, MG_EV_POLL, &now);
@@ -791,10 +792,12 @@ void mg_mgr_poll(struct mg_mgr *mgr, int ms) {
     } else {
       if (c->is_readable) read_conn(c);
       if (c->is_writable) write_conn(c);
-      if (c->is_tls && !c->is_tls_hs && c->send.len == 0) mg_tls_flush(c);
+      if (c->is_tls && !c->is_tls_hs && c->send.len == 0)
+        flush = mg_tls_flush(c);
     }
 
-    if (c->is_draining && c->send.len == 0) c->is_closing = 1;
+    if (flush == MG_IO_ERR) mg_error(c, "tx err");
+    if (c->is_draining && c->send.len == 0 && flush == 0) c->is_closing = 1;
     if (c->is_closing) close_conn(c);
   }
 }

src/tls.h

@@ -81,7 +81,7 @@ void mg_tls_free(struct mg_connection *);
 long mg_tls_send(struct mg_connection *, const void *buf, size_t len);
 long mg_tls_recv(struct mg_connection *, void *buf, size_t len);
 size_t mg_tls_pending(struct mg_connection *);
-void mg_tls_flush(struct mg_connection *);
+long mg_tls_flush(struct mg_connection *);
 void mg_tls_handshake(struct mg_connection *);
 void mg_tls_ctx_init(struct mg_mgr *);
 void mg_tls_ctx_free(struct mg_mgr *);

src/tls_builtin.c

@@ -2846,13 +2846,16 @@ size_t mg_tls_pending(struct mg_connection *c) {
   return tls != NULL ? tls->recv_len : 0;
 }
 
-void mg_tls_flush(struct mg_connection *c) {
+long mg_tls_flush(struct mg_connection *c) {
   struct tls_data *tls = (struct tls_data *) c->tls;
-  long n;
+  long n = 0;
   while (tls->send.len > 0 &&
          (n = mg_io_send(c, tls->send.buf, tls->send.len)) > 0) {
     mg_iobuf_del(&tls->send, 0, (size_t) n);
   }
+  if (n == MG_IO_ERR) return n;
+  c->is_tls_throttled = tls->send.len > 0;
+  return c->is_tls_throttled ? MG_IO_WAIT : 0;
 }
 
 void mg_tls_ctx_init(struct mg_mgr *mgr) {

src/tls_dummy.c

@@ -21,8 +21,9 @@ size_t mg_tls_pending(struct mg_connection *c) {
   (void) c;
   return 0;
 }
-void mg_tls_flush(struct mg_connection *c) {
+long mg_tls_flush(struct mg_connection *c) {
   (void) c;
+  return 0;
 }
 void mg_tls_ctx_init(struct mg_mgr *mgr) {
   (void) mgr;

src/tls_mbed.c

@@ -249,17 +249,20 @@ long mg_tls_send(struct mg_connection *c, const void *buf, size_t len) {
   return n;
 }
 
-void mg_tls_flush(struct mg_connection *c) {
+long mg_tls_flush(struct mg_connection *c) {
   struct mg_tls *tls = (struct mg_tls *) c->tls;
+  long n = 0;
   if (c->is_tls_throttled && c->is_draining) {
-    long n =
-        mbedtls_ssl_write(&tls->ssl, tls->throttled_buf, tls->throttled_len);
+    n = mbedtls_ssl_write(&tls->ssl, tls->throttled_buf, tls->throttled_len);
 #if defined(MBEDTLS_ERR_SSL_RECEIVED_NEW_SESSION_TICKET)
-    if (n == MBEDTLS_ERR_SSL_RECEIVED_NEW_SESSION_TICKET) return;
+    if (n == MBEDTLS_ERR_SSL_RECEIVED_NEW_SESSION_TICKET) return MG_IO_WAIT;
 #endif
     c->is_tls_throttled =
         (n == MBEDTLS_ERR_SSL_WANT_READ || n == MBEDTLS_ERR_SSL_WANT_WRITE);
+    if (c->is_tls_throttled) return MG_IO_WAIT;
+    if (n <= 0) return MG_IO_ERR;
   }
+  return 0;
 }
 
 void mg_tls_ctx_init(struct mg_mgr *mgr) {

src/tls_openssl.c

@@ -335,8 +335,9 @@ long mg_tls_send(struct mg_connection *c, const void *buf, size_t len) {
   return n;
 }
 
-void mg_tls_flush(struct mg_connection *c) {
+long mg_tls_flush(struct mg_connection *c) {
   (void) c;
+  return 0;
 }
 
 void mg_tls_ctx_init(struct mg_mgr *mgr) {

test/mip_test.c

@@ -119,6 +119,33 @@ static void tcpclosure_fn(struct mg_connection *c, int ev, void *ev_data) {
   (void) c, (void) ev_data;
 }
 
+static struct mg_connection *s_drain_conn;
+static bool s_drain_tls;
+
+static void tcpdrain_fn(struct mg_connection *c, int ev, void *ev_data) {
+  static const char data[2 * TCP_TEST_WIN] = {0};
+  if (ev == MG_EV_ACCEPT) {
+    s_drain_conn = c;
+#if MG_TLS == MG_TLS_BUILTIN
+    if (s_drain_tls) {
+      struct mg_tls_opts opts;
+      struct tls_data *tls;
+
+      memset(&opts, 0, sizeof(opts));
+      mg_tls_init(c, &opts);
+      tls = (struct tls_data *) c->tls;
+      ASSERT(tls != NULL);
+      ASSERT(mg_iobuf_add(&tls->send, 0, data, sizeof(data)) == sizeof(data));
+    } else
+#endif
+      mg_send(c, data, sizeof(data));
+    c->is_draining = 1;
+  } else if (ev == MG_EV_CLOSE && c == s_drain_conn) {
+    s_drain_conn = NULL;
+  }
+  (void) ev_data;
+}
+
 static void txwindow_fn(struct mg_connection *c, int ev, void *ev_data) {
   if (ev == MG_EV_ACCEPT) {
     char bigdata[2 * TCP_TEST_WIN + 256];
@@ -1024,6 +1051,41 @@ static void test_tcp_ackseq(void) {
   mg_mgr_free(&mgr);
 }
 
+static void test_tcp_drain(bool tls) {
+  struct mg_mgr mgr;
+  struct eth e;
+  struct ip ip;
+  struct ipp ipp;
+  struct tcp *t = (struct tcp *) (s_driver_data.buf + sizeof(e) + sizeof(ip));
+  struct connstate *s;
+  struct mg_tcpip_driver driver;
+  struct mg_tcpip_if mif;
+
+  ipp.ip4 = &ip;
+  ipp.ip6 = NULL;
+  s_drain_conn = NULL;
+  s_drain_tls = tls;
+  init_tcp_tests(&mgr, &e, &ipp, &driver, &mif, tcpdrain_fn);
+  init_tcp_handshake(&e, &ipp, &mgr);  // Send until peer window closes
+  ASSERT(s_drain_conn != NULL);
+  while (received_response(&s_driver_data)) {
+    s_driver_data.len = 0;  // Do not feed Mongoose its own TX packet
+    mg_mgr_poll(&mgr, 0);
+  }
+  ASSERT((t->flags & TH_FIN) == 0);
+
+  ASSERT(s_drain_conn != NULL);
+  s = (struct connstate *) (s_drain_conn + 1);
+  create_tcp_simpleseg(&e, &ipp, 1001, s->seq, TH_ACK, 0);  // Reopen window
+  do {
+    while (!received_response(&s_driver_data)) mg_mgr_poll(&mgr, 0);
+  } while (t->flags != (TH_FIN | TH_ACK));
+  ASSERT(t->flags == (TH_FIN | TH_ACK));
+
+  s_driver_data.len = 0;
+  mg_mgr_free(&mgr);
+}
+
 static void test_frag_recv_path(void) {
   struct mg_mgr mgr;
   struct eth e;
@@ -1186,6 +1248,10 @@ static void test_tcp(bool ipv6) {
     test_tcp_retransmit();
     test_tcp_txwindow();
     test_tcp_ackseq();
+    test_tcp_drain(false);
+#if MG_TLS == MG_TLS_BUILTIN
+    test_tcp_drain(true);
+#endif
   }
 }