Add ShellConfig system to support more shells (home directory detection)

src/connection.ts

@@ -6,6 +6,7 @@ import { configMatches, getFlagBoolean, loadConfigs } from './config';
 import type { EnvironmentVariable, FileSystemConfig } from './fileSystemConfig';
 import { Logging, LOGGING_NO_STACKTRACE } from './logging';
 import type { SSHPseudoTerminal } from './pseudoTerminal';
+import { calculateShellConfig, ShellConfig, tryEcho } from './shellConfig';
 import type { SSHFileSystem } from './sshFileSystem';
 import { mergeEnvironment, toPromise } from './utils';
 
@@ -14,22 +15,14 @@ export interface Connection {
     actualConfig: FileSystemConfig;
     client: Client;
     home: string;
+    shellConfig: ShellConfig;
     environment: EnvironmentVariable[];
     terminals: SSHPseudoTerminal[];
     filesystems: SSHFileSystem[];
     pendingUserCount: number;
     idleTimer: NodeJS.Timeout;
 }
 
-async function tryGetHome(ssh: Client): Promise<string | null> {
-    const exec = await toPromise<ClientChannel>(cb => ssh.exec('echo "::sshfs:home:`echo ~`:"', cb));
-    let home = '';
-    exec.stdout.on('data', (chunk: any) => home += chunk);
-    await toPromise(cb => exec.on('close', cb));
-    if (!home) return null;
-    return home.match(/::sshfs:home:(.*?):/)?.[1] || null;
-}
-
 const TMP_PROFILE_SCRIPT = `
 if type code > /dev/null 2> /dev/null; then
     return 0;
@@ -152,8 +145,10 @@ export class ConnectionManager {
         const client = await createSSH(actualConfig);
         if (!client) throw new Error(`Could not create SSH session for '${name}'`);
         logging.info`Remote version: ${(client as any)._remoteVer || 'N/A'}`;
+        // Calculate shell config
+        const shellConfig = await calculateShellConfig(client, logging);
         // Query home directory
-        let home = await tryGetHome(client).catch((e: Error) => e);
+        let home = await tryEcho(client, shellConfig, '~').catch((e: Error) => e);
         if (typeof home !== 'string') {
             const [flagCH] = getFlagBoolean('CHECK_HOME', true, config.flags);
             logging.error('Could not detect home directory', LOGGING_NO_STACKTRACE);
@@ -192,7 +187,7 @@ export class ConnectionManager {
         // Set up the Connection object
         let timeoutCounter = 0;
         const con: Connection = {
-            config, client, actualConfig, home, environment,
+            config, client, actualConfig, home, shellConfig, environment,
             terminals: [],
             filesystems: [],
             pendingUserCount: 0,

src/pseudoTerminal.ts

@@ -138,7 +138,7 @@ export async function replaceVariablesRecursive<T>(object: T, handler: (value: s
 
 export async function createTerminal(options: TerminalOptions): Promise<SSHPseudoTerminal> {
     const { connection } = options;
-    const { actualConfig, client } = connection;
+    const { actualConfig, client, shellConfig } = connection;
     const onDidWrite = new vscode.EventEmitter<string>();
     const onDidClose = new vscode.EventEmitter<number>();
     const onDidOpen = new vscode.EventEmitter<void>();
@@ -169,17 +169,12 @@ export async function createTerminal(options: TerminalOptions): Promise<SSHPseud
                 let SHELL = '$SHELL';
                 // Add exports for environment variables if needed
                 const env = mergeEnvironment(connection.environment, options.environment);
-                commands.push(environmentToExportString(env));
+                commands.push(environmentToExportString(env, shellConfig.setEnv));
                 // Beta feature to add a "code <file>" command in terminals to open the file locally
                 if (getFlagBoolean('REMOTE_COMMANDS', false, actualConfig.flags)[0]) {
                     const profilePathEnv = env.find(e => e.key === 'KELVIN_SSHFS_PROFILE_PATH');
                     if (!profilePathEnv) throw new Error(`Missing KELVIN_SSHFS_PROFILE_PATH environment variable`);
-                    // For bash
-                    commands.push(`export ORIG_PROMPT_COMMAND="$PROMPT_COMMAND"`);
-                    commands.push(`export PROMPT_COMMAND='source "${profilePathEnv.value}" PC; $ORIG_PROMPT_COMMAND'`);
-                    // For sh
-                    commands.push(`export OLD_ENV="$ENV"`); // not actually used (yet?)
-                    commands.push(`export ENV="${profilePathEnv.value}"`);
+                    commands.push(shellConfig.setupRemoteCommands(profilePathEnv.value));
                 }
                 // Push the actual command or (default) shell command with replaced variables
                 if (options.command) {

src/shellConfig.ts

@@ -0,0 +1,95 @@
+import { posix as path } from 'path';
+import type { Client, ClientChannel } from "ssh2";
+import type { Logger } from "./logging";
+import { toPromise } from "./utils";
+
+export interface ShellConfig {
+    shell: string;
+    setEnv(key: string, value: string): string;
+    setupRemoteCommands(path: string): string;
+    embedSubstitutions(command: TemplateStringsArray, ...substitutions: (string | number)[]): string;
+}
+const KNOWN_SHELL_CONFIGS: Record<string, ShellConfig> = {}; {
+    const add = (shell: string,
+        setEnv: (key: string, value: string) => string,
+        setupRemoteCommands: (path: string) => string,
+        embedSubstitution: (command: TemplateStringsArray, ...substitutions: (string | number)[]) => string) => {
+        KNOWN_SHELL_CONFIGS[shell] = { shell, setEnv, setupRemoteCommands, embedSubstitutions: embedSubstitution };
+    }
+    // Ways to set an environment variable
+    const setEnvExport = (key: string, value: string) => `export ${key}=${value}`;
+    const setEnvSetGX = (key: string, value: string) => `set -gx ${key} ${value}`;
+    const setEnvSetEnv = (key: string, value: string) => `setenv ${key} ${value}`;
+    // Ways to set up the remote commands script auto-execution
+    const setupRemoteCommandsENV = (path: string) => [
+        `export OLD_ENV="$ENV"`, // OLD_ENV ignored for now
+        `export ENV="${path}"`].join('; ');
+    const setupRemoteCommandsPROMPT_COMMAND = (path: string) => [
+        `export ORIG_PROMPT_COMMAND="$PROMPT_COMMAND"`,
+        `export PROMPT_COMMAND='source "${path}" PC; $ORIG_PROMPT_COMMAND'`].join('; ');
+    const setupRemoteCommandsUnknown = () => 'echo "This shell does not yet have REMOTE_COMMANDS support"';
+    // Ways to embed a substitution
+    const embedSubstitutionsBackticks = (command: TemplateStringsArray, ...substitutions: (string | number)[]): string =>
+        '"' + substitutions.reduce((str, sub, i) => `${str}\`${sub}\`${command[i + 1]}`, command[0]) + '"';
+    const embedSubstitutionsFish = (command: TemplateStringsArray, ...substitutions: (string | number)[]) =>
+        substitutions.reduce((str, sub, i) => `${str}"(${sub})"${command[i + 1]}`, '"' + command[0]) + '"';
+    // Register the known shells
+    add('sh', setEnvExport, setupRemoteCommandsENV, embedSubstitutionsBackticks);
+    add('bash', setEnvExport, setupRemoteCommandsPROMPT_COMMAND, embedSubstitutionsBackticks);
+    add('rbash', setEnvExport, setupRemoteCommandsPROMPT_COMMAND, embedSubstitutionsBackticks);
+    add('ash', setEnvExport, setupRemoteCommandsENV, embedSubstitutionsBackticks);
+    add('dash', setEnvExport, setupRemoteCommandsENV, embedSubstitutionsBackticks);
+    add('ksh', setEnvExport, setupRemoteCommandsENV, embedSubstitutionsBackticks);
+    // Shells that we know `setEnv` and `embedSubstitution` for, but don't support  `setupRemoteCommands` for yet
+    add('zsh', setEnvExport, setupRemoteCommandsUnknown, embedSubstitutionsBackticks);
+    add('fish', setEnvSetGX, setupRemoteCommandsUnknown, embedSubstitutionsFish); // https://fishshell.com/docs/current/tutorial.html#autoloading-functions
+    add('csh', setEnvSetEnv, setupRemoteCommandsUnknown, embedSubstitutionsBackticks);
+    add('tcsh', setEnvSetEnv, setupRemoteCommandsUnknown, embedSubstitutionsBackticks);
+}
+
+export async function tryCommand(ssh: Client, command: string): Promise<string | null> {
+    const exec = await toPromise<ClientChannel>(cb => ssh.exec(command, cb));
+    const output = ['', ''] as [string, string];
+    exec.stdout.on('data', (chunk: any) => output[0] += chunk);
+    exec.stderr.on('data', (chunk: any) => output[1] += chunk);
+    await toPromise(cb => {
+        exec.once('error', cb);
+        exec.once('close', cb);
+    }).catch(e => {
+        if (typeof e !== 'number') throw e;
+        throw new Error(`Command '${command}' failed with exit code ${e}${output[1] ? `:\n${output[1].trim()}` : ''}`);
+    });
+    if (!output[0]) {
+        if (!output[1]) return null;
+        throw new Error(`Command '${command}' only produced stderr:\n${output[1].trim()}`);
+    }
+    return output[0];
+}
+
+export async function tryEcho(ssh: Client, shellConfig: ShellConfig, variable: string): Promise<string | null> {
+    const uniq = Date.now() % 1e5;
+    const output = await tryCommand(ssh, `echo ${shellConfig.embedSubstitutions`::${'echo ' + uniq}:echo_result:${`echo ${variable}`}:${'echo ' + uniq}::`}`);
+    return output?.match(`::${uniq}:echo_result:(.*?):${uniq}::`)?.[1] || null;
+}
+
+export async function calculateShellConfig(client: Client, logging?: Logger): Promise<ShellConfig> {
+    try {
+        const shellStdout = await tryCommand(client, 'echo :::SHELL:$SHELL:SHELL:::');
+        const shell = shellStdout?.match(/:::SHELL:([^$].*?):SHELL:::/)?.[1];
+        if (!shell) {
+            if (shellStdout) logging?.error(`Could not get $SHELL from following output:\n${shellStdout}`);
+            throw new Error('Could not get $SHELL');
+        }
+        const known = KNOWN_SHELL_CONFIGS[path.basename(shell)];
+        if (known) {
+            logging?.debug(`Detected known $SHELL '${shell}' (${known.shell})`);
+            return known;
+        } else {
+            logging?.warning(`Unrecognized $SHELL '${shell}', using default ShellConfig instead`);
+            return { ...KNOWN_SHELL_CONFIGS['sh'], shell };
+        }
+    } catch (e) {
+        logging && logging.error`Error calculating ShellConfig: ${e}`;
+        return { ...KNOWN_SHELL_CONFIGS['sh'], shell: '???' };
+    }
+}

src/utils.ts

@@ -86,33 +86,28 @@ function escapeBashValue(value: string) {
 }
 
 /** Convert an {@link EnvironmentVariable} array to a `export var1=val; export var2='escaped$val'` etc */
-export function environmentToExportString(env: EnvironmentVariable[]): string {
-    return env.map(({ key, value }) => `export ${escapeBashValue(key)}=${escapeBashValue(value)}`).join('; ');
+export function environmentToExportString(env: EnvironmentVariable[], createSetEnv: (key: string, value: string) => string): string {
+    return env.map(({ key, value }) => createSetEnv(escapeBashValue(key), escapeBashValue(value))).join('; ');
 }
 
 /** Returns a new {@link EnvironmentVariable} array with all the given environments merged into it, overwriting same-key variables */
-export function mergeEnvironment(env: EnvironmentVariable[], ...others: (EnvironmentVariable[] | Record<string, string> | undefined)[]): EnvironmentVariable[] {
-    const result = [...env];
-    for (const other of others) {
+export function mergeEnvironment(...environments: (EnvironmentVariable[] | Record<string, string> | undefined)[]): EnvironmentVariable[] {
+    const result = new Map<string, EnvironmentVariable>();
+    for (let other of environments) {
         if (!other) continue;
         if (Array.isArray(other)) {
-            for (const variable of other) {
-                const index = result.findIndex(v => v.key === variable.key);
-                if (index === -1) result.push(variable);
-                else result[index] = variable;
-            }
+            for (const variable of other) result.set(variable.key, variable);
         } else {
             for (const [key, value] of Object.entries(other)) {
-                result.push({ key, value });
+                result.set(key, { key, value });
             }
         }
     }
-    return result;
+    return [...result.values()];
 }
 
 /** Joins the commands together using the given separator. Automatically ignores `undefined` and empty strings */
 export function joinCommands(commands: string | string[] | undefined, separator: string): string | undefined {
-    if (!commands) return undefined;
     if (typeof commands === 'string') return commands;
-    return commands.filter(c => c && c.trim()).join(separator);
+    return commands?.filter(c => c?.trim()).join(separator);
 }