add a function that will reject all pending incoming connections (twi…

…tter#154) * add a function that will reject all pending incoming connections * add some comment about tcp_reject
slyphon · Jul 3, 2018 · a2d5c01 · a2d5c01
1 parent dd56abd
commit a2d5c01
Show file tree

Hide file tree

Showing 3 changed files with 51 additions and 6 deletions.
diff --git a/include/channel/cc_tcp.h b/include/channel/cc_tcp.h
@@ -128,7 +128,8 @@ ssize_t tcp_recvv(struct tcp_conn *c, struct array *bufv, size_t nbyte);
 ssize_t tcp_sendv(struct tcp_conn *c, struct array *bufv, size_t nbyte);
 
 bool tcp_accept(struct tcp_conn *sc, struct tcp_conn *c);   /* channel_accept_fn */
-void tcp_reject(struct tcp_conn *sc);                   /* channel_reject_fn */
+void tcp_reject(struct tcp_conn *sc);                       /* channel_reject_fn */
+void tcp_reject_all(struct tcp_conn *sc);                   /* channel_reject_fn */
 
 /* functions getting/setting connection attribute */
 int tcp_set_blocking(int sd);

diff --git a/src/channel/cc_tcp.c b/src/channel/cc_tcp.c
@@ -377,6 +377,11 @@ tcp_accept(struct tcp_conn *sc, struct tcp_conn *c)
     return true;
 }
 
+
+/*
+ * due to lack of a direct rejection API in POSIX, tcp_reject accepts the
+ * frontmost connection and immediately closes it
+ */
 void
 tcp_reject(struct tcp_conn *sc)
 {
@@ -397,6 +402,45 @@ tcp_reject(struct tcp_conn *sc)
     }
 }
 
+/*
+ * due to lack of a direct rejection API in POSIX, tcp_reject_all accepts
+ * connections ready on the listening socket, and immediately closes them.
+ * It does so until there are no more pending connections.
+ */
+void
+tcp_reject_all(struct tcp_conn *sc)
+{
+    int ret;
+    int sd;
+
+    for (;;) {
+        sd = accept(sc->sd, NULL, NULL);
+        if (sd < 0) {
+            if (errno == EINTR) {
+                log_debug("sd %d not ready: eintr", sc->sd);
+                continue;
+            }
+
+            if (errno == EAGAIN || errno == EWOULDBLOCK) {
+                log_debug("sd %d has no more outstanding connections", sc->sd);
+                return;
+            }
+
+            log_error("accept on sd %d failed: %s", sc->sd, strerror(errno));
+            INCR(tcp_metrics, tcp_reject_ex);
+            return;
+        }
+
+        ret = close(sd);
+        if (ret < 0) {
+            INCR(tcp_metrics, tcp_reject_ex);
+            log_warn("close c %d failed, ignored: %s", sd, strerror(errno));
+        }
+
+        INCR(tcp_metrics, tcp_reject);
+    }
+}
+
 int
 tcp_set_blocking(int sd)
 {

diff --git a/test/channel/tcp/check_tcp.c b/test/channel/tcp/check_tcp.c
@@ -72,7 +72,7 @@ find_port_listen(struct tcp_conn **_conn_listen, struct addrinfo **_ai, uint16_t
         freeaddrinfo(ai);
     }
     /* for some reason this line is needed, I would appreciate some insight */
-    ck_assert_int_eq(tcp_connect(ai, conn_client), true);
+    ck_assert(tcp_connect(ai, conn_client));
     tcp_reject(conn_listen);
 
     if (_conn_listen) {
@@ -157,7 +157,7 @@ START_TEST(test_client_send_server_recv)
     conn_server = tcp_conn_create();
     ck_assert_ptr_ne(conn_server, NULL);
 
-    ck_assert_int_eq(tcp_accept(conn_listen, conn_server), true);
+    ck_assert(tcp_accept(conn_listen, conn_server));
     ck_assert_int_eq(tcp_send(conn_client, send_data, LEN), LEN);
     while ((recv = tcp_recv(conn_server, recv_data, LEN + 1)) == CC_EAGAIN) {}
     ck_assert_int_eq(recv, LEN);
@@ -199,7 +199,7 @@ START_TEST(test_server_send_client_recv)
     conn_server = tcp_conn_create();
     ck_assert_ptr_ne(conn_server, NULL);
 
-    ck_assert_int_eq(tcp_accept(conn_listen, conn_server), true);
+    ck_assert(tcp_accept(conn_listen, conn_server));
     ck_assert_int_eq(tcp_send(conn_server, send_data, LEN), LEN);
     while ((recv = tcp_recv(conn_client, recv_data, LEN + 1)) == CC_EAGAIN) {}
     ck_assert_int_eq(recv, LEN);
@@ -252,7 +252,7 @@ START_TEST(test_client_sendv_server_recvv)
     conn_server = tcp_conn_create();
     ck_assert_ptr_ne(conn_server, NULL);
 
-    ck_assert_int_eq(tcp_accept(conn_listen, conn_server), true);
+    ck_assert(tcp_accept(conn_listen, conn_server));
     ck_assert_int_eq(tcp_sendv(conn_client, send_array, LEN), LEN);
     while ((recv = tcp_recvv(conn_server, recv_array, LEN + 1)) == CC_EAGAIN) {}
     ck_assert_int_eq(recv, LEN);
@@ -315,7 +315,7 @@ START_TEST(test_nonblocking)
     conn_server = tcp_conn_create();
     ck_assert_ptr_ne(conn_server, NULL);
 
-    ck_assert_int_eq(tcp_accept(conn_listen, conn_server), true);
+    ck_assert(tcp_accept(conn_listen, conn_server));
 
     task.usleep = SLEEP_TIME;
     task.c = conn_server;