refactor: handle errors and closed events on servers

shellhub-io · Oct 9, 2024 · 39456e5 · 39456e5
1 parent 36dfbb6
commit 39456e5
Show file tree

Hide file tree

Showing 5 changed files with 134 additions and 85 deletions.
diff --git a/src/preload/index.d.ts b/src/preload/index.d.ts
@@ -4,8 +4,8 @@ import { ElectronAPI } from '@electron-toolkit/preload'
 import EventEmitter from 'events'
 
 export interface SSH {
-  localPortForward(): SSHConnection
-  dynamicPortForward(): SSHConnection
+  localPortForward(settings: any): SSHConnection
+  dynamicPortForward(settings: any): SSHConnection
 }
 
 declare global {

diff --git a/src/preload/index.ts b/src/preload/index.ts
@@ -11,8 +11,8 @@ import {
 const api = {}
 
 const ssh: SSH = {
-  localPortForward: (): SSHConnection => {
-    const localPortForwardInstance = new SSHConnectionLocalPortForward()
+  localPortForward: (settings: any): SSHConnection => {
+    const localPortForwardInstance = new SSHConnectionLocalPortForward(settings)
 
     return {
       events: localPortForwardInstance.events,
@@ -25,8 +25,8 @@ const ssh: SSH = {
       onDisconnect: localPortForwardInstance.onDisconnect.bind(localPortForwardInstance)
     }
   },
-  dynamicPortForward: (): SSHConnection => {
-    const dynamicPortForwardInstance = new SSHConnectionDynamicPortForward()
+  dynamicPortForward: (settings: any): SSHConnection => {
+    const dynamicPortForwardInstance = new SSHConnectionDynamicPortForward(settings)
 
     return {
       events: dynamicPortForwardInstance.events,

diff --git a/src/preload/ssh/index.d.ts b/src/preload/ssh/index.d.ts
@@ -33,7 +33,7 @@ export type SSHDynamicPortForwardSettings = {
 export interface SSHConnection {
   events: SSHEmitter
   client: ssh2.Client
-  connect(auth: SSHConnectionAuth, settings: any): void
+  connect(auth: SSHConnectionAuth): void
   disconnect(): void
   onAuth(callback: any): void
   onConnect(callback: any): void

diff --git a/src/preload/ssh/ssh.test.ts b/src/preload/ssh/ssh.test.ts
@@ -1,18 +1,18 @@
-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'
+import { SSHConnectionLocalPortForward, SSHConnectionDynamicPortForward } from './ssh'
+import { describe, it, expect, beforeEach, vi, afterEach } from 'vitest'
+import ssh2 from 'ssh2'
+import socks from 'socksv5'
 
 vi.mock('ssh2')
 vi.mock('node:net')
 vi.mock('socksv5')
 
 describe('SSH Interface', () => {
   it('localPortForward should return a SSHConnection instance', () => {
-    const ssh = new SSHConnectionLocalPortForward()
+    const ssh = new SSHConnectionLocalPortForward({})
     expect(ssh).toHaveProperty('connect')
     expect(ssh).toHaveProperty('disconnect')
     expect(ssh).toHaveProperty('onAuth')
@@ -22,7 +22,7 @@ describe('SSH Interface', () => {
   })
 
   it('dynamicPortForward should return a SSHConnection instance', () => {
-    const ssh = new SSHConnectionDynamicPortForward()
+    const ssh = new SSHConnectionDynamicPortForward({})
     expect(ssh).toHaveProperty('connect')
     expect(ssh).toHaveProperty('disconnect')
     expect(ssh).toHaveProperty('onAuth')
@@ -41,8 +41,6 @@ describe('SSHConnectionLocalPortForward', () => {
   let mockEvents
 
   beforeEach(() => {
-    connection = new SSHConnectionLocalPortForward()
-
     auth = {
       host: 'localhost',
       username: 'user',
@@ -58,6 +56,8 @@ describe('SSHConnectionLocalPortForward', () => {
       destinationPort: 80
     }
 
+    connection = new SSHConnectionLocalPortForward(settings)
+
     mockClient = new ssh2.Client()
     connection.client = mockClient
 
@@ -199,8 +199,6 @@ describe('SSHConnectionDynamicPortForward', () => {
   let mockEvents
 
   beforeEach(() => {
-    connection = new SSHConnectionDynamicPortForward()
-
     auth = {
       host: 'localhost',
       username: 'user',
@@ -214,6 +212,8 @@ describe('SSHConnectionDynamicPortForward', () => {
       destinationPort: 1080
     }
 
+    connection = new SSHConnectionDynamicPortForward(settings)
+
     mockClient = new ssh2.Client()
     connection.client = mockClient
 

diff --git a/src/preload/ssh/ssh.ts b/src/preload/ssh/ssh.ts
@@ -8,58 +8,92 @@ export class SSHConnectionLocalPortForward implements SSHConnection {
   client: any
   events: SSHEmitter
 
-  constructor() {
+  constructor(settings: any) {
     this.client = new ssh2.Client()
     this.events = new SSHEmitter()
-  }
 
-  connect(auth: SSHConnectionAuth, settings: any) {
-    try {
-      this.client.on('ready', () => {
-        console.info('SSH connection authenticated')
-        this.events.emit(SSHEvent.Auth)
-
-        this.client.forwardOut(
-          settings.sourceAddr,
-          settings.sourcePort,
-          settings.destinationAddr,
-          settings.destinationPort,
-          (err: Error, stream: any) => {
+    const server = net.createServer((socket: net.Socket) => {
+      this.client.forwardOut(
+        settings.sourceAddr,
+        settings.sourcePort,
+        settings.destinationAddr,
+        settings.destinationPort,
+        (err: Error, channel: any) => {
+          try {
             if (err) {
               this.events.emit(SSHEvent.Error, err)
 
               return
             }
 
-            const server = net.createServer((client) => {
-              client.pipe(stream).pipe(client)
+            channel.on('close', () => {
+              console.log('channel closed')
             })
 
-            server.on('close', () => {
-              this.events.emit(SSHEvent.Disconnect)
+            channel.on('error', (err: Error) => {
+              console.log('channel error', err)
             })
 
-            server.listen(settings.sourcePort, settings.sourceAddr, () => {
-              console.log(
-                `Local port forward started from ${settings.sourceAddr}:${settings.sourcePort} to ${settings.destinationAddr}:${settings.destinationPort}.`
-              )
-              this.events.emit(SSHEvent.Connect, settings.sourcePort, settings.sourceAddr)
-            })
+            socket.pipe(channel).pipe(socket)
+          } catch (err: unknown) {
+            console.debug('SSH channel throw an exception', err)
           }
-        )
-      })
+        }
+      )
+    })
+
+    server.on('close', () => {
+      console.log('Local port forward closed')
+    })
+
+    server.on('error', (err: Error) => {
+      console.debug('error called on server', err)
+    })
+
+    this.client.on('ready', () => {
+      try {
+        console.info('SSH connection authenticated')
+        this.events.emit(SSHEvent.Auth)
 
-      this.client.on('error', (err: unknown) => {
+        server.listen(settings.sourcePort, settings.sourceAddr, () => {
+          console.log(
+            `Local port forward started from ${settings.sourceAddr}:${settings.sourcePort} to ${settings.destinationAddr}:${settings.destinationPort}.`
+          )
+
+          this.events.emit(SSHEvent.Connect, settings.sourcePort, settings.sourceAddr)
+        })
+      } catch (err: unknown) {
+        console.error('SSH ready throw an exception', err)
+
+        server.close()
         this.events.emit(SSHEvent.Error, err)
-      })
+      }
+    })
+
+    this.client.on('close', () => {
+      console.debug('close called on client')
+
+      server.close()
+      this.events.emit(SSHEvent.Disconnect)
+    })
+
+    this.client.on('error', (err: unknown) => {
+      console.error('error called on client', err)
+
+      server.close()
+      this.events.emit(SSHEvent.Error, err)
+    })
+  }
 
+  connect(auth: SSHConnectionAuth) {
+    try {
       this.client.connect({
         host: auth.host,
         username: `${auth.username}@${auth.namespace}.${auth.device}`,
         password: auth.password
       })
     } catch (err: unknown) {
-      console.error('SSH client throw a expection', err)
+      console.error('SSH client throw an exception', err)
 
       this.events.emit(SSHEvent.Error, err)
     }
@@ -68,7 +102,6 @@ export class SSHConnectionLocalPortForward implements SSHConnection {
   disconnect() {
     if (this.client) {
       this.client.end()
-      this.events.emit(SSHEvent.Disconnect)
     }
   }
 
@@ -100,55 +133,71 @@ export class SSHConnectionDynamicPortForward implements SSHConnection {
   client: any
   events: SSHEmitter
 
-  constructor() {
+  constructor(settings: any) {
     this.client = new ssh2.Client()
     this.events = new SSHEmitter()
-  }
-
-  connect(auth: SSHConnectionAuth, settings: any) {
-    try {
-      this.client.on('ready', () => {
-        console.info('SSH connection established')
-        this.events.emit(SSHEvent.Auth)
 
-        const server = socks.createServer(
-          (info: Info, accept: (_: boolean) => any, deny: () => void) => {
-            this.client.forwardOut(
-              info.srcAddr,
-              info.srcPort,
-              info.dstAddr,
-              info.dstPort,
-              (err: Error, stream: Stream) => {
-                if (err) {
-                  console.debug('SSH forwarding error:', err)
-                  this.events.emit(SSHEvent.Error, err)
-                  deny()
-                  return
-                }
-
-                const client = accept(true)
-                stream.pipe(client).pipe(stream)
+    this.client.on('ready', () => {
+      console.info('SSH connection established')
+      this.events.emit(SSHEvent.Auth)
+
+      const server = socks.createServer(
+        (info: Info, accept: (_: boolean) => any, deny: () => void) => {
+          this.client.forwardOut(
+            info.srcAddr,
+            info.srcPort,
+            info.dstAddr,
+            info.dstPort,
+            (err: Error, channel: Stream) => {
+              if (err) {
+                console.debug('SSH forwarding error:', err)
+                this.events.emit(SSHEvent.Error, err)
+                deny()
+
+                return
               }
-            )
-          }
-        )
 
-        server.on('error', (err: unknown) => {
-          this.events.emit(SSHEvent.Error, err)
-        })
+              const client = accept(true)
+              channel.pipe(client).pipe(channel)
+            }
+          )
+        }
+      )
 
-        server.useAuth(socks.auth.None())
+      server.on('close', () => {
+        console.log('Dynamic port forward closed')
+      })
 
-        server.listen(settings.destinationPort, settings.destinationAddr, () => {
-          console.log('Server listening on', settings.destinationPort, settings.destinationAddr)
-          this.events.emit(SSHEvent.Connect, settings.destinationPort, settings.destinationAddr)
-        })
+      server.on('error', (err: Error) => {
+        console.debug('error called on server', err)
       })
 
-      this.client.on('error', (err: unknown) => {
-        this.events.emit(SSHEvent.Error, err)
+      server.useAuth(socks.auth.None())
+
+      server.listen(settings.destinationPort, settings.destinationAddr, () => {
+        console.log(
+          `Dynamic port forward started from ${settings.destinationAddr}:${settings.destinationPort}.`
+        )
+
+        this.events.emit(SSHEvent.Connect, settings.destinationPort, settings.destinationAddr)
       })
+    })
+
+    this.client.on('close', () => {
+      console.debug('close called on client')
+
+      this.events.emit(SSHEvent.Disconnect)
+    })
+
+    this.client.on('error', (err: unknown) => {
+      console.error('error called on client', err)
+
+      this.events.emit(SSHEvent.Error, err)
+    })
+  }
 
+  connect(auth: SSHConnectionAuth) {
+    try {
       this.client.connect({
         host: auth.host,
         username: `${auth.username}@${auth.namespace}.${auth.device}`,