fix: allow tcp initial packet rewrite

Author: abnegate

src/Adapter/TCP.php

@@ -4,6 +4,7 @@
 
 use Swoole\Coroutine\Client;
 use Utopia\Proxy\Adapter;
+use Utopia\Proxy\ConnectionResult;
 use Utopia\Proxy\Protocol;
 use Utopia\Proxy\Resolver;
 use Utopia\Proxy\Sockmap\Loader as Sockmap;
@@ -40,9 +41,14 @@ class TCP extends Adapter
     /** @var array<int, Client> */
     protected array $connections = [];
 
+    /** @var array<int, string> */
+    protected array $initialData = [];
+
     /** @var array<int, int> client fd => backend fd, for sockmap pairs handed to the kernel */
     protected array $sockmapPairs = [];
 
+    protected ?\Closure $initialDataTransformer = null;
+
     protected float $timeout = 30.0;
 
     protected float $connectTimeout = 5.0;
@@ -66,6 +72,25 @@ public function __construct(
         parent::__construct($resolver);
     }
 
+    /**
+     * Transform the first client packet before it is forwarded to the backend.
+     *
+     * The callback receives the raw initial packet and route result. This lets
+     * protocol-aware callers route by one identifier while forwarding a backend
+     * specific startup packet.
+     */
+    public function onInitialData(callable $callback): static
+    {
+        $this->initialDataTransformer = $callback(...);
+
+        return $this;
+    }
+
+    public function getInitialData(int $fd, string $default): string
+    {
+        return $this->initialData[$fd] ?? $default;
+    }
+
     public function setTimeout(float $timeout): static
     {
         $this->timeout = $timeout;
@@ -181,6 +206,7 @@ public function getConnection(string $data, int $fd): Client
 
         $result = $this->route($data);
         [$host, $port] = self::parseEndpoint($result->endpoint, $this->port);
+        $initialData = $this->transformInitialData($data, $result);
 
         $client = new Client(SWOOLE_SOCK_TCP);
         $client->set([
@@ -196,10 +222,25 @@ public function getConnection(string $data, int $fd): Client
 
         $this->applySocketOptions($client);
         $this->connections[$fd] = $client;
+        $this->initialData[$fd] = $initialData;
 
         return $client;
     }
 
+    protected function transformInitialData(string $data, ConnectionResult $result): string
+    {
+        if ($this->initialDataTransformer === null) {
+            return $data;
+        }
+
+        $transformed = ($this->initialDataTransformer)($data, $result);
+        if (!\is_string($transformed)) {
+            throw new \Exception('Initial data transformer must return string');
+        }
+
+        return $transformed;
+    }
+
     /**
      * Hand the (client fd, backend fd) pair to the kernel via sockmap.
      *
@@ -270,6 +311,7 @@ public function closeConnection(int $fd): void
         }
 
         unset($this->connections[$fd]);
+        unset($this->initialData[$fd]);
         $client->close();
     }
 

src/Server/TCP/Swoole.php

@@ -344,7 +344,7 @@ public function onReceive(Server $server, int $fd, string $data, int $port): voi
             // sockmap — once the pair is in the map, any send() on the
             // backend fd would be redirected back to the client by the
             // sk_msg program.
-            $backend->send($data);
+            $backend->send($adapter->getInitialData($fd, $data));
 
             // Activate kernel zero-copy relay. If successful the kernel
             // handles all subsequent bytes and we skip the userspace

tests/Integration/TCPServerTest.php

@@ -59,18 +59,23 @@ public function testTCPAdapterGetConnectionAndClose(): void
                 $adapter->setSkipValidation(true);
                 $adapter->setTimeout(5.0);
                 $adapter->setConnectTimeout(2.0);
+                $adapter->onInitialData(function (string $data): string {
+                    return str_replace('routing', 'backend', $data);
+                });
 
                 // getConnection dials the backend and caches by fd
                 $client = $adapter->getConnection('routing-data', 1);
                 $this->assertInstanceOf(Client::class, $client);
                 $this->assertTrue($client->isConnected());
+                $this->assertSame('backend-data', $adapter->getInitialData(1, 'routing-data'));
 
                 // Same fd returns cached connection
                 $cached = $adapter->getConnection('ignored', 1);
                 $this->assertSame($client, $cached);
 
                 // closeConnection cleans up
                 $adapter->closeConnection(1);
+                $this->assertSame('fallback', $adapter->getInitialData(1, 'fallback'));
 
                 // Closing again is a no-op
                 $adapter->closeConnection(1);