fix: Clean up sockets when closed in startLocalProxy and CloudSQLInstance

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);
     });
   }

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);
     });

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<unknown>;
+      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<unknown>;
+      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'
+    );
+  });
 });

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<unknown>;
+    sockets: Set<unknown>;
+    startLocalProxy(opts: unknown): Promise<void>;
+    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'
+  );
+});