From e5955541bf722edb5e026b95b2e6df9b39e621ab Mon Sep 17 00:00:00 2001
From: Henry Barreto <henry.barreto@ossystems.com.br>
Date: Tue, 8 Oct 2024 16:38:22 -0300
Subject: [PATCH] chore: add tests for dynamic port forward

---
 src/preload/ssh/ssh.test.ts | 196 ++++++++++++++++++++++++++++++++++++
 1 file changed, 196 insertions(+)

diff --git a/src/preload/ssh/ssh.test.ts b/src/preload/ssh/ssh.test.ts
index 4394e38..3126153 100644
--- a/src/preload/ssh/ssh.test.ts
+++ b/src/preload/ssh/ssh.test.ts
@@ -2,11 +2,13 @@ import { describe, it, expect, beforeEach, vi, afterEach } from 'vitest'
 import { SSHConnectionLocalPortForward, SSHConnectionDynamicPortForward } from './ssh'
 import ssh2 from 'ssh2'
 import net from 'node:net'
+import socks from 'socksv5'
 import { Stream } from 'node:stream'
 import EventEmitter from 'node:events'
 
 vi.mock('ssh2')
 vi.mock('node:net')
+vi.mock('socksv5')
 
 describe('SSH Interface', () => {
   it('localPortForward should return a SSHConnection instance', () => {
@@ -187,3 +189,197 @@ describe('SSHConnectionLocalPortForward', () => {
     expect(emitSpy).toHaveBeenCalledWith('connect', expect.any(Number), expect.any(String))
   })
 })
+
+describe('SSHConnectionDynamicPortForward', () => {
+  let connection
+  let auth
+  let settings
+
+  let mockClient
+  let mockEvents
+
+  beforeEach(() => {
+    connection = new SSHConnectionDynamicPortForward()
+
+    auth = {
+      host: 'localhost',
+      username: 'user',
+      namespace: 'namespace',
+      device: 'device',
+      password: 'password'
+    }
+
+    settings = {
+      destinationAddr: '127.0.0.1',
+      destinationPort: 1080
+    }
+
+    mockClient = new ssh2.Client()
+    connection.client = mockClient
+
+    mockEvents = new EventEmitter()
+    connection.events = mockEvents
+  })
+
+  afterEach(() => {
+    vi.restoreAllMocks()
+  })
+
+  it('should emit error event on failed connection', () => {
+    const connectSpy = vi.spyOn(mockClient, 'connect')
+
+    const onSpy = vi.spyOn(mockClient, 'on').mockImplementation((event, callback: any) => {
+      if (event == 'error') callback(new Error('error from ssh client'))
+    })
+
+    connection.onError((err: any) => {
+      expect(err).toBeInstanceOf(Error)
+      expect(err.message).toBe('error from ssh client')
+    })
+
+    connection.connect(auth, settings)
+
+    expect(connectSpy).toHaveBeenCalled()
+    expect(onSpy).toHaveBeenCalledWith('error', expect.any(Function))
+  })
+
+  it('should emit auth event when authentication happens', () => {
+    const connectSpy = vi.spyOn(mockClient, 'connect')
+
+    const emitSpy = vi.spyOn(mockEvents, 'emit')
+
+    const onSpy = vi.spyOn(mockClient, 'on').mockImplementation((event, callback: any) => {
+      switch (event) {
+        case 'ready':
+          callback()
+          break
+      }
+    })
+
+    const listenMock = vi.fn().mockImplementation((_port, _addr, callback) => {
+      callback()
+    })
+
+    const createServerSpy = vi.spyOn(socks, 'createServer').mockImplementation((): any => {
+      return {
+        listen: listenMock,
+        on: vi.fn().mockReturnThis(),
+        close: vi.fn().mockReturnThis(),
+        useAuth: vi.fn().mockReturnThis()
+      }
+    })
+
+    connection.connect(auth, settings)
+
+    expect(connectSpy).toHaveBeenCalled()
+    expect(onSpy).toHaveBeenCalledWith('ready', expect.any(Function))
+    expect(emitSpy).toHaveBeenCalledWith('auth')
+    expect(createServerSpy).toHaveBeenCalled()
+  })
+
+  it('should emit error when forward out fails', () => {
+    const connectSpy = vi.spyOn(mockClient, 'connect')
+
+    const emitSpy = vi.spyOn(mockEvents, 'emit')
+
+    const onSpy = vi.spyOn(mockClient, 'on').mockImplementation((event, callback: any) => {
+      switch (event) {
+        case 'ready':
+          callback()
+          break
+      }
+    })
+
+    const listenMock = vi.fn().mockImplementation((_port, _addr, callback) => {
+      callback()
+    })
+
+    const denySpy = vi.fn()
+
+    const createServerSpy = vi
+      .spyOn(socks, 'createServer')
+      .mockImplementation((callback: (_info, _accept, _deny) => void): any => {
+        callback({}, undefined, denySpy)
+
+        return {
+          listen: listenMock,
+          on: vi.fn().mockReturnThis(),
+          close: vi.fn().mockReturnThis(),
+          useAuth: vi.fn().mockReturnThis()
+        }
+      })
+
+    const forwardOutSpy = vi
+      .spyOn(mockClient, 'forwardOut')
+      .mockImplementation((_srcPort, _srcAddr, _dstPort, _dstAddr, callback: any) => {
+        callback(new Error('failed to forward out'), null)
+      })
+
+    connection.onError((err: any) => {
+      expect(err).toBeInstanceOf(Error)
+      expect(err.message).toBe('failed to forward out')
+    })
+
+    connection.connect(auth, settings)
+
+    expect(connectSpy).toHaveBeenCalled()
+    expect(onSpy).toHaveBeenCalledWith('ready', expect.any(Function))
+    expect(emitSpy).toHaveBeenCalledWith('auth')
+    expect(createServerSpy).toHaveBeenCalled()
+    expect(forwardOutSpy).toHaveBeenCalled()
+    expect(onSpy).toHaveBeenCalledWith('error', expect.any(Function))
+    expect(denySpy).toHaveBeenCalled()
+  })
+
+  it('should emit listen when server is listening', () => {
+    const connectSpy = vi.spyOn(mockClient, 'connect')
+
+    const emitSpy = vi.spyOn(mockEvents, 'emit')
+
+    const onSpy = vi.spyOn(mockClient, 'on').mockImplementation((event, callback: any) => {
+      switch (event) {
+        case 'ready':
+          callback()
+          break
+      }
+    })
+
+    const listenMock = vi.fn().mockImplementation((_port, _addr, callback) => {
+      callback()
+    })
+
+    const denySpy = vi.fn()
+    const acceptSpy = vi.fn().mockImplementation(() => {
+      return new Stream()
+    })
+
+    const createServerSpy = vi
+      .spyOn(socks, 'createServer')
+      .mockImplementation((callback: (_info, _accept, _deny) => void): any => {
+        callback({}, acceptSpy, denySpy)
+
+        return {
+          listen: listenMock,
+          on: vi.fn().mockReturnThis(),
+          close: vi.fn().mockReturnThis(),
+          useAuth: vi.fn().mockReturnThis()
+        }
+      })
+
+    const forwardOutSpy = vi
+      .spyOn(mockClient, 'forwardOut')
+      .mockImplementation((_srcPort, _srcAddr, _dstPort, _dstAddr, callback: any) => {
+        callback(null, new Stream())
+      })
+
+    connection.connect(auth, settings)
+
+    expect(connectSpy).toHaveBeenCalled()
+    expect(onSpy).toHaveBeenCalledWith('ready', expect.any(Function))
+    expect(emitSpy).toHaveBeenCalledWith('auth')
+    expect(createServerSpy).toHaveBeenCalled()
+    expect(forwardOutSpy).toHaveBeenCalled()
+    expect(acceptSpy).toHaveBeenCalled()
+    expect(listenMock).toHaveBeenCalledWith(1080, '127.0.0.1', expect.any(Function))
+  })
+})