test(ssh): improve SSH proxy test coverage

test/ssh/AgentForwarding.test.ts

@@ -14,10 +14,15 @@
  * limitations under the License.
  */
 
-import { describe, it, expect, vi, beforeEach, type Mock } from 'vitest';
-import { LazySSHAgent, createLazyAgent } from '../../src/proxy/ssh/AgentForwarding';
+import { describe, it, expect, vi, beforeEach, afterEach, type Mock } from 'vitest';
+import {
+  LazySSHAgent,
+  createLazyAgent,
+  openTemporaryAgentChannel,
+} from '../../src/proxy/ssh/AgentForwarding';
 import { SSHAgentProxy } from '../../src/proxy/ssh/AgentProxy';
 import { ClientWithUser } from '../../src/proxy/ssh/types';
+import * as sshInternals from '../../src/proxy/ssh/sshInternals';
 
 describe('AgentForwarding', () => {
   let mockClient: Partial<ClientWithUser>;
@@ -322,8 +327,6 @@ describe('AgentForwarding', () => {
 
   describe('openTemporaryAgentChannel', () => {
     it('should return null when client has no protocol', async () => {
-      const { openTemporaryAgentChannel } = await import('../../src/proxy/ssh/AgentForwarding');
-
       const clientWithoutProtocol: any = {
         agentForwardingEnabled: true,
       };
@@ -334,105 +337,267 @@ describe('AgentForwarding', () => {
     });
 
     it('should handle timeout when channel confirmation not received', async () => {
-      const { openTemporaryAgentChannel } = await import('../../src/proxy/ssh/AgentForwarding');
+      vi.useFakeTimers();
+      try {
+        const mockClient: any = {
+          agentForwardingEnabled: true,
+          _protocol: {
+            _handlers: {},
+            openssh_authAgent: vi.fn(),
+            channelSuccess: vi.fn(),
+          },
+          _chanMgr: {
+            _channels: {},
+            _count: 0,
+          },
+        };
+
+        const promise = openTemporaryAgentChannel(mockClient);
+
+        // Advance past the 5s confirmation window instead of waiting in real time
+        vi.advanceTimersByTime(5001);
+        const result = await promise;
+
+        expect(result).toBeNull();
+      } finally {
+        vi.useRealTimers();
+      }
+    });
+
+    it('should find next available channel ID when channels exist', async () => {
+      vi.useFakeTimers();
+      try {
+        const mockClient: any = {
+          agentForwardingEnabled: true,
+          _protocol: {
+            _handlers: {},
+            openssh_authAgent: vi.fn(),
+            channelSuccess: vi.fn(),
+          },
+          _chanMgr: {
+            _channels: {
+              1: 'occupied',
+              2: 'occupied',
+              // Channel 3 should be used
+            },
+            _count: 2,
+          },
+        };
+
+        const promise = openTemporaryAgentChannel(mockClient);
+
+        // openssh_authAgent is called synchronously, before the promise settles
+        expect(mockClient._protocol.openssh_authAgent).toHaveBeenCalledWith(
+          3,
+          expect.any(Number),
+          expect.any(Number),
+        );
+
+        // Advance past the confirmation timeout so the promise settles cleanly
+        vi.advanceTimersByTime(5001);
+        await promise;
+      } finally {
+        vi.useRealTimers();
+      }
+    });
+
+    it('should use channel ID 1 when no channels exist', async () => {
+      vi.useFakeTimers();
+      try {
+        const mockClient: any = {
+          agentForwardingEnabled: true,
+          _protocol: {
+            _handlers: {},
+            openssh_authAgent: vi.fn(),
+            channelSuccess: vi.fn(),
+          },
+          _chanMgr: {
+            _channels: {},
+            _count: 0,
+          },
+        };
+
+        const promise = openTemporaryAgentChannel(mockClient);
 
+        expect(mockClient._protocol.openssh_authAgent).toHaveBeenCalledWith(
+          1,
+          expect.any(Number),
+          expect.any(Number),
+        );
+
+        vi.advanceTimersByTime(5001);
+        await promise;
+      } finally {
+        vi.useRealTimers();
+      }
+    });
+
+    it('should return null when client has no chanMgr', async () => {
       const mockClient: any = {
         agentForwardingEnabled: true,
         _protocol: {
           _handlers: {},
           openssh_authAgent: vi.fn(),
           channelSuccess: vi.fn(),
         },
-        _chanMgr: {
-          _channels: {},
-          _count: 0,
-        },
       };
 
       const result = await openTemporaryAgentChannel(mockClient);
 
-      // Should timeout and return null after 5 seconds
       expect(result).toBeNull();
-    }, 6000);
+      expect(mockClient._protocol.openssh_authAgent).not.toHaveBeenCalled();
+    });
 
-    it('should find next available channel ID when channels exist', async () => {
-      const { openTemporaryAgentChannel } = await import('../../src/proxy/ssh/AgentForwarding');
+    describe('CHANNEL_OPEN_CONFIRMATION handler', () => {
+      let getChannelModuleSpy: ReturnType<typeof vi.spyOn>;
 
-      const mockClient: any = {
-        agentForwardingEnabled: true,
-        _protocol: {
-          _handlers: {},
-          openssh_authAgent: vi.fn(),
-          channelSuccess: vi.fn(),
-        },
-        _chanMgr: {
-          _channels: {
-            1: 'occupied',
-            2: 'occupied',
-            // Channel 3 should be used
+      afterEach(() => {
+        getChannelModuleSpy?.mockRestore();
+      });
+
+      function makeMockClient(existingHandler?: (...args: unknown[]) => void) {
+        const handlers: Record<string, (...args: unknown[]) => void> = {};
+        if (existingHandler) {
+          handlers.CHANNEL_OPEN_CONFIRMATION = existingHandler;
+        }
+        return {
+          _protocol: {
+            _handlers: handlers,
+            openssh_authAgent: vi.fn(),
+            channelSuccess: vi.fn(),
           },
-          _count: 2,
-        },
-      };
+          _chanMgr: { _channels: {} as Record<number, any>, _count: 0 },
+        } as any;
+      }
+
+      function mockChannelModule(channelImpl?: (...args: unknown[]) => unknown) {
+        const mockChannel = { on: vi.fn(), write: vi.fn(), end: vi.fn() };
+        getChannelModuleSpy = vi.spyOn(sshInternals, 'getChannelModule').mockReturnValue({
+          Channel: channelImpl ?? vi.fn().mockReturnValue(mockChannel),
+          MAX_WINDOW: 2 * 1024 * 1024,
+          PACKET_SIZE: 32 * 1024,
+        });
+        return mockChannel;
+      }
 
-      // Start the operation but don't wait for completion (will timeout)
-      const promise = openTemporaryAgentChannel(mockClient);
+      it('should create AgentProxy on successful confirmation', async () => {
+        mockChannelModule();
+        const client = makeMockClient();
 
-      // Verify openssh_authAgent was called with the next available channel (3)
-      expect(mockClient._protocol.openssh_authAgent).toHaveBeenCalledWith(
-        3,
-        expect.any(Number),
-        expect.any(Number),
-      );
+        const promise = openTemporaryAgentChannel(client);
 
-      // Clean up - wait for timeout
-      await promise;
-    }, 6000);
+        const handler = client._protocol._handlers.CHANNEL_OPEN_CONFIRMATION;
+        handler(null, { recipient: 1, sender: 42, window: 65536, packetSize: 32768 });
 
-    it('should use channel ID 1 when no channels exist', async () => {
-      const { openTemporaryAgentChannel } = await import('../../src/proxy/ssh/AgentForwarding');
+        const result = await promise;
 
-      const mockClient: any = {
-        agentForwardingEnabled: true,
-        _protocol: {
-          _handlers: {},
-          openssh_authAgent: vi.fn(),
-          channelSuccess: vi.fn(),
-        },
-        _chanMgr: {
-          _channels: {},
-          _count: 0,
-        },
-      };
+        expect(result).toBeInstanceOf(SSHAgentProxy);
+        expect(client._chanMgr._channels[1]).toBeDefined();
+        expect(client._chanMgr._count).toBe(1);
+      });
 
-      const promise = openTemporaryAgentChannel(mockClient);
+      it('should call and restore original handler on confirmation', async () => {
+        mockChannelModule();
+        const originalHandler = vi.fn();
+        const client = makeMockClient(originalHandler);
 
-      expect(mockClient._protocol.openssh_authAgent).toHaveBeenCalledWith(
-        1,
-        expect.any(Number),
-        expect.any(Number),
-      );
+        const promise = openTemporaryAgentChannel(client);
 
-      await promise;
-    }, 6000);
+        const handler = client._protocol._handlers.CHANNEL_OPEN_CONFIRMATION;
+        const confirmInfo = { recipient: 1, sender: 42, window: 65536, packetSize: 32768 };
+        handler(null, confirmInfo);
 
-    it('should return null when client has no chanMgr', async () => {
-      const { openTemporaryAgentChannel } = await import('../../src/proxy/ssh/AgentForwarding');
+        await promise;
 
-      const mockClient: any = {
-        agentForwardingEnabled: true,
-        _protocol: {
-          _handlers: {},
-          openssh_authAgent: vi.fn(),
-          channelSuccess: vi.fn(),
-        },
-        // No _chanMgr — the internals guard should reject and return null
-      };
+        expect(originalHandler).toHaveBeenCalledWith(null, confirmInfo);
+        expect(client._protocol._handlers.CHANNEL_OPEN_CONFIRMATION).toBe(originalHandler);
+      });
 
-      const result = await openTemporaryAgentChannel(mockClient);
+      it('should delete handler when no original existed', async () => {
+        mockChannelModule();
+        const client = makeMockClient();
 
-      expect(result).toBeNull();
-      expect(mockClient._protocol.openssh_authAgent).not.toHaveBeenCalled();
+        const promise = openTemporaryAgentChannel(client);
+
+        const handler = client._protocol._handlers.CHANNEL_OPEN_CONFIRMATION;
+        handler(null, { recipient: 1, sender: 42, window: 65536, packetSize: 32768 });
+
+        await promise;
+
+        expect(client._protocol._handlers.CHANNEL_OPEN_CONFIRMATION).toBeUndefined();
+      });
+
+      it('should ignore confirmation for non-matching channel recipient', async () => {
+        vi.useFakeTimers();
+        try {
+          mockChannelModule();
+          const client = makeMockClient();
+
+          const promise = openTemporaryAgentChannel(client);
+
+          const handler = client._protocol._handlers.CHANNEL_OPEN_CONFIRMATION;
+          handler(null, { recipient: 999, sender: 42, window: 65536, packetSize: 32768 });
+
+          vi.advanceTimersByTime(5001);
+
+          const result = await promise;
+          expect(result).toBeNull();
+        } finally {
+          vi.useRealTimers();
+        }
+      });
+
+      it('should resolve null when Channel constructor throws', async () => {
+        mockChannelModule(function () {
+          throw new Error('Channel creation failed');
+        });
+        const client = makeMockClient();
+
+        const promise = openTemporaryAgentChannel(client);
+
+        const handler = client._protocol._handlers.CHANNEL_OPEN_CONFIRMATION;
+        handler(null, { recipient: 1, sender: 42, window: 65536, packetSize: 32768 });
+
+        const result = await promise;
+        expect(result).toBeNull();
+      });
+    });
+  });
+
+  describe('LazySSHAgent - lock recovery', () => {
+    it('should continue operating after a previous operation fails', () => {
+      return new Promise<void>((resolve) => {
+        const openChannelFn = vi.fn();
+        const client = {
+          agentForwardingEnabled: true,
+          clientIp: '127.0.0.1',
+          authenticatedUser: { username: 'testuser' },
+        } as unknown as ClientWithUser;
+
+        openChannelFn.mockRejectedValueOnce(new Error('Channel open failed'));
+
+        const identities = [
+          { publicKeyBlob: Buffer.from('key1'), comment: 'k', algorithm: 'ssh-ed25519' },
+        ];
+        const mockProxy = {
+          getIdentities: vi.fn().mockResolvedValue(identities),
+          sign: vi.fn(),
+          close: vi.fn(),
+        };
+        openChannelFn.mockResolvedValueOnce(mockProxy);
+
+        const agent = new LazySSHAgent(openChannelFn, client);
+
+        agent.getIdentities((err) => {
+          expect(err).toBeDefined();
+
+          agent.getIdentities((err2, keys) => {
+            expect(err2).toBeNull();
+            expect(keys).toHaveLength(1);
+            resolve();
+          });
+        });
+      });
     });
   });
 });