Skip to content

Commit

Permalink
Refactor REMOTE_COMMANDS to support most shells (#267)
Browse files Browse the repository at this point in the history
  • Loading branch information
@SchoofsKelvin
SchoofsKelvin committed Oct 24, 2021
1 parent a57886a commit 06a397e
Show file tree
Hide file tree
Showing 3 changed files with 68 additions and 59 deletions.
40 changes: 7 additions & 33 deletions src/connection.ts
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import { posix as path } from 'path';
import * as readline from 'readline';
import type { Client, ClientChannel, SFTPWrapper } from 'ssh2';
import type { Client, ClientChannel } from 'ssh2';
import * as vscode from 'vscode';
import { configMatches, getFlagBoolean, loadConfigs } from './config';
import type { EnvironmentVariable, FileSystemConfig } from './fileSystemConfig';
Expand All @@ -19,32 +19,11 @@ export interface Connection {
environment: EnvironmentVariable[];
terminals: SSHPseudoTerminal[];
filesystems: SSHFileSystem[];
cache: Record<string, any>;
pendingUserCount: number;
idleTimer: NodeJS.Timeout;
}

const TMP_PROFILE_SCRIPT = `
if type code > /dev/null 2> /dev/null; then
return 0;
fi
code() {
if [ "$#" -ne 1 ] || [ $1 = "help" ] || [ $1 = "--help" ] || [ $1 = "-h" ] || [ $1 = "-?" ]; then
echo "Usage:";
echo " code <path_to_existing_file> Will make VS Code open the file";
echo " code <path_to_existing_folder> Will make VS Code add the folder as an additional workspace folder";
echo " code <path_to_nonexisting_file> Will prompt VS Code to create an empty file, then open it afterwards";
elif [ ! -n "$KELVIN_SSHFS_CMD_PATH" ]; then
echo "Not running in a terminal spawned by SSH FS? Failed to sent!"
elif [ -c "$KELVIN_SSHFS_CMD_PATH" ]; then
echo "::sshfs:code:$(pwd):::$1" >> $KELVIN_SSHFS_CMD_PATH;
echo "Command sent to SSH FS extension";
else
echo "Missing command shell pty of SSH FS extension? Failed to sent!"
fi
}
echo "Injected 'code' alias";
`;

export class ConnectionManager {
protected onConnectionAddedEmitter = new vscode.EventEmitter<Connection>();
protected onConnectionRemovedEmitter = new vscode.EventEmitter<Connection>();
Expand All @@ -70,10 +49,11 @@ export class ConnectionManager {
public getPendingConnections(): [string, FileSystemConfig | undefined][] {
return Object.keys(this.pendingConnections).map(name => [name, this.pendingConnections[name][1]]);
}
protected async _createCommandTerminal(client: Client, authority: string, debugLogging: boolean): Promise<string> {
protected async _createCommandTerminal(client: Client, shellConfig: ShellConfig, authority: string, debugLogging: boolean): Promise<string> {
const logging = Logging.scope(`CmdTerm(${authority})`);
const shell = await toPromise<ClientChannel>(cb => client.shell({}, cb));
shell.write('echo ::sshfs:TTY:`tty`\n');
logging.debug(`TTY COMMAND: ${`echo ${shellConfig.embedSubstitutions`::sshfs:${'echo TTY'}:${'tty'}`}\n`}`);
shell.write(`echo ${shellConfig.embedSubstitutions`::sshfs:${'echo TTY'}:${'tty'}`}\n`);
return new Promise((resolvePath, rejectPath) => {
setTimeout(() => rejectPath(new Error('Timeout fetching command path')), 10e3);
const rl = readline.createInterface(shell.stdout);
Expand Down Expand Up @@ -173,15 +153,8 @@ export class ConnectionManager {
const [flagRCDV, flagRCDR] = getFlagBoolean('DEBUG_REMOTE_COMMANDS', false, actualConfig.flags);
const withDebugStr = flagRCDV ? ` with debug logging enabled by '${flagRCDR}'` : '';
logging.info`Flag REMOTE_COMMANDS provided in '${flagRCR}', setting up command terminal${withDebugStr}`;
const cmdPath = await this._createCommandTerminal(client, name, flagRCDV);
const cmdPath = await this._createCommandTerminal(client, shellConfig, name, flagRCDV);
environment.push({ key: 'KELVIN_SSHFS_CMD_PATH', value: cmdPath });
const profilePath = `/tmp/.Kelvin_sshfs.${actualConfig.username || Date.now()}`;
environment.push({ key: 'KELVIN_SSHFS_PROFILE_PATH', value: profilePath });
const sftp = await toPromise<SFTPWrapper>(cb => client.sftp(cb));
await toPromise(cb => sftp.writeFile(profilePath, TMP_PROFILE_SCRIPT, { mode: 0o666 }, cb)).catch(e => {
logging.error`Failed to write profile script to '${profilePath}':\n${e}\nDisabling REMOTE_COMMANDS flag`;
actualConfig.flags = ['-REMOTE_COMMANDS', ...(actualConfig.flags || [])];
});
}
logging.debug`Environment: ${environment}`;
// Set up the Connection object
Expand All @@ -190,6 +163,7 @@ export class ConnectionManager {
config, client, actualConfig, home, shellConfig, environment,
terminals: [],
filesystems: [],
cache: {},
pendingUserCount: 0,
idleTimer: setInterval(() => { // Automatically close connection when idle for a while
timeoutCounter = timeoutCounter ? timeoutCounter - 1 : 0;
Expand Down
5 changes: 2 additions & 3 deletions src/pseudoTerminal.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -179,9 +179,8 @@ export async function createTerminal(options: TerminalOptions): Promise<SSHPseud
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`);
commands.push(shellConfig.setupRemoteCommands(profilePathEnv.value));
const rcCmds = await shellConfig.setupRemoteCommands(connection);
if (rcCmds?.length) commands.push(joinCommands(rcCmds, separator)!);
}
// Push the actual command or (default) shell command with replaced variables
if (options.command) {
Expand Down
82 changes: 59 additions & 23 deletions src/shellConfig.ts
View file
Original file line number Diff line number Diff line change
@@ -1,50 +1,86 @@
import { posix as path } from 'path';
import type { Client, ClientChannel } from "ssh2";
import type { Logger } from "./logging";
import type { Client, ClientChannel, SFTPWrapper } from "ssh2";
import type { Connection } from './connection';
import { Logger, Logging } from "./logging";
import { toPromise } from "./utils";

const SCRIPT_COMMAND_CODE = `#!/bin/sh
if [ "$#" -ne 1 ] || [ $1 = "help" ] || [ $1 = "--help" ] || [ $1 = "-h" ] || [ $1 = "-?" ]; then
echo "Usage:";
echo " code <path_to_existing_file> Will make VS Code open the file";
echo " code <path_to_existing_folder> Will make VS Code add the folder as an additional workspace folder";
echo " code <path_to_nonexisting_file> Will prompt VS Code to create an empty file, then open it afterwards";
elif [ ! -n "$KELVIN_SSHFS_CMD_PATH" ]; then
echo "Not running in a terminal spawned by SSH FS? Failed to sent!"
elif [ -c "$KELVIN_SSHFS_CMD_PATH" ]; then
echo "::sshfs:code:$(pwd):::$1" >> $KELVIN_SSHFS_CMD_PATH;
echo "Command sent to SSH FS extension";
else
echo "Missing command shell pty of SSH FS extension? Failed to sent!"
fi
`;

type RemoteCommandInitializer = (connection: Connection) => void
| string | string[] | undefined
| Promise<void | string | string[] | undefined>;

async function ensureCachedFile(connection: Connection, key: string, path: string, content: string, sftp?: SFTPWrapper):
Promise<[written: boolean, path: string | null]> {
const rc_files: Record<string, string> = connection.cache.rc_files ||= {};
if (rc_files[key]) return [false, rc_files[key]];
try {
sftp ||= await toPromise<SFTPWrapper>(cb => connection.client.sftp(cb));
await toPromise(cb => sftp!.writeFile(path, content, { mode: 0o755 }, cb));
return [true, rc_files[key] = path];
} catch (e) {
Logging.error`Failed to write ${key} file to '${path}':\n${e}`;
return [false, null];
}
}

async function rcInitializePATH(connection: Connection): Promise<string[] | string> {
const dir = `/tmp/.Kelvin_sshfs.RcBin.${connection.actualConfig.username || Date.now()}`;
const sftp = await toPromise<SFTPWrapper>(cb => connection.client.sftp(cb));
await toPromise(cb => sftp!.mkdir(dir, { mode: 0o755 }, cb)).catch(() => { });
const [, path] = await ensureCachedFile(connection, 'CmdCode', `${dir}/code`, SCRIPT_COMMAND_CODE, sftp);
return path ? [
connection.shellConfig.setEnv('PATH', `${dir}:$PATH`),
] : 'echo "An error occured while adding REMOTE_COMMANDS support"';
}

export interface ShellConfig {
shell: string;
setEnv(key: string, value: string): string;
setupRemoteCommands(path: string): string;
setupRemoteCommands: RemoteCommandInitializer;
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,
setupRemoteCommands: RemoteCommandInitializer,
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);
add('sh', setEnvExport, rcInitializePATH, embedSubstitutionsBackticks);
add('bash', setEnvExport, rcInitializePATH, embedSubstitutionsBackticks);
add('rbash', setEnvExport, rcInitializePATH, embedSubstitutionsBackticks);
add('ash', setEnvExport, rcInitializePATH, embedSubstitutionsBackticks);
add('dash', setEnvExport, rcInitializePATH, embedSubstitutionsBackticks);
add('ksh', setEnvExport, rcInitializePATH, embedSubstitutionsBackticks);
add('zsh', setEnvExport, rcInitializePATH, embedSubstitutionsBackticks);
add('fish', setEnvSetGX, rcInitializePATH, embedSubstitutionsFish); // https://fishshell.com/docs/current/tutorial.html#autoloading-functions
add('csh', setEnvSetEnv, rcInitializePATH, embedSubstitutionsBackticks);
add('tcsh', setEnvSetEnv, rcInitializePATH, embedSubstitutionsBackticks);
}

export async function tryCommand(ssh: Client, command: string): Promise<string | null> {
Expand Down

0 comments on commit 06a397e

Please sign in to comment.