Fix/ssh forwarding agent

src/authResolver.ts

@@ -17,7 +17,6 @@ import { disposeAll } from './common/disposable';
 import { installCodeServer, ServerInstallError, findServerInstallPath } from './serverSetup';
 import { isWindows } from './common/platform';
 import * as os from 'os';
-import { isNullable } from '@zokugun/is-it-type';
 import { ServerVersion } from './serverConfig';
 
 const PASSWORD_RETRY_COUNT = 3;
@@ -110,6 +109,7 @@ export class RemoteSSHResolver implements vscode.RemoteAuthorityResolver, vscode
     private sshConnection: SSHConnection | undefined;
     private sshAgentSock: string | undefined;
     private proxyCommandProcess: cp.ChildProcessWithoutNullStreams | undefined;
+    private agentForwardSession: ssh2.ClientChannel | undefined;
 
     private socksTunnel: SSHTunnelConfig | undefined;
     private tunnels: TunnelInfo[] = [];
@@ -246,7 +246,21 @@ export class RemoteSSHResolver implements vscode.RemoteAuthorityResolver, vscode
 
                 const envVariables: Record<string, string | null> = {};
                 if (agentForward) {
-                    envVariables['SSH_AUTH_SOCK'] = null;
+                    // The agent-forwarding socket sshd creates is scoped to the ssh channel that
+                    // requested it and is torn down as soon as that channel closes. The server
+                    // install/start script runs on its own short-lived exec channel, so any
+                    // SSH_AUTH_SOCK it reports is already stale by the time we get here. Keep a
+                    // dedicated channel open for the lifetime of the connection instead, and use
+                    // its socket path everywhere else (terminals, extension host). Agent
+                    // forwarding is best-effort: a failure here must not prevent connecting.
+                    try {
+                        const remoteAgentSock = await this.openAgentForwardSession();
+                        if (remoteAgentSock) {
+                            envVariables['SSH_AUTH_SOCK'] = remoteAgentSock;
+                        }
+                    } catch (e) {
+                        this.logger.error(`Failed to setup agent forwarding`, e);
+                    }
                 }
 
                 // Find the custom install path for this hostname (supports wildcards)
@@ -257,20 +271,14 @@ export class RemoteSSHResolver implements vscode.RemoteAuthorityResolver, vscode
                     serverDownloadUrlTemplate,
                     serverVersion,
                     defaultExtensions,
-                    Object.keys(envVariables),
+                    [],
                     remotePlatformMap[sshDest.hostname],
                     remoteServerListenOnSocket,
                     customInstallPath,
                     this.logger,
                     this.context.extensionPath
                 );
 
-                for (const key of Object.keys(envVariables)) {
-                    if (!isNullable(installResult[key])) {
-                        envVariables[key] = String(installResult[key]);
-                    }
-                }
-
                 // Update terminal env variables
                 this.context.environmentVariableCollection.persistent = false;
                 for (const [key, value] of Object.entries(envVariables)) {
@@ -335,6 +343,58 @@ export class RemoteSSHResolver implements vscode.RemoteAuthorityResolver, vscode
         });
     }
 
+    private openAgentForwardSession(): Promise<string | undefined> {
+        // No pty here on purpose: a pty echoes back whatever is written to the
+        // channel before the remote shell executes it, which would otherwise be
+        // mistaken for the command's actual output. `exec cat` keeps the process
+        // (and therefore the channel's agent-forwarding socket) alive indefinitely
+        // after printing the socket path once.
+        return this.sshConnection!.execChannel('echo "$SSH_AUTH_SOCK"; exec cat').then(channel => {
+            this.agentForwardSession?.close();
+            this.agentForwardSession = channel;
+
+            return new Promise<string | undefined>(resolve => {
+                let buffer = '';
+                let resolved = false;
+
+                const finish = (value: string | undefined) => {
+                    if (!resolved) {
+                        resolved = true;
+                        channel.removeListener('data', onData);
+                        channel.removeListener('close', onClose);
+                        clearTimeout(timer);
+                        resolve(value);
+                    }
+                };
+
+                const onData = (data: Buffer) => {
+                    buffer += data.toString();
+                    const newlineIdx = buffer.indexOf('\n');
+                    if (newlineIdx < 0) {
+                        return;
+                    }
+                    // A forwarded SSH_AUTH_SOCK is always an absolute path. Anything else
+                    // (e.g. a non-POSIX remote echoing the command back verbatim) is rejected
+                    // rather than exported as a bogus value.
+                    const value = buffer.slice(0, newlineIdx).trim();
+                    finish(value.startsWith('/') ? value : undefined);
+                };
+
+                // On a non-POSIX remote the `echo`d line ends the command and the channel
+                // closes without a usable path; resolve now instead of waiting for the timeout.
+                const onClose = () => finish(undefined);
+
+                const timer = setTimeout(() => {
+                    this.logger.trace('Timed out waiting for remote SSH_AUTH_SOCK');
+                    finish(undefined);
+                }, 5000);
+
+                channel.on('data', onData);
+                channel.on('close', onClose);
+            });
+        });
+    }
+
     private async openTunnel(localPort: number, remotePortOrSocketPath: number | string) {
         localPort = localPort > 0 ? localPort : await findRandomPort();
 
@@ -522,6 +582,8 @@ export class RemoteSSHResolver implements vscode.RemoteAuthorityResolver, vscode
 
     dispose() {
         disposeAll(this.tunnels);
+        this.agentForwardSession?.close();
+        this.agentForwardSession = undefined;
         // If there's proxy connections then just close the parent connection
         if (this.proxyConnections.length) {
             this.proxyConnections[0].close();

src/ssh/sshConnection.ts

@@ -90,6 +90,19 @@ export default class SSHConnection extends EventEmitter {
         });
     }
 
+    /**
+     * Start a command and return its channel as soon as it's opened, without
+     * waiting for it to finish. Unlike shell(), no pty is allocated, so there's
+     * no local echo of anything written to the channel.
+     */
+    execChannel(cmd: string, options: ExecOptions = {}): Promise<ClientChannel> {
+        return this.connect().then(() => {
+            return new Promise<ClientChannel>((resolve, reject) => {
+                this.sshConnection!.exec(cmd, options, (err, stream) => err ? reject(err) : resolve(stream));
+            });
+        });
+    }
+
     /**
      * Exec a command
      */

test/forward-agent.test.ts

@@ -0,0 +1,129 @@
+import { execFileSync } from 'node:child_process';
+import { randomUUID } from 'node:crypto';
+import fse from '@zokugun/fs-extra-plus/sync';
+import { vol } from 'memfs';
+import { afterAll, beforeAll, expect, it } from 'vitest';
+import SSHConnection from '../src/ssh/sshConnection';
+import { RemoteSSHResolver, getRemoteAuthority } from './rewires/remote';
+import { Log } from './mocks/logger';
+import * as vscode from './mocks/vscode';
+import { runDocker } from './utils/run-docker';
+import { getMappedPort } from './utils/get-mapped-port';
+import { waitForSSHReady } from './utils/wait-for-ssh-ready';
+
+const SERVER_SETUP = fse.readFile('./src/scripts/server-setup.sh', 'utf8').value!;
+
+const PRODUCT_JSON = JSON.stringify({
+  nameShort: 'VSCodium',
+  nameLong: 'VSCodium',
+  applicationName: 'codium',
+  quality: 'stable',
+  commit: '4c0b0c6cc561d2d3636d1ec250935431876ce4dc',
+  version: '1.126.04524',
+  serverApplicationName: 'codium-server',
+  serverDataFolderName: '.vscodium-server',
+  serverDownloadUrlTemplate: 'https://github.com/VSCodium/vscodium/releases/download/1.126.04524/vscodium-reh-${os}-${arch}-1.126.04524.tar.gz',
+});
+
+const IMAGE = 'local-ubuntu-bash';
+const USERNAME = 'openremotessh';
+const PASSWORD = 'openremotessh';
+
+const containerName = `open-remote-ssh-test-${randomUUID()}`;
+
+let authSock: string;
+let agentPid: string;
+let hostPort: number;
+
+beforeAll(async () => {
+  vol.reset();
+
+  const agentOutput = execFileSync('ssh-agent', ['-s'], { encoding: 'utf8' });
+  authSock = /SSH_AUTH_SOCK=([^;]+);/.exec(agentOutput)![1];
+  agentPid = /SSH_AGENT_PID=(\d+);/.exec(agentOutput)![1];
+
+  runDocker(['rm', '-f', containerName], true);
+
+  runDocker([
+    'run',
+    '--detach',
+    '--rm',
+    '--name',
+    containerName,
+    '--publish',
+    '2222',
+    '--env',
+    `USER_NAME=${USERNAME}`,
+    '--env',
+    `USER_PASSWORD=${PASSWORD}`,
+    '--env',
+    'PASSWORD_ACCESS=true',
+    '--env',
+    'SUDO_ACCESS=false',
+    '--env',
+    'LOG_STDOUT=true',
+    IMAGE,
+  ]);
+
+  hostPort = getMappedPort(containerName);
+
+  await waitForSSHReady(USERNAME, PASSWORD, hostPort, 60_000);
+}, 120_000);
+
+afterAll(() => {
+  runDocker(['rm', '-f', containerName], true);
+  execFileSync('ssh-agent', ['-k'], { env: { ...process.env, SSH_AGENT_PID: agentPid, SSH_AUTH_SOCK: authSock } });
+});
+
+it('forwards the agent through a socket that stays alive', async () => {
+  vol.fromJSON({
+    '/etc/ssh/ssh_config': [
+      'Host test',
+      '  HostName 127.0.0.1',
+      `  Port ${hostPort}`,
+      `  User ${USERNAME}`,
+      `  Password ${PASSWORD}`,
+      '  ForwardAgent yes',
+      `  IdentityAgent ${authSock}`,
+    ].join('\n'),
+    '/bin/vscodium/app/product.json': PRODUCT_JSON,
+    '/data/vscodium/extensions/open-remote-ssh/src/scripts/server-setup.sh': SERVER_SETUP,
+  });
+
+  vscode.window.setPassword(PASSWORD);
+
+  const logger = new Log('Remote - SSH');
+  const extContext = new vscode.ExtensionContext();
+  const remoteSSHResolver = new RemoteSSHResolver(extContext, logger);
+  const remoteContext = new vscode.RemoteAuthorityResolverContext();
+  const authority = getRemoteAuthority('test');
+  const result = await remoteSSHResolver.resolve(authority, remoteContext);
+
+  expect(result).toBeDefined();
+  expect(result.host).to.eql('127.0.0.1');
+
+  const remoteAuthSock = result.extensionHostEnv?.SSH_AUTH_SOCK;
+  expect(remoteAuthSock, 'SSH_AUTH_SOCK should be exported to the extension host').toBeTypeOf('string');
+  expect(remoteAuthSock!.startsWith('/'), 'forwarded SSH_AUTH_SOCK should be an absolute path').toBe(true);
+
+  expect(extContext.environmentVariableCollection.replace).toHaveBeenCalledWith('SSH_AUTH_SOCK', remoteAuthSock);
+
+  const probe = new SSHConnection({
+    host: '127.0.0.1',
+    port: hostPort,
+    username: USERNAME,
+    password: PASSWORD,
+    reconnect: false,
+    readyTimeout: 10_000,
+    strictVendor: false,
+  });
+
+  try {
+    const { stdout } = await probe.exec(`test -S "${remoteAuthSock}" && echo LIVE || echo DEAD`);
+    expect(stdout.trim(), 'forwarded socket should still exist after resolve').toContain('LIVE');
+  } finally {
+    await probe.close();
+  }
+
+  remoteSSHResolver.dispose();
+}, 60_000);