Add support for environment variables (closes #241)

src/connection.ts

@@ -1,7 +1,7 @@
 import type { Client } from 'ssh2';
 import * as vscode from 'vscode';
-import { configMatches, loadConfigs } from './config';
-import type { FileSystemConfig } from './fileSystemConfig';
+import { configMatches, getFlagBoolean, loadConfigs } from './config';
+import type { EnvironmentVariable, FileSystemConfig } from './fileSystemConfig';
 import { Logging } from './logging';
 import type { SSHPseudoTerminal } from './pseudoTerminal';
 import type { SSHFileSystem } from './sshFileSystem';
@@ -10,12 +10,48 @@ export interface Connection {
     config: FileSystemConfig;
     actualConfig: FileSystemConfig;
     client: Client;
+    environment: EnvironmentVariable[];
     terminals: SSHPseudoTerminal[];
     filesystems: SSHFileSystem[];
     pendingUserCount: number;
     idleTimer: NodeJS.Timeout;
 }
 
+export function mergeEnvironment(env: EnvironmentVariable[], ...others: (EnvironmentVariable[] | Record<string, string> | undefined)[]): EnvironmentVariable[] {
+    const result = [...env];
+    for (const other of others) {
+        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;
+            }
+        } else {
+            for (const [key, value] of Object.entries(other)) {
+                result.push({ key, value });
+            }
+        }
+    }
+    return result;
+}
+
+// https://stackoverflow.com/a/20053121 way 1
+const CLEAN_BASH_VALUE_REGEX = /^[\w-/\\]+$/;
+function escapeBashValue(value: string) {
+    if (CLEAN_BASH_VALUE_REGEX.test(value)) return value;
+    return `'${value.replace(/'/g, `'\\''`)}'`;
+}
+export function environmentToExportString(env: EnvironmentVariable[]): string {
+    return env.map(({ key, value }) => `export ${escapeBashValue(key)}=${escapeBashValue(value)}`).join('; ');
+}
+
+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);
+}
+
 export class ConnectionManager {
     protected onConnectionAddedEmitter = new vscode.EventEmitter<Connection>();
     protected onConnectionRemovedEmitter = new vscode.EventEmitter<Connection>();
@@ -51,9 +87,10 @@ export class ConnectionManager {
         if (!actualConfig) throw new Error('Connection cancelled');
         const client = await createSSH(actualConfig);
         if (!client) throw new Error(`Could not create SSH session for '${name}'`);
+        const environment: EnvironmentVariable[] = mergeEnvironment([], config.environment);
         let timeoutCounter = 0;
         const con: Connection = {
-            config, client, actualConfig,
+            config, client, actualConfig, environment,
             terminals: [],
             filesystems: [],
             pendingUserCount: 0,

src/fileSystemConfig.ts

@@ -8,6 +8,12 @@ export interface ProxyConfig {
 
 export type ConfigLocation = number | string;
 
+/** Might support conditional stuff later, although ssh2/OpenSSH might not support that natively */
+export interface EnvironmentVariable {
+    key: string;
+    value: string;
+}
+
 export function formatConfigLocation(location?: ConfigLocation): string {
   if (!location) return 'Unknown location';
   if (typeof location === 'number') {
@@ -100,6 +106,8 @@ export interface FileSystemConfig extends ConnectConfig {
   terminalCommand?: string | string[];
   /** The command(s) to run when a `ssh-shell` task gets run. Defaults to the placeholder `$COMMAND`. Internally the command `cd ...` is run first */
   taskCommand?: string | string[];
+  /** An object with environment variables to add to the SSH connection. Affects the whole connection thus all terminals */
+  environment?: EnvironmentVariable[] | Record<string, string>;
   /** The filemode to assign to created files */
   newFileMode?: number | string;
   /** Whether this config was created from an instant connection string. Enables fuzzy matching for e.g. PuTTY, config-by-host, ... */

src/manager.ts

@@ -3,7 +3,7 @@ import * as path from 'path';
 import type { Client, ClientChannel } from 'ssh2';
 import * as vscode from 'vscode';
 import { getConfig, getFlagBoolean, loadConfigsRaw } from './config';
-import { Connection, ConnectionManager } from './connection';
+import { Connection, ConnectionManager, joinCommands } from './connection';
 import type { FileSystemConfig } from './fileSystemConfig';
 import { getRemotePath } from './fileSystemRouter';
 import { Logging, LOGGING_NO_STACKTRACE } from './logging';
@@ -143,7 +143,7 @@ export class Manager implements vscode.TaskProvider, vscode.TerminalLinkProvider
     const workingDirectory = uri && getRemotePath(con.actualConfig, uri);
     // Create pseudo terminal
     this.connectionManager.update(con, con => con.pendingUserCount++);
-    const pty = await createTerminal({ client: con.client, config: con.actualConfig, workingDirectory });
+    const pty = await createTerminal({ connection: con, workingDirectory });
     pty.onDidClose(() => this.connectionManager.update(con, con => con.terminals = con.terminals.filter(t => t !== pty)));
     this.connectionManager.update(con, con => (con.terminals.push(pty), con.pendingUserCount--));
     // Create and show the graphical representation
@@ -176,7 +176,7 @@ export class Manager implements vscode.TaskProvider, vscode.TerminalLinkProvider
       `SSH Task '${task.name}'`,
       'ssh',
       new vscode.CustomExecution(async (resolved: SSHShellTaskOptions) => {
-        const { createTerminal, createTextTerminal, joinCommands } = await import('./pseudoTerminal');
+        const { createTerminal, createTextTerminal } = await import('./pseudoTerminal');
         try {
           if (!resolved.host) throw new Error('Missing field \'host\' in task description');
           if (!resolved.command) throw new Error('Missing field \'command\' in task description');
@@ -196,11 +196,7 @@ export class Manager implements vscode.TaskProvider, vscode.TerminalLinkProvider
           }
           //if (workingDirectory) workingDirectory = this.getRemotePath(config, workingDirectory);
           this.connectionManager.update(connection, con => con.pendingUserCount++);
-          const pty = await createTerminal({
-            command, workingDirectory,
-            client: connection.client,
-            config: connection.actualConfig,
-          });
+          const pty = await createTerminal({ command, workingDirectory, connection });
           this.connectionManager.update(connection, con => (con.pendingUserCount--, con.terminals.push(pty)));
           pty.onDidClose(() => this.connectionManager.update(connection,
             con => con.terminals = con.terminals.filter(t => t !== pty)));

src/pseudoTerminal.ts

@@ -1,9 +1,10 @@
 import * as path from 'path';
-import type { Client, ClientChannel, PseudoTtyOptions } from "ssh2";
+import type { ClientChannel, PseudoTtyOptions } from "ssh2";
 import type { Readable } from "stream";
 import * as vscode from "vscode";
 import { getFlagBoolean } from './config';
-import type { FileSystemConfig } from "./fileSystemConfig";
+import { Connection, environmentToExportString, joinCommands, mergeEnvironment } from './connection';
+import type { EnvironmentVariable, FileSystemConfig } from "./fileSystemConfig";
 import { getRemotePath } from './fileSystemRouter';
 import { Logging, LOGGING_NO_STACKTRACE } from "./logging";
 import { toPromise } from "./toPromise";
@@ -18,8 +19,7 @@ export interface SSHPseudoTerminal extends vscode.Pseudoterminal {
     onDidOpen: vscode.Event<void>;
     handleInput(data: string): void; // We don't support/need read-only terminals for now
     status: 'opening' | 'open' | 'closed';
-    config: FileSystemConfig;
-    client: Client;
+    connection: Connection;
     /** Could be undefined if it only gets created during psy.open() instead of beforehand */
     channel?: ClientChannel;
     /** Either set by the code calling createTerminal, otherwise "calculated" and hopefully found */
@@ -28,25 +28,18 @@ export interface SSHPseudoTerminal extends vscode.Pseudoterminal {
 
 export function isSSHPseudoTerminal(terminal: vscode.Pseudoterminal): terminal is SSHPseudoTerminal {
     const term = terminal as SSHPseudoTerminal;
-    return !!(term.config && term.status && term.client);
+    return !!(term.connection && term.status && term.handleInput);
 }
 
 export interface TerminalOptions {
-    client: Client;
-    config: FileSystemConfig;
+    connection: Connection;
+    environment?: EnvironmentVariable[];
     /** If absent, this defaults to config.root if present, otherwise whatever the remote shell picks as default */
     workingDirectory?: string;
     /** The command to run in the remote shell. If undefined, a (regular interactive) shell is started instead by running $SHELL*/
     command?: string;
 }
 
-export function joinCommands(commands: string | string[] | undefined, separator: string): string | undefined {
-    if (!commands) return undefined;
-    if (typeof commands === 'string') return commands;
-    return commands.join(separator);
-}
-
-
 export function replaceVariables(value: string, config: FileSystemConfig): string {
     return value.replace(/\$\{(.*?)\}/g, (str, match: string) => {
         if (!match.startsWith('remote')) return str; // Our variables always start with "remote"
@@ -145,7 +138,8 @@ export async function replaceVariablesRecursive<T>(object: T, handler: (value: s
 }
 
 export async function createTerminal(options: TerminalOptions): Promise<SSHPseudoTerminal> {
-    const { client, config } = options;
+    const { connection } = options;
+    const { actualConfig, client } = connection;
     const onDidWrite = new vscode.EventEmitter<string>();
     const onDidClose = new vscode.EventEmitter<number>();
     const onDidOpen = new vscode.EventEmitter<void>();
@@ -153,7 +147,7 @@ export async function createTerminal(options: TerminalOptions): Promise<SSHPseud
     // Won't actually open the remote terminal until pseudo.open(dims) is called
     const pseudo: SSHPseudoTerminal = {
         status: 'opening',
-        config, client,
+        connection,
         onDidWrite: onDidWrite.event,
         onDidClose: onDidClose.event,
         onDidOpen: onDidOpen.event,
@@ -168,20 +162,24 @@ export async function createTerminal(options: TerminalOptions): Promise<SSHPseud
             pseudo.channel = undefined;
         },
         async open(dims) {
-            onDidWrite.fire(`Connecting to ${config.label || config.name}...\r\n`);
+            onDidWrite.fire(`Connecting to ${actualConfig.label || actualConfig.name}...\r\n`);
             try {
-                const [useWinCmdSep] = getFlagBoolean('WINDOWS_COMMAND_SEPARATOR', false, config.flags);
+                const [useWinCmdSep] = getFlagBoolean('WINDOWS_COMMAND_SEPARATOR', false, actualConfig.flags);
                 const separator = useWinCmdSep ? ' && ' : '; ';
                 let commands: string[] = [];
+                // Add exports for environment variables if needed
+                const env = mergeEnvironment(connection.environment, options.environment);
+                commands.push(environmentToExportString(env));
+                // Push the actual command or (default) shell command with replaced variables
                 if (options.command) {
-                    commands.push(options.command);
+                    commands.push(replaceVariables(options.command, actualConfig));
                 } else {
-                    const tc = joinCommands(config.terminalCommand, separator);
-                    commands.push(tc ? replaceVariables(tc, config) : '$SHELL');
+                    const tc = joinCommands(actualConfig.terminalCommand, separator);
+                    commands.push(tc ? replaceVariables(tc, actualConfig) : '$SHELL');
                 }
                 // There isn't a proper way of setting the working directory, but this should work in most cases
                 let { workingDirectory } = options;
-                workingDirectory = workingDirectory || config.root;
+                workingDirectory = workingDirectory || actualConfig.root;
                 if (workingDirectory) {
                     if (workingDirectory.startsWith('~')) {
                         // So `cd "~/a/b/..." apparently doesn't work, but `~/"a/b/..."` does

webview/src/types/fileSystemConfig.ts

@@ -8,6 +8,12 @@ export interface ProxyConfig {
 
 export type ConfigLocation = number | string;
 
+/** Might support conditional stuff later, although ssh2/OpenSSH might not support that natively */
+export interface EnvironmentVariable {
+    key: string;
+    value: string;
+}
+
 export function formatConfigLocation(location?: ConfigLocation): string {
   if (!location) return 'Unknown location';
   if (typeof location === 'number') {
@@ -100,6 +106,8 @@ export interface FileSystemConfig extends ConnectConfig {
   terminalCommand?: string | string[];
   /** The command(s) to run when a `ssh-shell` task gets run. Defaults to the placeholder `$COMMAND`. Internally the command `cd ...` is run first */
   taskCommand?: string | string[];
+  /** An object with environment variables to add to the SSH connection. Affects the whole connection thus all terminals */
+  environment?: EnvironmentVariable[] | Record<string, string>;
   /** The filemode to assign to created files */
   newFileMode?: number | string;
   /** Whether this config was created from an instant connection string. Enables fuzzy matching for e.g. PuTTY, config-by-host, ... */