Merge pull request #300 from skaiea13-ai/codex/nstun-aggregate-budget

Limit nstun TCP receive buffers with a shared payload budget

Makefile

@@ -45,6 +45,7 @@ endif
 
 BIN = nsjail
 LIBS = kafel/libkafel.a
+TEST_BIN = tests/nstun_buffer_budget_test
 
 # If PASTA_BIN_PATH is not provided in env, dynamically search for it if EMBED_PASTA is requested
 # or fallback to it naturally.
@@ -109,7 +110,7 @@ kafel/libkafel.a: kafel_init
 
 # Utilities
 clean:
-	$(RM) core Makefile.bak $(OBJS) $(SRCS_PB_CXX) $(SRCS_PB_H) $(SRCS_PB_O) $(BIN)
+	$(RM) core Makefile.bak $(OBJS) $(SRCS_PB_CXX) $(SRCS_PB_H) $(SRCS_PB_O) $(BIN) $(TEST_BIN)
 ifneq ("$(wildcard kafel/Makefile)","")
 	+$(MAKE) -C kafel clean
 endif
@@ -148,7 +149,8 @@ NEW_EF := --experimental_mnt=new
 UID := $(shell id -u)
 
 .PHONY: test
-test: $(BIN)
+test: $(BIN) $(TEST_BIN)
+	$(call run_test, ./$(TEST_BIN), 0)
 	# --- Basic sanity tests ---
 	$(call run_test, ./nsjail -q -Mo --chroot / --user 99999 --group 99999 -- /bin/true, 0)
 	$(call run_test, ./nsjail -q -Mo --chroot / --user 99999 --group 99999 -- /bin/false, 1)
@@ -238,6 +240,9 @@ endif
 	@echo "========================================"
 	@echo ""
 
+$(TEST_BIN): tests/nstun_buffer_budget_test.cc nstun/buffer_budget.h
+	$(CXX) $(filter-out -c,$(CXXFLAGS)) $< -o $@
+
 # Dependencies (Generated by makedepend)
 # DO NOT DELETE THIS LINE -- make depend depends on it.
 

nstun/buffer_budget.h

@@ -0,0 +1,38 @@
+#ifndef NSTUN_BUFFER_BUDGET_H_
+#define NSTUN_BUFFER_BUDGET_H_
+
+#include <stddef.h>
+
+namespace nstun {
+
+class BufferBudget {
+public:
+	explicit constexpr BufferBudget(size_t limit) : limit_(limit) {}
+
+	[[nodiscard]] bool try_reserve(size_t bytes) {
+		if (bytes > limit_ - used_) {
+			return false;
+		}
+		used_ += bytes;
+		return true;
+	}
+
+	[[nodiscard]] bool release(size_t bytes) {
+		if (bytes > used_) {
+			return false;
+		}
+		used_ -= bytes;
+		return true;
+	}
+
+	[[nodiscard]] size_t used() const { return used_; }
+	[[nodiscard]] size_t limit() const { return limit_; }
+
+private:
+	size_t limit_;
+	size_t used_ = 0;
+};
+
+} /* namespace nstun */
+
+#endif /* NSTUN_BUFFER_BUDGET_H_ */

nstun/core.h

@@ -12,12 +12,14 @@
 #include <memory>
 #include <vector>
 
+#include "buffer_budget.h"
 #include "net_defs.h"
 #include "nstun.h"
 
 namespace nstun {
 
 constexpr size_t NSTUN_MAX_FLOWS = 1024;
+constexpr size_t NSTUN_MAX_TCP_RX_BUFFERED_BYTES = 64 * 1024 * 1024;
 
 // Removed MemcmpLess in favor of C++20 operator<=>
 
@@ -166,6 +168,7 @@ struct Context {
 
 	/* Unified host mapping for all encapsulated flows */
 	std::map<int, Flow*> flows_by_fd; // Observer pointer
+	BufferBudget tcp_rx_buffer_budget{NSTUN_MAX_TCP_RX_BUFFERED_BYTES};
 
 	/* IPv6 maps (Owning) */
 	std::map<FlowKey6, std::unique_ptr<UdpFlow>> ipv6_udp_flows_by_key;

nstun/tcp.cc

@@ -232,7 +232,27 @@ static void tcp_rst_and_destroy(Context* ctx, TcpFlow* flow) {
 	tcp_destroy_flow(ctx, flow);
 }
 
+static void release_tcp_rx_buffer(Context* ctx, TcpFlow* flow) {
+	size_t bytes = flow->rx_buffer.size();
+	if (bytes == 0) {
+		flow->rx_sent_offset = 0;
+		return;
+	}
+	if (!ctx->tcp_rx_buffer_budget.release(bytes)) {
+		LOG_F("TCP rx buffer budget accounting underflow: used=%zu release=%zu",
+		    ctx->tcp_rx_buffer_budget.used(), bytes);
+		abort();
+	}
+	flow->rx_buffer.clear();
+	flow->rx_sent_offset = 0;
+	/* Keep normal-flow allocations reusable, but release large retained buffers. */
+	if (flow->rx_buffer.capacity() > TCP_RECV_BUF_SIZE) {
+		std::vector<uint8_t>().swap(flow->rx_buffer);
+	}
+}
+
 void tcp_destroy_flow(Context* ctx, TcpFlow* flow) {
+	release_tcp_rx_buffer(ctx, flow);
 	if (flow->host_fd != -1) {
 		epoll_ctl(ctx->epoll_fd, EPOLL_CTL_DEL, flow->host_fd, nullptr);
 		ctx->flows_by_fd.erase(flow->host_fd);
@@ -304,8 +324,7 @@ bool flush_to_host(Context* ctx, TcpFlow* flow) {
 	if (written > 0) {
 		flow->rx_sent_offset += written;
 		if (flow->rx_sent_offset >= flow->rx_buffer.size()) {
-			flow->rx_buffer.clear();
-			flow->rx_sent_offset = 0;
+			release_tcp_rx_buffer(ctx, flow);
 		}
 
 		/* We made progress, remove EPOLLOUT if empty */
@@ -614,12 +633,18 @@ static void tcp_process_data(Context* ctx, TcpFlow* flow, const tcp_hdr* tcp,
 				const uint8_t* new_data = data + overlap;
 				size_t new_data_len = data_len - overlap;
 
-				if (flow->rx_buffer.size() + new_data_len >
-				    TCP_RX_BUFFER_HARD_CAP) {
+				if (new_data_len >
+				    TCP_RX_BUFFER_HARD_CAP - flow->rx_buffer.size()) {
 					LOG_D("TCP rx_buffer reached 8MB limit (DoS protection), "
 					      "dropping");
 					return;
 				}
+				if (!ctx->tcp_rx_buffer_budget.try_reserve(new_data_len)) {
+					LOG_D("Aggregate TCP rx_buffer reached %zuMB limit (DoS "
+					      "protection), dropping",
+					    ctx->tcp_rx_buffer_budget.limit() / (1024 * 1024));
+					return;
+				}
 
 				flow->seq_from_guest += new_data_len;
 				flow->ack_to_guest = flow->seq_from_guest;

tests/nstun_buffer_budget_test.cc

@@ -0,0 +1,17 @@
+#include <limits>
+
+#include "nstun/buffer_budget.h"
+
+int main() {
+	nstun::BufferBudget budget(64);
+	if (budget.limit() != 64 || budget.used() != 0) return 1;
+	if (!budget.try_reserve(32) || budget.used() != 32) return 2;
+	if (budget.try_reserve(33) || budget.used() != 32) return 3;
+	if (!budget.release(16) || budget.used() != 16) return 4;
+	if (!budget.try_reserve(48) || budget.used() != 64) return 5;
+	if (budget.try_reserve(1) || budget.used() != 64) return 6;
+	if (!budget.release(64) || budget.used() != 0) return 7;
+	if (budget.release(1) || budget.used() != 0) return 8;
+	if (budget.try_reserve(std::numeric_limits<size_t>::max()) || budget.used() != 0) return 9;
+	return 0;
+}