exposedev · mpociot · Jul 9, 2026
diff --git a/app/Connections/ConnectionManager.php b/app/Connections/ConnectionManager.php
@@ -17,6 +17,12 @@ class ConnectionManager implements ConnectionManagerContract
     /** @var array */
     protected $connections = [];
 
+    /** @var array */
+    protected $connectionsByClientId = [];
+
+    /** @var array */
+    protected $httpConnectionsBySubdomainAndServerHost = [];
+
     /** @var array */
     protected $httpConnections = [];
 
@@ -82,6 +88,8 @@ public function storeConnection(string $host, ?string $subdomain, ?string $serve
         );
 
         $this->connections[] = $storedConnection;
+        $this->connectionsByClientId[$clientId] = $storedConnection;
+        $this->httpConnectionsBySubdomainAndServerHost[$this->httpConnectionKey($storedConnection->subdomain, $storedConnection->serverHost)][$clientId] = $storedConnection;
 
         $this->statisticsCollector->siteShared($this->getAuthTokenFromConnection($connection));
 
@@ -116,6 +124,7 @@ public function storeTcpConnection(int $port, ConnectionInterface $connection):
         );
 
         $this->connections[] = $storedConnection;
+        $this->connectionsByClientId[$clientId] = $storedConnection;
 
         $this->statisticsCollector->portShared($this->getAuthTokenFromConnection($connection));
 
@@ -158,44 +167,59 @@ public function getHttpConnectionForRequestId(string $requestId): ?HttpConnectio
         return $this->httpConnections[$requestId] ?? null;
     }
 
+    public function removeHttpConnection(string $requestId): void
+    {
+        unset($this->httpConnections[$requestId]);
+    }
+
     public function removeControlConnection($connection)
     {
         if (isset($connection->request_id)) {
-            if (isset($this->httpConnections[$connection->request_id])) {
-                unset($this->httpConnections[$connection->request_id]);
-            }
+            $this->removeHttpConnection($connection->request_id);
         }
 
         if (isset($connection->client_id)) {
             $clientId = $connection->client_id;
 
-            $controlConnection = collect($this->connections)->first(function ($connection) use ($clientId) {
-                return $connection->client_id == $clientId;
-            });
+            $controlConnection = $this->connectionsByClientId[$clientId] ?? null;
 
             if ($controlConnection instanceof TcpControlConnection) {
                 $controlConnection->stop();
-                $controlConnection = null;
             }
 
-            $this->connections = collect($this->connections)->reject(function ($connection) use ($clientId) {
-                return $connection->client_id == $clientId;
-            })->toArray();
+            if ($controlConnection instanceof ControlConnection && ! $controlConnection instanceof TcpControlConnection) {
+                $httpConnectionKey = $this->httpConnectionKey($controlConnection->subdomain, $controlConnection->serverHost);
+
+                unset($this->httpConnectionsBySubdomainAndServerHost[$httpConnectionKey][$clientId]);
+
+                if (empty($this->httpConnectionsBySubdomainAndServerHost[$httpConnectionKey])) {
+                    unset($this->httpConnectionsBySubdomainAndServerHost[$httpConnectionKey]);
+                }
+            }
+
+            unset($this->connectionsByClientId[$clientId]);
+
+            foreach ($this->connections as $index => $storedConnection) {
+                if ($storedConnection->client_id == $clientId) {
+                    unset($this->connections[$index]);
+                    break;
+                }
+            }
+
+            $this->connections = array_values($this->connections);
         }
     }
 
     public function findControlConnectionForSubdomainAndServerHost($subdomain, $serverHost): ?ControlConnection
     {
-        return collect($this->connections)->last(function ($connection) use ($subdomain, $serverHost) {
-            return $connection->subdomain == $subdomain && $connection->serverHost === $serverHost;
-        });
+        $connections = $this->httpConnectionsBySubdomainAndServerHost[$this->httpConnectionKey($subdomain, $serverHost)] ?? [];
+
+        return empty($connections) ? null : end($connections);
     }
 
     public function findControlConnectionForClientId(string $clientId): ?ControlConnection
     {
-        return collect($this->connections)->last(function ($connection) use ($clientId) {
-            return $connection->client_id == $clientId;
-        });
+        return $this->connectionsByClientId[$clientId] ?? null;
     }
 
     public function findControlConnectionsForIp(string $ip): array
@@ -222,6 +246,11 @@ protected function getAuthTokenFromConnection(ConnectionInterface $connection):
         return QueryParameters::create($connection->httpRequest)->get('authToken');
     }
 
+    protected function httpConnectionKey($subdomain, $serverHost): string
+    {
+        return $serverHost."\0".$subdomain;
+    }
+
     public function getConnectionsForAuthToken(string $authToken): array
     {
         return collect($this->connections)

diff --git a/app/Contracts/ConnectionManager.php b/app/Contracts/ConnectionManager.php
@@ -18,6 +18,8 @@ public function storeHttpConnection(ConnectionInterface $httpConnection, $reques
 
     public function getHttpConnectionForRequestId(string $requestId): ?HttpConnection;
 
+    public function removeHttpConnection(string $requestId): void;
+
     public function removeControlConnection($connection);
 
     public function findControlConnectionForSubdomainAndServerHost($subdomain, $serverHost): ?ControlConnection;

diff --git a/app/Http/Controllers/ControlMessageController.php b/app/Http/Controllers/ControlMessageController.php
@@ -57,7 +57,9 @@ public function onClose(ConnectionInterface $connection)
     {
         if (isset($connection->request_id)) {
             $httpConnection = $this->connectionManager->getHttpConnectionForRequestId($connection->request_id);
-            $httpConnection->close();
+            if ($httpConnection !== null) {
+                $httpConnection->close();
+            }
         }
 
         $this->connectionManager->removeControlConnection($connection);
@@ -92,7 +94,9 @@ protected function sendResponseToHttpConnection(string $requestId, $response)
     {
         $httpConnection = $this->connectionManager->getHttpConnectionForRequestId($requestId);
 
-        $httpConnection->send($response);
+        if ($httpConnection !== null) {
+            $httpConnection->send($response);
+        }
     }
 
     protected function authenticate(ConnectionInterface $connection, $data)

diff --git a/app/Http/Controllers/TunnelMessageController.php b/app/Http/Controllers/TunnelMessageController.php
@@ -107,6 +107,9 @@ protected function sendRequestToClient(Request $request, ControlConnection $cont
         $requestId = $request->header('X-Expose-Request-ID');
 
         $httpConnection = $this->connectionManager->storeHttpConnection($httpConnection, $requestId);
+        $httpConnection->getConnection()->on('close', function () use ($requestId) {
+            $this->connectionManager->removeHttpConnection($requestId);
+        });
 
         transform($this->passRequestThroughModifiers($request, $httpConnection), function (Request $request) use ($httpConnection, $controlConnection, $requestId) {
             $controlConnection->once('proxy_ready_'.$requestId, function (ConnectionInterface $proxy) use ($httpConnection, $request) {

diff --git a/tests/Feature/Server/ConnectionManagerTest.php b/tests/Feature/Server/ConnectionManagerTest.php
@@ -9,6 +9,7 @@
 use Expose\Server\Contracts\SubdomainGenerator;
 use Expose\Server\Contracts\UserRepository;
 use Mockery;
+use Ratchet\ConnectionInterface;
 use React\EventLoop\LoopInterface;
 use Tests\Feature\TestCase;
 
@@ -96,4 +97,66 @@ public function it_still_applies_connection_length_limits_to_users_without_custo
 
         $timerCallback();
     }
+
+    /** @test */
+    public function it_uses_indexes_for_client_and_subdomain_lookups()
+    {
+        $statisticsCollector = Mockery::mock(StatisticsCollector::class);
+        $statisticsCollector->shouldReceive('siteShared')->twice();
+
+        $logger = Mockery::mock(LoggerRepository::class);
+        $logger->shouldReceive('logSubdomain')->twice();
+
+        $manager = new ConnectionManager(
+            Mockery::mock(SubdomainGenerator::class),
+            $statisticsCollector,
+            $logger,
+            Mockery::mock(LoopInterface::class)
+        );
+
+        $firstConnection = $this->mockSocketConnection();
+        $secondConnection = $this->mockSocketConnection();
+
+        $first = $manager->storeConnection('127.0.0.1:8085', 'shared', 'localhost', $firstConnection);
+        $second = $manager->storeConnection('127.0.0.1:8086', 'shared', 'localhost', $secondConnection);
+
+        $this->assertSame($second, $manager->findControlConnectionForSubdomainAndServerHost('shared', 'localhost'));
+        $this->assertSame($first, $manager->findControlConnectionForClientId($first->client_id));
+
+        $manager->removeControlConnection($secondConnection);
+
+        $this->assertSame($first, $manager->findControlConnectionForSubdomainAndServerHost('shared', 'localhost'));
+        $this->assertNull($manager->findControlConnectionForClientId($second->client_id));
+    }
+
+    /** @test */
+    public function it_can_remove_http_connections_by_request_id()
+    {
+        $manager = new ConnectionManager(
+            Mockery::mock(SubdomainGenerator::class),
+            Mockery::mock(StatisticsCollector::class),
+            Mockery::mock(LoggerRepository::class),
+            Mockery::mock(LoopInterface::class)
+        );
+
+        $connection = $this->mockSocketConnection();
+
+        $manager->storeHttpConnection($connection, 'request-one');
+        $this->assertNotNull($manager->getHttpConnectionForRequestId('request-one'));
+
+        $manager->removeHttpConnection('request-one');
+
+        $this->assertNull($manager->getHttpConnectionForRequestId('request-one'));
+    }
+
+    protected function mockSocketConnection(): ConnectionInterface
+    {
+        $connection = Mockery::mock(ConnectionInterface::class);
+        $connection->httpRequest = new \GuzzleHttp\Psr7\Request('GET', '/expose/control?authToken=test-token&version=test-version');
+        $connection->remoteAddress = '127.0.0.1';
+        $connection->shouldReceive('send')->byDefault();
+        $connection->shouldReceive('close')->byDefault();
+
+        return $connection;
+    }
 }