diff --git a/src/cloud-sql-instance.ts b/src/cloud-sql-instance.ts index 8d9f0e91..00db0774 100644 --- a/src/cloud-sql-instance.ts +++ b/src/cloud-sql-instance.ts @@ -461,7 +461,7 @@ export class CloudSQLInstance { // Add the socket to the list this.sockets.add(socket); // When the socket is closed, remove it. - socket.once('closed', () => { + socket.once('close', () => { this.sockets.delete(socket); }); } diff --git a/src/connector.ts b/src/connector.ts index a665fad5..c998ac06 100644 --- a/src/connector.ts +++ b/src/connector.ts @@ -337,12 +337,22 @@ export class Connector { console.error(err); }); + server.once('close', () => { + this.localProxies.delete(server); + }); + // When a connection is established, pipe data from the // local proxy server to the secure TCP Socket and vice-versa. server.on('connection', c => { const s = stream(); this.sockets.add(s); this.sockets.add(c); + s.once('close', () => { + this.sockets.delete(s); + }); + c.once('close', () => { + this.sockets.delete(c); + }); c.pipe(s); s.pipe(c); }); diff --git a/test/cloud-sql-instance.ts b/test/cloud-sql-instance.ts index 9b4c2471..f412634e 100644 --- a/test/cloud-sql-instance.ts +++ b/test/cloud-sql-instance.ts @@ -628,4 +628,79 @@ t.test('CloudSQLInstance', async t => { }))(); } ); + + t.test('addSocket manages and cleans up sockets on close event', async t => { + const {EventEmitter} = await import('node:events'); + interface InstanceWithInternals { + sockets: Set; + addSocket(socket: unknown): void; + close(): void; + } + const instance = new CloudSQLInstance({ + options: { + ipType: IpAddressTypes.PUBLIC, + authType: AuthTypes.PASSWORD, + domainName: 'db.example.com', + sqlAdminFetcher: fetcher, + }, + }) as unknown as InstanceWithInternals; + t.after(() => instance.close()); + + class MockSocket extends EventEmitter { + destroyed = false; + destroy() { + this.destroyed = true; + this.emit('close'); + } + } + + const socket1 = new MockSocket(); + const socket2 = new MockSocket(); + + instance.addSocket(socket1); + instance.addSocket(socket2); + + t.equal(instance.sockets.size, 2, 'both sockets added'); + + // Close socket1 + socket1.emit('close'); + t.equal(instance.sockets.size, 1, 'socket1 removed on close'); + t.ok(instance.sockets.has(socket2), 'socket2 still tracked'); + + // close instance destroys remaining sockets + instance.close(); + t.equal( + socket2.destroyed, + true, + 'remaining socket destroyed on instance close' + ); + }); + + t.test('addSocket ignores sockets when domainName is not set', async t => { + const {EventEmitter} = await import('node:events'); + interface InstanceWithInternals { + sockets: Set; + addSocket(socket: unknown): void; + close(): void; + } + const instance = new CloudSQLInstance({ + options: { + ipType: IpAddressTypes.PUBLIC, + authType: AuthTypes.PASSWORD, + instanceConnectionName: 'my-project:us-east1:my-instance', + sqlAdminFetcher: fetcher, + }, + }) as unknown as InstanceWithInternals; + t.after(() => instance.close()); + + const socket = new EventEmitter(); + Object.assign(socket, {destroy: () => {}}); + + instance.addSocket(socket); + t.equal( + instance.sockets.size, + 0, + 'socket not added when domainName not set' + ); + }); }); diff --git a/test/connector.ts b/test/connector.ts index 6210a603..dda4eb95 100644 --- a/test/connector.ts +++ b/test/connector.ts @@ -694,3 +694,160 @@ t.test( t.same(mockSocket.destroyed, true, 'old instance closed its sockets'); } ); + +t.test('Connector startLocalProxy manages and cleans up sockets', async t => { + setupCredentials(t); + + class MockSocket extends EventEmitter { + destroyed = false; + pipedTo: unknown = null; + pipe(dest: unknown) { + this.pipedTo = dest; + return dest; + } + destroy() { + this.destroyed = true; + this.emit('close'); + return this; + } + } + + let serverListenOptions: unknown = null; + let serverClosed = false; + let mockServer: EventEmitter & {listen: Function; close: Function}; + + const {Connector} = t.mockRequire('../src/connector', { + 'node:net': { + createServer() { + mockServer = Object.assign(new EventEmitter(), { + listen(opts: unknown, cb: Function) { + serverListenOptions = opts; + if (cb) cb(); + }, + close(cb?: Function) { + serverClosed = true; + mockServer.emit('close'); + if (cb) cb(); + }, + }); + return mockServer; + }, + }, + '../src/sqladmin-fetcher': { + SQLAdminFetcher: class { + getInstanceMetadata() { + return Promise.resolve({ + ipAddresses: { + public: '127.0.0.1', + }, + serverCaCert: { + cert: CA_CERT, + expirationTime: '2033-01-06T10:00:00.232Z', + }, + }); + } + getEphemeralCertificate() { + return Promise.resolve({ + cert: CLIENT_CERT, + expirationTime: '2033-01-06T10:00:00.232Z', + }); + } + }, + }, + '../src/cloud-sql-instance': t.mockRequire('../src/cloud-sql-instance', { + '../src/crypto': { + generateKeys: async () => ({ + publicKey: '-----BEGIN PUBLIC KEY-----', + privateKey: CLIENT_KEY, + }), + }, + }), + }); + + interface ConnectorWithInternals { + localProxies: Set; + sockets: Set; + startLocalProxy(opts: unknown): Promise; + close(): void; + getOptions: (opts: unknown) => Promise<{stream: () => unknown}>; + } + + const connector = new Connector() as unknown as ConnectorWithInternals; + const mockStreamSockets: MockSocket[] = []; + + // Mock getOptions to return mock stream sockets + connector.getOptions = async () => { + return { + stream() { + const s = new MockSocket(); + mockStreamSockets.push(s); + return s; + }, + }; + }; + + await connector.startLocalProxy({ + ipType: 'PUBLIC', + instanceConnectionName: 'my-project:us-east1:my-instance', + listenOptions: {path: '/tmp/test.sock'}, + }); + + t.same( + serverListenOptions, + {path: '/tmp/test.sock', readableAll: undefined, writableAll: undefined}, + 'server should listen with options' + ); + t.equal( + connector.localProxies.size, + 1, + 'server should be tracked in localProxies' + ); + + // Simulate a client connecting + const clientSocket1 = new MockSocket(); + mockServer.emit('connection', clientSocket1); + + t.equal(mockStreamSockets.length, 1, 'stream() should have been called'); + const streamSocket1 = mockStreamSockets[0]; + + t.equal( + clientSocket1.pipedTo, + streamSocket1, + 'client socket piped to stream socket' + ); + t.equal( + streamSocket1.pipedTo, + clientSocket1, + 'stream socket piped to client socket' + ); + t.equal(connector.sockets.size, 2, 'both client and stream sockets tracked'); + + // Emit close on stream socket + streamSocket1.emit('close'); + t.equal(connector.sockets.size, 1, 'stream socket removed on close'); + t.ok(connector.sockets.has(clientSocket1), 'client socket still tracked'); + + // Emit close on client socket + clientSocket1.emit('close'); + t.equal(connector.sockets.size, 0, 'client socket removed on close'); + + // Simulate another connection + const clientSocket2 = new MockSocket(); + mockServer.emit('connection', clientSocket2); + t.equal(connector.sockets.size, 2, 'second connection pair tracked'); + + // Close connector + connector.close(); + t.equal(serverClosed, true, 'server should be closed'); + t.equal( + connector.localProxies.size, + 0, + 'server removed from localProxies on close' + ); + t.equal(clientSocket2.destroyed, true, 'remaining client socket destroyed'); + t.equal( + mockStreamSockets[1].destroyed, + true, + 'remaining stream socket destroyed' + ); +});