From 5721f1ca0935217c02e036af64089527be49c7e9 Mon Sep 17 00:00:00 2001
From: Kelvin Schoofs <schoofs.kelvin@gmail.com>
Date: Tue, 7 Jun 2022 21:58:14 +0200
Subject: [PATCH] Add `SHELL_CONFIG` flag and fix too eager PowerShell version
 detection (#331)

---
 CHANGELOG.md       |  1 +
 src/config.ts      |  3 +++
 src/connection.ts  | 14 +++++++++++---
 src/shellConfig.ts |  4 ++--
 4 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 0064b9e..dcaf03c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -22,6 +22,7 @@
   - The `write` flag will show a notification should an error happen for a "write" operation
   - Write operations are: `createDirectory`, `writeFile`, `delete`, and `rename`
   - Since `readDirectory`, `readFile` and `stat` are disabled by default, it should prevent extension detection spam (see #341)
+- Added the `SHELL_CONFIG` flag to force a specific remote shell configuration (#331)
 
 ## v1.25.0 (2022-06-01)
 
diff --git a/src/config.ts b/src/config.ts
index e65f822..d952050 100644
--- a/src/config.ts
+++ b/src/config.ts
@@ -405,6 +405,9 @@ function parseFlagList(list: string[] | undefined, origin: string): Record<strin
     - The presence of `write` is equal to `createDirectory,writeFile,delete,rename`
     - Besides those provided by  `write`, there's also `readDirectory`, `readFile` and `stat`
     - Automatically set to `write` for VS Code 1.56 and later (see issue #282)
+  SHELL_CONFIG (string)
+    - Forces the use of a specific shell configuration. Check shellConfig.ts for possible values
+    - By default, when this flag is absent (or an empty or not a string), the extension will try to detect the correct type to use
 */
 export type FlagValue = string | boolean | null;
 export type FlagCombo<V extends FlagValue = FlagValue> = [value: V, origin: string];
diff --git a/src/connection.ts b/src/connection.ts
index f68ccc5..73f3d3b 100644
--- a/src/connection.ts
+++ b/src/connection.ts
@@ -3,10 +3,10 @@ import { posix as path } from 'path';
 import * as readline from 'readline';
 import type { Client, ClientChannel } from 'ssh2';
 import * as vscode from 'vscode';
-import { configMatches, getFlagBoolean, loadConfigs } from './config';
+import { configMatches, getFlag, getFlagBoolean, loadConfigs } from './config';
 import { Logging, LOGGING_NO_STACKTRACE } from './logging';
 import type { SSHPseudoTerminal } from './pseudoTerminal';
-import { calculateShellConfig, ShellConfig, tryCommand, tryEcho } from './shellConfig';
+import { calculateShellConfig, KNOWN_SHELL_CONFIGS, ShellConfig, tryCommand, tryEcho } from './shellConfig';
 import type { SSHFileSystem } from './sshFileSystem';
 import { mergeEnvironment, toPromise } from './utils';
 
@@ -127,7 +127,15 @@ export class ConnectionManager {
         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);
+        let shellConfig: ShellConfig;
+        const [flagSCV, flagSCR] = getFlag("SHELL_CONFIG", config.flags) || [];
+        if (flagSCV && typeof flagSCV === 'string') {
+            logging.info`Using forced shell config '${flagSCV}' set by ${flagSCR}`;
+            shellConfig = KNOWN_SHELL_CONFIGS[flagSCV];
+            if (!shellConfig) throw new Error(`The forced shell config '${flagSCV}' does not exist`);
+        } else {
+            shellConfig = await calculateShellConfig(client, logging);
+        }
         // Query home directory
         let home: string | Error | null;
         if (shellConfig.isWindows) {
diff --git a/src/shellConfig.ts b/src/shellConfig.ts
index 8315522..be0af75 100644
--- a/src/shellConfig.ts
+++ b/src/shellConfig.ts
@@ -55,7 +55,7 @@ export interface ShellConfig {
     setupRemoteCommands?: RemoteCommandInitializer;
     embedSubstitutions?(command: TemplateStringsArray, ...substitutions: (string | number)[]): string;
 }
-const KNOWN_SHELL_CONFIGS: Record<string, ShellConfig> = {}; {
+export const KNOWN_SHELL_CONFIGS: Record<string, ShellConfig> = {}; {
     const add = (shell: string,
         setEnv: (key: string, value: string) => string,
         setupRemoteCommands?: RemoteCommandInitializer,
@@ -122,7 +122,7 @@ async function getPowershellVersion(client: Client): Promise<string | null> {
         console.error(e);
         return null;
     });
-    return !version?.includes('PSVersion') ? version : null;
+    return version?.match(/\d+\.\d+\.\d+\.\d+/)?.[0] || null;
 }
 
 async function getWindowsVersion(client: Client): Promise<string | null> {