From c3025eb98e93246274d4c8a3df8a8e753b76b7bc Mon Sep 17 00:00:00 2001
From: DeFiDude <59237470+DeFiDude@users.noreply.github.com>
Date: Sat, 28 Mar 2026 15:43:49 -0600
Subject: [PATCH] Fix TCP stack overflow risk and cap WiFi AP client
 connections
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

TCP: Move 2KB stack-allocated packet rewrite buffers in
send_outgoing() to a PSRAM-allocated member (_wrapBuffer).
The ESP32-S3 loop task has 8KB stack — a single 2KB VLA used
25%, risking overflow under deep call chains during TCP load.

WiFi AP: Cap concurrent client connections at 4. Previously
unbounded — any number of clients could connect and exhaust
memory. New clients beyond the limit get a clean TCP RST.
---
 src/transport/TCPClientInterface.cpp | 23 +++++++++++++----------
 src/transport/TCPClientInterface.h   |  3 ++-
 src/transport/WiFiInterface.cpp      |  5 +++++
 src/transport/WiFiInterface.h        |  2 ++
 4 files changed, 22 insertions(+), 11 deletions(-)

diff --git a/src/transport/TCPClientInterface.cpp b/src/transport/TCPClientInterface.cpp
index 404308e..bd7a797 100644
--- a/src/transport/TCPClientInterface.cpp
+++ b/src/transport/TCPClientInterface.cpp
@@ -12,12 +12,15 @@ TCPClientInterface::TCPClientInterface(const char* host, uint16_t port, const ch
     if (!_rxBuffer) _rxBuffer = (uint8_t*)malloc(RX_BUFFER_SIZE);
     _txBuffer = (uint8_t*)ps_malloc(TX_BUFFER_SIZE);
     if (!_txBuffer) _txBuffer = (uint8_t*)malloc(TX_BUFFER_SIZE);
+    _wrapBuffer = (uint8_t*)ps_malloc(RX_BUFFER_SIZE);
+    if (!_wrapBuffer) _wrapBuffer = (uint8_t*)malloc(RX_BUFFER_SIZE);
 }
 
 TCPClientInterface::~TCPClientInterface() {
     stop();
     if (_rxBuffer) { free(_rxBuffer); _rxBuffer = nullptr; }
     if (_txBuffer) { free(_txBuffer); _txBuffer = nullptr; }
+    if (_wrapBuffer) { free(_wrapBuffer); _wrapBuffer = nullptr; }
 }
 
 bool TCPClientInterface::start() {
@@ -160,28 +163,28 @@ void TCPClientInterface::send_outgoing(const RNS::Bytes& data) {
                     _txDropCount++;
                     return;
                 }
-                uint8_t wrapped[RX_BUFFER_SIZE];
-                wrapped[0] = new_flags;
-                wrapped[1] = data.data()[1];  // hops
-                memcpy(wrapped + 2, _hubTransportId, 16);  // transport_id
-                memcpy(wrapped + 18, data.data() + 2, data.size() - 2);  // dest_hash + context + payload
+                if (!_wrapBuffer) return;
+                _wrapBuffer[0] = new_flags;
+                _wrapBuffer[1] = data.data()[1];  // hops
+                memcpy(_wrapBuffer + 2, _hubTransportId, 16);  // transport_id
+                memcpy(_wrapBuffer + 18, data.data() + 2, data.size() - 2);  // dest_hash + context + payload
 
                 Serial.printf("[TCP] TX %d->%d bytes (H1->H2 wrap) to %s:%d\n",
                               (int)data.size(), (int)new_len, _host.c_str(), _port);
-                sendFrame(wrapped, new_len);
+                sendFrame(_wrapBuffer, new_len);
                 InterfaceImpl::handle_outgoing(data);  // Stats use original size
                 return;
             }
             else if (data.size() >= 35 && memcmp(data.data() + 2, _hubTransportId, 16) != 0) {
                 // Header2 with wrong transport_id → fix it
                 // Transport::outbound() may have used _received_from=destination_hash
-                uint8_t fixed[RX_BUFFER_SIZE];
-                memcpy(fixed, data.data(), data.size());
-                memcpy(fixed + 2, _hubTransportId, 16);
+                if (!_wrapBuffer) return;
+                memcpy(_wrapBuffer, data.data(), data.size());
+                memcpy(_wrapBuffer + 2, _hubTransportId, 16);
 
                 Serial.printf("[TCP] TX %d bytes (H2 transport_id fixed) to %s:%d\n",
                               (int)data.size(), _host.c_str(), _port);
-                sendFrame(fixed, data.size());
+                sendFrame(_wrapBuffer, data.size());
                 InterfaceImpl::handle_outgoing(data);
                 return;
             }
diff --git a/src/transport/TCPClientInterface.h b/src/transport/TCPClientInterface.h
index 7cef646..583c67f 100644
--- a/src/transport/TCPClientInterface.h
+++ b/src/transport/TCPClientInterface.h
@@ -36,7 +36,8 @@ private:
     unsigned long _lastAttempt = 0;
     unsigned long _lastRxTime = 0;
     uint8_t* _rxBuffer = nullptr;
-    uint8_t* _txBuffer = nullptr;  // PSRAM-allocated send frame buffer
+    uint8_t* _txBuffer = nullptr;    // PSRAM-allocated send frame buffer
+    uint8_t* _wrapBuffer = nullptr;  // PSRAM-allocated packet rewrite buffer
     static constexpr size_t RX_BUFFER_SIZE = 2048;
     static constexpr size_t TX_BUFFER_SIZE = RX_BUFFER_SIZE * 2 + 2;
 
diff --git a/src/transport/WiFiInterface.cpp b/src/transport/WiFiInterface.cpp
index 1e11f29..0c44779 100644
--- a/src/transport/WiFiInterface.cpp
+++ b/src/transport/WiFiInterface.cpp
@@ -77,6 +77,11 @@ void WiFiInterface::stopFull() {
 void WiFiInterface::acceptClients() {
     WiFiClient newClient = _server.available();
     if (newClient) {
+        if ((int)_clients.size() >= MAX_AP_CLIENTS) {
+            newClient.stop();
+            Serial.printf("[WIFI] Client rejected (max %d reached)\n", MAX_AP_CLIENTS);
+            return;
+        }
         _clients.push_back(newClient);
         Serial.printf("[WIFI] Client connected (%d total)\n", (int)_clients.size());
     }
diff --git a/src/transport/WiFiInterface.h b/src/transport/WiFiInterface.h
index 22b1248..1f9d016 100644
--- a/src/transport/WiFiInterface.h
+++ b/src/transport/WiFiInterface.h
@@ -66,6 +66,8 @@ private:
     std::vector<WiFiClient> _clients;
     uint8_t _rxBuffer[600];
 
+    static constexpr int MAX_AP_CLIENTS = 4;
+
     static constexpr uint8_t FRAME_START = 0x7E;
     static constexpr uint8_t FRAME_ESC   = 0x7D;
     static constexpr uint8_t FRAME_XOR   = 0x20;