diff --git a/app/Connections/ConnectionManager.php b/app/Connections/ConnectionManager.php index 94dafc3..25a23e1 100644 --- 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 index 77fda40..15e5ce8 100644 --- 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 index d748e8a..02feea3 100644 --- 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 index 34d597e..f359ce3 100644 --- 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 index d727be1..fd1ac67 100644 --- 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; + } }