nstun: add missing TCP state - TcpState::CLOSING

Author: robertswiecki

nstun/icmp.cc

@@ -1,5 +1,6 @@
 #include "icmp.h"
 
+#include <errno.h>
 #include <netinet/in.h>
 #include <string.h>
 #include <sys/epoll.h>

nstun/net_defs.h

@@ -102,10 +102,10 @@ inline uint16_t compute_checksum(const void* buf, size_t len, uint32_t sum = 0)
 	return finalize_checksum(compute_checksum_part(buf, len, sum));
 }
 
-}  // namespace nstun
-
 inline bool is_loopback_addr(uint32_t addr_net) {
 	return (ntohl(addr_net) & 0xFF000000) == 0x7F000000;
 }
 
+}  // namespace nstun
+
 #endif /* NSTUN_NET_DEFS_H_ */

nstun/nstun.cc

@@ -35,7 +35,9 @@ namespace nstun {
 
 Context::~Context() {
 	for (auto& pair : udp_flows_by_key) {
-		if (pair.second->host_fd != -1) close(pair.second->host_fd);
+		if (pair.second->host_fd != -1 && !pair.second->host_fd_is_listener) {
+			close(pair.second->host_fd);
+		}
 		if (pair.second->tcp_fd != -1) close(pair.second->tcp_fd);
 		delete pair.second;
 	}
@@ -47,6 +49,9 @@ Context::~Context() {
 		if (pair.second->host_fd != -1) close(pair.second->host_fd);
 		delete pair.second;
 	}
+	for (auto& [fd, _] : host_listener_fd_to_rule) {
+		close(fd);
+	}
 }
 
 RuleResult evaluate_rules(Context* ctx, nstun_direction_t dir, nstun_proto_t proto, uint32_t src_ip,
@@ -83,7 +88,9 @@ static void garbage_collect(Context* ctx) {
 			timeout = 5;
 		} else if (flow->state == TcpState::CLOSE_WAIT ||
 			   flow->state == TcpState::LAST_ACK ||
-			   flow->state == TcpState::TIME_WAIT) {
+			   flow->state == TcpState::FIN_WAIT_1 ||
+			   flow->state == TcpState::FIN_WAIT_2 ||
+			   flow->state == TcpState::CLOSING || flow->state == TcpState::TIME_WAIT) {
 			timeout = 10;
 		}
 		if (now - flow->last_active > timeout) {

nstun/tcp.cc

@@ -33,27 +33,34 @@ void tcp_send_packet(Context* ctx, TcpFlow* flow, uint8_t flags, const uint8_t*
 		return;
 	}
 
+	constexpr uint8_t TCP_OPT_NOP = 1;
+	constexpr uint8_t TCP_OPT_MSS = 2;
+	constexpr uint8_t TCP_OPT_MSS_LEN = 4;
+	constexpr uint8_t TCP_OPT_WSCALE = 3;
+	constexpr uint8_t TCP_OPT_WSCALE_LEN = 3;
+
 	size_t opt_len = 0;
 	uint8_t options[40];
 	if (flags & NSTUN_TCP_FLAG_SYN) {
 		/* Add MSS option (Kind=2, Length=4, MSS=65495) */
-		options[opt_len++] = 2;
-		options[opt_len++] = 4;
+		options[opt_len++] = TCP_OPT_MSS;
+		options[opt_len++] = TCP_OPT_MSS_LEN;
 		uint16_t mss = htons(65495);
 		memcpy(&options[opt_len], &mss, 2);
 		opt_len += 2;
 
-		/* Add NOP (Kind=1) for 32-bit alignment */
-		options[opt_len++] = 1;
+		/* Add NOP for 32-bit alignment */
+		options[opt_len++] = TCP_OPT_NOP;
 
 		/* Add Window Scale option (Kind=3, Length=3, Shift=8) */
-		options[opt_len++] = 3;
-		options[opt_len++] = 3;
+		options[opt_len++] = TCP_OPT_WSCALE;
+		options[opt_len++] = TCP_OPT_WSCALE_LEN;
 		options[opt_len++] = 8;
 	}
 
 	size_t frame_len = sizeof(ip4_hdr) + sizeof(tcp_hdr) + opt_len + len;
-	uint8_t frame_buf[sizeof(ip4_hdr) + sizeof(tcp_hdr) + sizeof(options) + NSTUN_MTU];
+	/* Single-threaded network loop: use static buffer to avoid 63KB stack allocation */
+	static thread_local uint8_t frame_buf[sizeof(ip4_hdr) + sizeof(tcp_hdr) + 40 + NSTUN_MTU];
 
 	ip4_hdr* r_ip = reinterpret_cast<ip4_hdr*>(frame_buf);
 	tcp_hdr* r_tcp = reinterpret_cast<tcp_hdr*>(frame_buf + sizeof(ip4_hdr));
@@ -496,6 +503,10 @@ void handle_tcp(Context* ctx, const ip4_hdr* ip, const uint8_t* payload, size_t
 
 			if (flow->state == TcpState::ESTABLISHED) {
 				flow->state = TcpState::CLOSE_WAIT;
+			} else if (flow->state == TcpState::FIN_WAIT_1) {
+				flow->state = TcpState::CLOSING;
+			} else if (flow->state == TcpState::FIN_WAIT_2) {
+				flow->state = TcpState::TIME_WAIT;
 			}
 
 			push_to_guest(ctx, flow);
@@ -630,7 +641,6 @@ void handle_host_tcp_data(Context* ctx, TcpFlow* flow, int fd) {
 	    (flow->inbound && flow->state == TcpState::SYN_SENT)) {
 		uint8_t buf[65536];
 		ssize_t recv_len = recv(fd, buf, sizeof(buf), 0);
-		LOG_D("recv_len=%zd errno=%d", recv_len, errno);
 		if (recv_len == 0) goto eof;
 		if (recv_len < 0) {
 			if (errno == EAGAIN || errno == EWOULDBLOCK) return;
@@ -757,6 +767,10 @@ void handle_host_tcp_accept(Context* ctx, int listen_fd, const nstun_rule_t& rul
 	}
 
 	LOG_D("Accepted fd=%d", fd);
+
+	int opt = 1;
+	setsockopt(fd, IPPROTO_TCP, TCP_NODELAY, &opt, sizeof(opt));
+
 	struct sockaddr_in server_addr = {};
 	socklen_t servlen = sizeof(server_addr);
 	getsockname(fd, (struct sockaddr*)&server_addr, &servlen);

nstun/tcp.h

@@ -16,6 +16,7 @@ enum class TcpState {
 	ESTABLISHED,
 	FIN_WAIT_1,
 	FIN_WAIT_2,
+	CLOSING,
 	TIME_WAIT,
 	CLOSE_WAIT,
 	LAST_ACK

nstun/udp.cc

@@ -1,6 +1,7 @@
 #include "udp.h"
 
 #include <arpa/inet.h>
+#include <errno.h>
 #include <netinet/in.h>
 #include <string.h>
 #include <sys/epoll.h>