Close conn based on piped conn IDE-50 (#50)

* Close conn based on piped conn IDE-50

* Close session after auth failed

* 💄

---------

Co-authored-by: Jean Pierre <[email protected]>

Author: mustard-mh

src/daemonStarter.ts

@@ -61,8 +61,8 @@ export async function tryStartDaemon(logService: ILogService, options?: Partial<
     const daemon = spawn(process.execPath, args, {
         detached: true,
         stdio: 'ignore',
-        env: process.env
+        env: { ...process.env, ELECTRON_RUN_AS_NODE: '1' },
     });
     daemon.unref();
     return daemon;
-}
+}

src/local-ssh/ipc/localssh.ts

@@ -81,10 +81,10 @@ export class LocalSSHServiceImpl implements LocalSSHServiceImplementation {
     }
 
     public async getWorkspaceAuthInfo(workspaceId: string) {
-        return retryWithStop(async (stop) => {
-            const getAuthInfo = async (id: string, client: Client<ExtensionServiceDefinition>) => {
+        return retryWithStop((stop) => {
+            const getAuthInfo = (id: string, client: Client<ExtensionServiceDefinition>) => {
                 try {
-                    return await client.getWorkspaceAuthInfo({ workspaceId });
+                    return client.getWorkspaceAuthInfo({ workspaceId });
                 } catch (e) {
                     if (e instanceof ServerError) {
                         if (e.code === Status.UNAVAILABLE && e.details.startsWith('workspace is not running')) {
@@ -95,7 +95,7 @@ export class LocalSSHServiceImpl implements LocalSSHServiceImplementation {
                     throw e;
                 }
             };
-            return await Promise.any(this.extensionServices.map(ext => getAuthInfo(ext.id, ext.client)));
+            return Promise.any(this.extensionServices.map(ext => getAuthInfo(ext.id, ext.client)));
         }, 200, 3);
     }
 

src/local-ssh/logger.ts

@@ -6,14 +6,8 @@
 import winston from 'winston';
 import { ILogService } from '../services/logService';
 
-export enum ExitCode {
-    OK = 0,
-    ListenPortFailed = 100,
-    UnexpectedError = 101,
-    InvalidOptions = 102,
-}
-
 const DefaultLogFormatter = winston.format.combine(winston.format.timestamp(), winston.format.errors({ stack: true }), winston.format.simple());
+
 export class Logger implements ILogService {
     private logger: winston.Logger;
 
@@ -47,4 +41,4 @@ export class Logger implements ILogService {
     show(): void {
         // no-op
     }
-}
+}

src/local-ssh/server.ts

@@ -9,7 +9,7 @@ import { LocalSSHServiceImpl, startLocalSSHService } from './ipc/localssh';
 import { SupervisorSSHTunnel } from './sshTunnel';
 import { ILogService } from '../services/logService';
 import { SshServer, SshClient } from '@microsoft/dev-tunnels-ssh-tcp';
-import { NodeStream, SshClientCredentials, SshClientSession, SshSessionConfiguration } from '@microsoft/dev-tunnels-ssh';
+import { NodeStream, SshClientCredentials, SshClientSession, SshDisconnectReason, SshSessionConfiguration } from '@microsoft/dev-tunnels-ssh';
 import { importKeyBytes } from '@microsoft/dev-tunnels-ssh-keys';
 import { parsePrivateKey } from 'sshpk';
 import { PipeExtensions } from './patch/pipeExtension';
@@ -23,14 +23,15 @@ export class LocalSSHGatewayServer {
 	private localsshService!: LocalSSHServiceImpl;
 	private localsshServiceServer?: GrpcServer;
 	private server?: SshServer;
+	private clientCount = 0;
 
 	constructor(
 		private readonly logger: ILogService,
 		private readonly port: number,
 		private readonly ipcPort: number,
 	) { }
 
-	async authenticateClient(clientUsername: string) {
+	private async authenticateClient(clientUsername: string) {
 		const workspaceInfo = await this.localsshService.getWorkspaceAuthInfo(clientUsername).catch(e => {
 			this.logger.error(e, 'failed to get workspace auth info');
 			/*
@@ -69,22 +70,32 @@ export class LocalSSHGatewayServer {
 
 			server.onSessionOpened((session) => {
 				let pipeSession: SshClientSession;
+				this.clientCount += 1;
 				session.onAuthenticating((e) => {
-					e.authenticationPromise = new Promise((resolve, reject) => {
-						this.authenticateClient(e.username!).then(async s => {
+					e.authenticationPromise = this.authenticateClient(e.username!).then(s => {
 							this.logger.info('authenticate with ' + e.username);
 							pipeSession = s;
-							resolve(new Object());
+							return {};
 						}).catch(e => {
 							this.logger.error(e, 'failed to authenticate client');
 							// TODO not sure how to get gitpod host here
 							// this.localsshService.sendErrorReport(e.username, undefined, e, 'failed to authenticate client');
-							reject(null);
+							session.close(SshDisconnectReason.hostNotAllowedToConnect, 'auth failed or workspace is not running');
+							return null;
 						});
-					});
 				});
-				session.onClientAuthenticated(() => {
-					PipeExtensions.pipeSession(session, pipeSession);
+				session.onClientAuthenticated(async () => {
+					try {
+						await PipeExtensions.pipeSession(session, pipeSession);
+					} catch (e) {
+						this.logger.error(e, 'pipe session ended with error');
+					} finally {
+						session.close(SshDisconnectReason.connectionLost, 'pipe session ended');
+					}
+				});
+				session.onClosed(() => {
+					this.clientCount -= 1;
+					this.logger.debug('current connecting client count: ' + this.clientCount);
 				});
 			});
 			await server.acceptSessions(this.port, '127.0.0.1');
@@ -180,4 +191,3 @@ export class LocalSSHGatewayServer {
 		this.localsshServiceServer?.shutdown();
 	}
 }
-