From 3109e977a5ac068aac6cf9081cfcc6c5b655de2f Mon Sep 17 00:00:00 2001 From: Kelvin Schoofs Date: Mon, 28 Jun 2021 23:16:17 +0200 Subject: [PATCH] Add support for environment variables (closes #241) --- src/connection.ts | 43 +++++++++++++++++++++++++-- src/fileSystemConfig.ts | 8 +++++ src/manager.ts | 12 +++----- src/pseudoTerminal.ts | 42 +++++++++++++------------- webview/src/types/fileSystemConfig.ts | 8 +++++ 5 files changed, 80 insertions(+), 33 deletions(-) diff --git a/src/connection.ts b/src/connection.ts index b0dbc2c..f766780 100644 --- a/src/connection.ts +++ b/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 | 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(); protected onConnectionRemovedEmitter = new vscode.EventEmitter(); @@ -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, diff --git a/src/fileSystemConfig.ts b/src/fileSystemConfig.ts index 2bf4977..d11fea2 100644 --- a/src/fileSystemConfig.ts +++ b/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; /** 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, ... */ diff --git a/src/manager.ts b/src/manager.ts index a4b0f49..67a9406 100644 --- a/src/manager.ts +++ b/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))); diff --git a/src/pseudoTerminal.ts b/src/pseudoTerminal.ts index 9e4680b..9ca2aee 100644 --- a/src/pseudoTerminal.ts +++ b/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; 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(object: T, handler: (value: s } export async function createTerminal(options: TerminalOptions): Promise { - const { client, config } = options; + const { connection } = options; + const { actualConfig, client } = connection; const onDidWrite = new vscode.EventEmitter(); const onDidClose = new vscode.EventEmitter(); const onDidOpen = new vscode.EventEmitter(); @@ -153,7 +147,7 @@ export async function createTerminal(options: TerminalOptions): Promise; /** 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, ... */