fix: bound per-connection port-forward work (#298)

src/terminal/forwarding/ForwardDestinationHandler.cpp

@@ -44,9 +44,10 @@ void ForwardDestinationHandler::update(vector<PortForwardData>* retval,
     return;
   }
 
-  while (socketHandler->hasData(fd)) {
-    char buf[1024];
-    int bytesRead = socketHandler->read(fd, buf, 1024);
+  size_t bytesThisUpdate = 0;
+  while (bytesThisUpdate < MAX_BYTES_PER_UPDATE && socketHandler->hasData(fd)) {
+    char buf[16 * 1024];
+    int bytesRead = socketHandler->read(fd, buf, sizeof(buf));
     auto readErrno = GetErrno();
     if (bytesRead == -1 && (readErrno == EAGAIN || readErrno == EWOULDBLOCK)) {
       // Bail for now
@@ -65,6 +66,7 @@ void ForwardDestinationHandler::update(vector<PortForwardData>* retval,
     } else {
       VLOG(1) << "Reading " << bytesRead << " bytes from socket " << socketId;
       pwd.set_buffer(string(buf, bytesRead));
+      bytesThisUpdate += bytesRead;
     }
     retval->push_back(pwd);
     if (bytesRead < 1) {

src/terminal/forwarding/ForwardDestinationHandler.hpp

@@ -11,6 +11,7 @@ namespace et {
  */
 class ForwardDestinationHandler {
  public:
+  static constexpr size_t MAX_BYTES_PER_UPDATE = 64 * 1024;
   /** @brief Binds the handler to a destination fd so data can be sent
    * downstream. */
   ForwardDestinationHandler(shared_ptr<SocketHandler> _socketHandler, int _fd,

src/terminal/forwarding/ForwardSourceHandler.cpp

@@ -193,6 +193,7 @@ bool ForwardSourceHandler::update(vector<PortForwardData>* data,
   }
 
   vector<int> socketsToRemove;
+  size_t bytesThisUpdate = 0;
 
   for (auto& it : socketFdMap) {
     int socketId = it.first;
@@ -215,25 +216,25 @@ bool ForwardSourceHandler::update(vector<PortForwardData>* data,
       continue;
     }
 
-    while (true) {
+    while (bytesThisUpdate < MAX_BYTES_PER_UPDATE) {
       if (!stdioMode && !socketHandler->hasData(fd)) {
         break;
       }
 
-      char buf[1024];
+      char buf[16 * 1024];
       int bytesRead = -1;
       int readErrno = 0;
       if (stdioMode) {
 #ifndef WIN32
-        bytesRead = ::read(fd, buf, 1024);
+        bytesRead = ::read(fd, buf, sizeof(buf));
         readErrno = errno;
         SetErrno(readErrno);
 #else
-        bytesRead = socketHandler->read(fd, buf, 1024);
+        bytesRead = socketHandler->read(fd, buf, sizeof(buf));
         readErrno = GetErrno();
 #endif
       } else {
-        bytesRead = socketHandler->read(fd, buf, 1024);
+        bytesRead = socketHandler->read(fd, buf, sizeof(buf));
         readErrno = GetErrno();
       }
       if (bytesRead == -1 &&
@@ -260,6 +261,7 @@ bool ForwardSourceHandler::update(vector<PortForwardData>* data,
       } else {
         VLOG(1) << "Reading " << bytesRead << " bytes from socket " << socketId;
         pwd.set_buffer(string(buf, bytesRead));
+        bytesThisUpdate += bytesRead;
       }
       data->push_back(pwd);
       if (bytesRead < 1) {

src/terminal/forwarding/ForwardSourceHandler.hpp

@@ -12,6 +12,7 @@ namespace et {
  */
 class ForwardSourceHandler {
  public:
+  static constexpr size_t MAX_BYTES_PER_UPDATE = 64 * 1024;
   /**
    * @brief Creates source/destination handlers used for local port forwarding.
    * @param alreadyListening If true, skip listen(); caller already registered

test/unit_tests/ForwardDestinationHandlerTest.cpp

@@ -187,3 +187,21 @@ TEST_CASE("ForwardDestinationHandler ignores transient EAGAIN reads",
   REQUIRE(handler.getFd() == 70);
   REQUIRE(socketHandler->closedFds.empty());
 }
+
+TEST_CASE("ForwardDestinationHandler yields after a bounded data batch",
+          "[ForwardDestinationHandler][issue298]") {
+  auto socketHandler = std::make_shared<MockSocketHandler>();
+  ForwardDestinationHandler handler(socketHandler, /*fd=*/71, /*socketId=*/6);
+
+  const string chunk(16 * 1024, 'x');
+  for (int i = 0; i < 5; ++i) {
+    socketHandler->enqueueHasData(true);
+    socketHandler->enqueueRead(chunk.size(), chunk);
+  }
+
+  std::vector<PortForwardData> responses;
+  handler.update(&responses);
+
+  REQUIRE(responses.size() == 4);
+  REQUIRE(socketHandler->performedReads.size() == 4);
+}

test/unit_tests/ForwardSourceHandlerTest.cpp

@@ -386,6 +386,36 @@ TEST_CASE("ForwardSourceHandler update ignores transient EAGAIN reads",
   REQUIRE(socketHandler->closedFds.empty());
 }
 
+TEST_CASE("ForwardSourceHandler update yields after a bounded data batch",
+          "[ForwardSourceHandler][issue298]") {
+  auto socketHandler = std::make_shared<MockSocketHandler>();
+  socketHandler->setEndpointFds({100});
+  socketHandler->enqueueAccept(42);
+
+  SocketEndpoint source;
+  source.set_name("localhost");
+  source.set_port(8080);
+  SocketEndpoint destination;
+  destination.set_name("remote");
+  destination.set_port(9090);
+  ForwardSourceHandler handler(socketHandler, source, destination);
+
+  int fd = handler.listen();
+  handler.addSocket(123, fd);
+
+  const string chunk(16 * 1024, 'x');
+  for (int i = 0; i < 5; ++i) {
+    socketHandler->enqueueHasData(true);
+    socketHandler->enqueueRead(chunk.size(), chunk);
+  }
+
+  std::vector<PortForwardData> data;
+  handler.update(&data);
+
+  REQUIRE(data.size() == 4);
+  REQUIRE(socketHandler->performedReads.size() == 4);
+}
+
 TEST_CASE("ForwardSourceHandler update reads only ready fds",
           "[ForwardSourceHandler]") {
   auto socketHandler = std::make_shared<MockSocketHandler>();