From cda2aac71b644c8c72c468e3afc6af7ad64151b7 Mon Sep 17 00:00:00 2001
From: Jin-Seong Kim <jseong82.kim@samsung.com>
Date: Thu, 15 Mar 2018 19:22:39 +0900
Subject: [PATCH] introduce LWIP_TCP_MAX_TW_BUCKETS to limit timewait socket

---
 src/apps/snmp/snmp_mib2_tcp.c    |  4 ++--
 src/core/tcp.c                   | 44 +++++++++++++++++++++++++++++-----------
 src/core/tcp_in.c                | 17 ++++++++++++----
 src/core/timeouts.c              |  4 ++--
 src/include/lwip/opt.h           | 13 ++++++++++++
 src/include/lwip/priv/tcp_priv.h | 14 ++++++++++++-
 test/unit/api/test_sockets.c     |  4 ++--
 test/unit/tcp/tcp_helper.c       |  4 ++--
 8 files changed, 79 insertions(+), 25 deletions(-)

diff --git a/src/apps/snmp/snmp_mib2_tcp.c b/src/apps/snmp/snmp_mib2_tcp.c
index 8a321f3..1be4ded 100644
--- a/src/apps/snmp/snmp_mib2_tcp.c
+++ b/src/apps/snmp/snmp_mib2_tcp.c
@@ -334,7 +334,7 @@ tcp_ConnectionTable_get_cell_value(const u32_t *column, const u32_t *row_oid, u8
   struct tcp_pcb *pcb;
   u8_t idx = 0;
   u8_t i;
-  struct tcp_pcb **const tcp_pcb_nonlisten_lists[] = {&tcp_bound_pcbs, &tcp_active_pcbs, &tcp_tw_pcbs};
+  struct tcp_pcb ** const tcp_pcb_nonlisten_lists[] = {&tcp_bound_pcbs, &tcp_active_pcbs, &tcp_tw_pcbs.pcbs};
 
   LWIP_UNUSED_ARG(value_len);
 
@@ -379,7 +379,7 @@ tcp_ConnectionTable_get_next_cell_instance_and_value(const u32_t *column, struct
    * 1x tcpConnectionRemAddressType   + 1x OID len + 16x tcpConnectionRemAddress    + 1x tcpConnectionRemPort */
   u32_t  result_temp[38];
   u8_t i;
-  struct tcp_pcb **const tcp_pcb_nonlisten_lists[] = {&tcp_bound_pcbs, &tcp_active_pcbs, &tcp_tw_pcbs};
+  struct tcp_pcb ** const tcp_pcb_nonlisten_lists[] = {&tcp_bound_pcbs, &tcp_active_pcbs, &tcp_tw_pcbs.pcbs};
 
   LWIP_UNUSED_ARG(value_len);
 
diff --git a/src/core/tcp.c b/src/core/tcp.c
index f38f7a1..1bce379 100644
--- a/src/core/tcp.c
+++ b/src/core/tcp.c
@@ -175,12 +175,11 @@ union tcp_listen_pcbs_t tcp_listen_pcbs;
  * they accept or send data. */
 struct tcp_pcb *tcp_active_pcbs;
 /** List of all TCP PCBs in TIME-WAIT state */
-struct tcp_pcb *tcp_tw_pcbs;
+struct tcp_tw_pcbs_t tcp_tw_pcbs;
 
 /** An array with all (non-temporary) PCB lists, mainly used for smaller code size */
-struct tcp_pcb **const tcp_pcb_lists[] = {&tcp_listen_pcbs.pcbs, &tcp_bound_pcbs,
-         &tcp_active_pcbs, &tcp_tw_pcbs
-};
+struct tcp_pcb ** const tcp_pcb_lists[] = {&tcp_listen_pcbs.pcbs, &tcp_bound_pcbs,
+  &tcp_active_pcbs, &tcp_tw_pcbs.pcbs};
 
 u8_t tcp_active_pcbs_changed;
 
@@ -565,7 +564,7 @@ tcp_abandon(struct tcp_pcb *pcb, int reset)
      are in an active state, call the receive function associated with
      the PCB with a NULL argument, and send an RST to the remote end. */
   if (pcb->state == TIME_WAIT) {
-    tcp_pcb_remove(&tcp_tw_pcbs, pcb);
+    TCP_RMV_TIMEWAIT(pcb);
     tcp_free(pcb);
   } else {
     int send_rst = 0;
@@ -1408,7 +1407,7 @@ tcp_slowtmr_start:
 
   /* Steps through all of the TIME-WAIT PCBs. */
   prev = NULL;
-  pcb = tcp_tw_pcbs;
+  pcb = tcp_tw_pcbs.pcbs;
   while (pcb != NULL) {
     LWIP_ASSERT("tcp_slowtmr: TIME-WAIT pcb->state == TIME-WAIT", pcb->state == TIME_WAIT);
     pcb_remove = 0;
@@ -1424,12 +1423,12 @@ tcp_slowtmr_start:
       tcp_pcb_purge(pcb);
       /* Remove PCB from tcp_tw_pcbs list. */
       if (prev != NULL) {
-        LWIP_ASSERT("tcp_slowtmr: middle tcp != tcp_tw_pcbs", pcb != tcp_tw_pcbs);
+        LWIP_ASSERT("tcp_slowtmr: middle tcp != tcp_tw_pcbs", pcb != tcp_tw_pcbs.pcbs);
         prev->next = pcb->next;
       } else {
         /* This PCB was the first. */
-        LWIP_ASSERT("tcp_slowtmr: first pcb == tcp_tw_pcbs", tcp_tw_pcbs == pcb);
-        tcp_tw_pcbs = pcb->next;
+        LWIP_ASSERT("tcp_slowtmr: first pcb == tcp_tw_pcbs", tcp_tw_pcbs.pcbs == pcb);
+        tcp_tw_pcbs.pcbs = pcb->next;
       }
       pcb2 = pcb;
       pcb = pcb->next;
@@ -1748,7 +1747,7 @@ tcp_kill_timewait(void)
   inactivity = 0;
   inactive = NULL;
   /* Go through the list of TIME_WAIT pcbs and get the oldest pcb. */
-  for (pcb = tcp_tw_pcbs; pcb != NULL; pcb = pcb->next) {
+  for (pcb = tcp_tw_pcbs.pcbs; pcb != NULL; pcb = pcb->next) {
     if ((u32_t)(tcp_ticks - pcb->tmr) >= inactivity) {
       inactivity = tcp_ticks - pcb->tmr;
       inactive = pcb;
@@ -2157,6 +2156,27 @@ tcp_pcb_remove(struct tcp_pcb **pcblist, struct tcp_pcb *pcb)
   LWIP_ASSERT("tcp_pcb_remove: tcp_pcbs_sane()", tcp_pcbs_sane());
 }
 
+err_t tcp_tw_pcb_register(struct tcp_pcb *pcb)
+{
+#if LWIP_TCP_MAX_TW_BUCKETS > 0
+  if (tcp_tw_pcbs.counter >= LWIP_TCP_MAX_TW_BUCKETS) {
+    LWIP_DEBUGF(TCP_DEBUG, ("exceed acceptable number of timewait pcbs " SZT_F "\n", tcp_tw_pcbs.counter));
+    return ERR_MEM;
+  }
+  ++tcp_tw_pcbs.counter;
+#endif
+  TCP_REG(&tcp_tw_pcbs.pcbs, pcb);
+
+  return ERR_OK;
+}
+
+void tcp_tw_pcb_remove(struct tcp_pcb *pcb)
+{
+#if LWIP_TCP_MAX_TW_BUCKETS > 0
+  tcp_tw_pcbs.counter = (tcp_tw_pcbs.counter > 0) ? (tcp_tw_pcbs.counter - 1) : 0;
+#endif
+  tcp_pcb_remove(&tcp_tw_pcbs.pcbs, pcb);
+}
 /**
  * Calculates a new initial sequence number for new connections.
  *
@@ -2446,7 +2466,7 @@ tcp_debug_print_pcbs(void)
   }
 
   LWIP_DEBUGF(TCP_DEBUG, ("TIME-WAIT PCB states:\n"));
-  for (pcb = tcp_tw_pcbs; pcb != NULL; pcb = pcb->next) {
+  for (pcb = tcp_tw_pcbs.pcbs; pcb != NULL; pcb = pcb->next) {
     LWIP_DEBUGF(TCP_DEBUG, ("Local port %"U16_F", foreign port %"U16_F" snd_nxt %"U32_F" rcv_nxt %"U32_F" ",
                             pcb->local_port, pcb->remote_port,
                             pcb->snd_nxt, pcb->rcv_nxt));
@@ -2466,7 +2486,7 @@ tcp_pcbs_sane(void)
     LWIP_ASSERT("tcp_pcbs_sane: active pcb->state != LISTEN", pcb->state != LISTEN);
     LWIP_ASSERT("tcp_pcbs_sane: active pcb->state != TIME-WAIT", pcb->state != TIME_WAIT);
   }
-  for (pcb = tcp_tw_pcbs; pcb != NULL; pcb = pcb->next) {
+  for (pcb = tcp_tw_pcbs.pcbs; pcb != NULL; pcb = pcb->next) {
     LWIP_ASSERT("tcp_pcbs_sane: tw pcb->state == TIME-WAIT", pcb->state == TIME_WAIT);
   }
   return 1;
diff --git a/src/core/tcp_in.c b/src/core/tcp_in.c
index f1d3a11..6cfd1d8 100644
--- a/src/core/tcp_in.c
+++ b/src/core/tcp_in.c
@@ -282,7 +282,7 @@ tcp_input(struct pbuf *p, struct netif *inp)
   if (pcb == NULL) {
     /* If it did not go to an active connection, we check the connections
        in the TIME-WAIT state. */
-    for (pcb = tcp_tw_pcbs; pcb != NULL; pcb = pcb->next) {
+    for (pcb = tcp_tw_pcbs.pcbs; pcb != NULL; pcb = pcb->next) {
       LWIP_ASSERT("tcp_input: TIME-WAIT pcb->state == TIME-WAIT", pcb->state == TIME_WAIT);
 
       /* check if PCB is bound to specific netif */
@@ -986,7 +986,10 @@ tcp_process(struct tcp_pcb *pcb)
           tcp_pcb_purge(pcb);
           TCP_RMV_ACTIVE(pcb);
           pcb->state = TIME_WAIT;
-          TCP_REG(&tcp_tw_pcbs, pcb);
+          if (TCP_REG_TIMEWAIT(pcb) != ERR_OK) {
+            tcp_free(pcb);
+            return ERR_ABRT;
+          }
         } else {
           tcp_ack_now(pcb);
           pcb->state = CLOSING;
@@ -1004,7 +1007,10 @@ tcp_process(struct tcp_pcb *pcb)
         tcp_pcb_purge(pcb);
         TCP_RMV_ACTIVE(pcb);
         pcb->state = TIME_WAIT;
-        TCP_REG(&tcp_tw_pcbs, pcb);
+        if (TCP_REG_TIMEWAIT(pcb) != ERR_OK) {
+          tcp_free(pcb);
+          return ERR_ABRT;
+        }
       }
       break;
     case CLOSING:
@@ -1014,7 +1020,10 @@ tcp_process(struct tcp_pcb *pcb)
         tcp_pcb_purge(pcb);
         TCP_RMV_ACTIVE(pcb);
         pcb->state = TIME_WAIT;
-        TCP_REG(&tcp_tw_pcbs, pcb);
+        if (TCP_REG_TIMEWAIT(pcb) != ERR_OK) {
+          tcp_free(pcb);
+          return ERR_ABRT;
+        }
       }
       break;
     case LAST_ACK:
diff --git a/src/core/timeouts.c b/src/core/timeouts.c
index f37acfe..88d1de8 100644
--- a/src/core/timeouts.c
+++ b/src/core/timeouts.c
@@ -148,7 +148,7 @@ tcpip_tcp_timer(void *arg)
   /* call TCP timer handler */
   tcp_tmr();
   /* timer still needed? */
-  if (tcp_active_pcbs || tcp_tw_pcbs) {
+  if (tcp_active_pcbs || (tcp_tw_pcbs.pcbs)) {
     /* restart timer */
     sys_timeout(TCP_TMR_INTERVAL, tcpip_tcp_timer, NULL);
   } else {
@@ -168,7 +168,7 @@ tcp_timer_needed(void)
   LWIP_ASSERT_CORE_LOCKED();
 
   /* timer is off but needed again? */
-  if (!tcpip_tcp_timer_active && (tcp_active_pcbs || tcp_tw_pcbs)) {
+  if (!tcpip_tcp_timer_active && (tcp_active_pcbs || (tcp_tw_pcbs.pcbs))) {
     /* enable and start timer */
     tcpip_tcp_timer_active = 1;
     sys_timeout(TCP_TMR_INTERVAL, tcpip_tcp_timer, NULL);
diff --git a/src/include/lwip/opt.h b/src/include/lwip/opt.h
index c4144c2..4777975 100644
--- a/src/include/lwip/opt.h
+++ b/src/include/lwip/opt.h
@@ -1509,6 +1509,19 @@
 #endif
 
 /**
+ * LWIP_TCP_MAX_TW_BUCKETS
+ * Set maximum number of tw (TIMED WAIT) pcbs in tcp_tw_tpcbs.
+ * It can be useful in case of memory constrined network devices
+ * to reduce memory usage due to TCP sessions
+ * When LWIP_TCP_MAX_TW_BUCKETS is 0, we can add tw_pcbs without limitation
+ *
+ */
+
+#ifndef LWIP_TCP_MAX_TW_BUCKETS
+#define LWIP_TCP_MAX_TW_BUCKETS		0
+#endif
+
+/**
  * @}
  */
 
diff --git a/src/include/lwip/priv/tcp_priv.h b/src/include/lwip/priv/tcp_priv.h
index ee7f9b3..99356cc 100644
--- a/src/include/lwip/priv/tcp_priv.h
+++ b/src/include/lwip/priv/tcp_priv.h
@@ -332,12 +332,17 @@ union tcp_listen_pcbs_t { /* List of all TCP PCBs in LISTEN state. */
   struct tcp_pcb_listen *listen_pcbs;
   struct tcp_pcb *pcbs;
 };
+struct tcp_tw_pcbs_t {
+  u8_t counter;          /** The number of TCP timewait pcbs to
+                             check maximum number of timewait pcbs */
+  struct tcp_pcb *pcbs;  /** List of all pcbs in time-wait state */
+};
 extern struct tcp_pcb *tcp_bound_pcbs;
 extern union tcp_listen_pcbs_t tcp_listen_pcbs;
 extern struct tcp_pcb *tcp_active_pcbs;  /* List of all TCP PCBs that are in a
               state in which they accept or send
               data. */
-extern struct tcp_pcb *tcp_tw_pcbs;      /* List of all TCP PCBs in TIME-WAIT. */
+extern struct tcp_tw_pcbs_t tcp_tw_pcbs;	/* List of all TCP PCBs in TIME-WAIT. */
 
 #define NUM_TCP_PCB_LISTS_NO_TIME_WAIT  3
 #define NUM_TCP_PCB_LISTS               4
@@ -435,11 +440,18 @@ extern struct tcp_pcb ** const tcp_pcb_lists[NUM_TCP_PCB_LISTS];
     tcp_active_pcbs_changed = 1;                   \
   } while (0)
 
+#define TCP_REG_TIMEWAIT(pcb)	\
+    tcp_tw_pcb_register(pcb)	\
+
+#define TCP_RMV_TIMEWAIT(pcb)	\
+    tcp_tw_pcb_remove(pcb)	\
 
 /* Internal functions: */
 struct tcp_pcb *tcp_pcb_copy(struct tcp_pcb *pcb);
 void tcp_pcb_purge(struct tcp_pcb *pcb);
 void tcp_pcb_remove(struct tcp_pcb **pcblist, struct tcp_pcb *pcb);
+err_t tcp_tw_pcb_register(struct tcp_pcb *pcb);
+void tcp_tw_pcb_remove(struct tcp_pcb *pcb);
 
 void tcp_segs_free(struct tcp_seg *seg);
 void tcp_seg_free(struct tcp_seg *seg);
diff --git a/test/unit/api/test_sockets.c b/test/unit/api/test_sockets.c
index 472fa48..6015fb9 100644
--- a/test/unit/api/test_sockets.c
+++ b/test/unit/api/test_sockets.c
@@ -44,8 +44,8 @@ sockets_teardown(void)
   fail_unless(test_sockets_get_used_count() == 0);
   /* poll until all memory is released... */
   tcpip_thread_poll_one();
-  while (tcp_tw_pcbs) {
-    tcp_abort(tcp_tw_pcbs);
+  while (tcp_tw_pcbs.pcbs) {
+    tcp_abort(tcp_tw_pcbs.pcbs);
     tcpip_thread_poll_one();
   }
   tcpip_thread_poll_one();
diff --git a/test/unit/tcp/tcp_helper.c b/test/unit/tcp/tcp_helper.c
index 8689b71..e3a7ec2 100644
--- a/test/unit/tcp/tcp_helper.c
+++ b/test/unit/tcp/tcp_helper.c
@@ -35,7 +35,7 @@ tcp_remove_all(void)
   tcp_remove(tcp_listen_pcbs.pcbs);
   tcp_remove(tcp_bound_pcbs);
   tcp_remove(tcp_active_pcbs);
-  tcp_remove(tcp_tw_pcbs);
+  tcp_remove(tcp_tw_pcbs.pcbs);
   fail_unless(MEMP_STATS_GET(used, MEMP_TCP_PCB) == 0);
   fail_unless(MEMP_STATS_GET(used, MEMP_TCP_PCB_LISTEN) == 0);
   fail_unless(MEMP_STATS_GET(used, MEMP_TCP_SEG) == 0);
@@ -169,7 +169,7 @@ tcp_set_state(struct tcp_pcb* pcb, enum tcp_state state, const ip_addr_t* local_
     ip_addr_copy(pcb->local_ip, *local_ip);
     pcb->local_port = local_port;
   } else if(state == TIME_WAIT) {
-    TCP_REG(&tcp_tw_pcbs, pcb);
+    TCP_REG(&tcp_tw_pcbs.pcbs, pcb);
     ip_addr_copy(pcb->local_ip, *local_ip);
     pcb->local_port = local_port;
     ip_addr_copy(pcb->remote_ip, *remote_ip);
-- 
1.9.1

