Skip to content

Commit

Permalink
chore: add tests for dynamic port forward
Browse files Browse the repository at this point in the history
  • Loading branch information
henrybarreto committed Oct 8, 2024
1 parent aebe888 commit e595554
Showing 1 changed file with 196 additions and 0 deletions.
196 changes: 196 additions & 0 deletions src/preload/ssh/ssh.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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', () => {
Expand Down Expand Up @@ -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))
})
})

0 comments on commit e595554

Please sign in to comment.