From 196a276f52450cc435c9506c55f3e21c7a05b4ad Mon Sep 17 00:00:00 2001
From: Kelvin Schoofs <schoofs.kelvin@gmail.com>
Date: Sun, 13 May 2018 16:09:37 +0200
Subject: [PATCH] Add support for SOCKS4/5 proxy

---
 src/manager.ts | 22 +++++++++++++++++++++-
 src/proxy.ts   | 36 ++++++++++++++++++++++++++++++++++++
 src/putty.ts   |  6 ++++++
 3 files changed, 63 insertions(+), 1 deletion(-)
 create mode 100644 src/proxy.ts

diff --git a/src/manager.ts b/src/manager.ts
index 752ac4f..beff0f2 100644
--- a/src/manager.ts
+++ b/src/manager.ts
@@ -4,6 +4,8 @@ import { parse as parseJsonc, ParseError } from 'jsonc-parser';
 import * as path from 'path';
 import { Client, ConnectConfig } from 'ssh2';
 import * as vscode from 'vscode';
+
+import * as proxy from './proxy';
 import { getSession as getPuttySession } from './putty';
 import SSHFileSystem, { EMPTY_FILE_SYSTEM } from './sshFileSystem';
 import { catchingPromise, toPromise } from './toPromise';
@@ -15,11 +17,18 @@ async function assertFs(man: Manager, uri: vscode.Uri) {
   // throw new Error(`A SSH filesystem with the name '${uri.authority}' doesn't exists`);
 }
 
+export interface ProxyConfig {
+  type: 'socks4' | 'socks5';
+  host: string;
+  port: number;
+}
+
 export interface FileSystemConfig extends ConnectConfig {
   name: string;
   label?: string;
   root?: string;
   putty?: string | boolean;
+  proxy?: ProxyConfig;
 }
 
 export enum ConfigStatus {
@@ -192,7 +201,7 @@ export class Manager implements vscode.FileSystemProvider, vscode.TreeDataProvid
             return reject(new Error(`Error while reading the keyfile at:\n${session.publickeyfile}`));
           }
         }
-      }
+        }
       if (!config.username || (config.username as any) === true) {
         config.username = await vscode.window.showInputBox({
           ignoreFocusOut: true,
@@ -226,6 +235,17 @@ export class Manager implements vscode.FileSystemProvider, vscode.TreeDataProvid
         }
       }
       if (config.password) config.agent = undefined;
+      switch (config.proxy && config.proxy.type) {
+        case null:
+        case undefined:
+          break;
+        case 'socks4':
+        case 'socks5':
+          config = await proxy.socks(config);
+          break;
+        default:
+          return reject(new Error(`Unknown proxy method`));
+      }
       const client = new Client();
       client.on('ready', () => {
         client.sftp((err, sftp) => {
diff --git a/src/proxy.ts b/src/proxy.ts
new file mode 100644
index 0000000..1ff7427
--- /dev/null
+++ b/src/proxy.ts
@@ -0,0 +1,36 @@
+
+import * as dns from 'dns';
+import { SocksClient } from 'socks';
+
+import { FileSystemConfig } from './manager';
+import { toPromise } from './toPromise';
+
+export async function socks(config: FileSystemConfig): Promise<FileSystemConfig> {
+  if (!config.proxy) throw new Error(`Missing field 'config.proxy'`);
+  if (!config.proxy.host) throw new Error(`Missing field 'config.proxy.host'`);
+  if (!config.proxy.port) throw new Error(`Missing field 'config.proxy.port'`);
+  if (!config.proxy.type) throw new Error(`Missing field 'config.proxy.type'`);
+  if (config.proxy.type !== 'socks4' && config.proxy.type !== 'socks5') {
+    throw new Error(`Expected config.proxy.type' to be 'socks4 or 'socks5'`);
+  }
+  try {
+    const ipaddress = (await toPromise<string[]>(cb => dns.resolve(config.proxy!.host, cb)))[0];
+    if (!ipaddress) throw new Error(`Couldn't resolve '${config.proxy.host}'`);
+    const con = await SocksClient.createConnection({
+      command: 'connect',
+      destination: {
+        host: config.host!,
+        port: config.port!,
+      },
+      proxy: {
+        ipaddress,
+        port: config.proxy.port,
+        type: config.proxy.type === 'socks4' ? 4 : 5,
+      },
+    });
+    config.sock = con.socket as NodeJS.ReadableStream;
+    return config;
+  } catch (e) {
+    throw new Error(`Error while connecting to the the proxy: ${e.message}`);
+  }
+}
diff --git a/src/putty.ts b/src/putty.ts
index 5deab03..5af47dd 100644
--- a/src/putty.ts
+++ b/src/putty.ts
@@ -11,6 +11,7 @@ const winreg = new Winreg({
 export type NumberAsBoolean = 0 | 1;
 export interface PuttySession {
   [key: string]: string | number | undefined;
+  // General settings
   name: string;
   hostname: string;
   protocol: string;
@@ -19,6 +20,11 @@ export interface PuttySession {
   usernamefromenvironment: NumberAsBoolean;
   tryagent: NumberAsBoolean;
   publickeyfile?: string;
+  // Proxy settings
+  proxyhost?: string;
+  proxyport: number;
+  proxylocalhost: NumberAsBoolean;
+  proxymethod: number; // Key of ['None', 'SOCKS 4', 'SOCKS 5', 'HTTP', 'Telnet', 'Local'] // Only checked first 3
 }
 
 function valueFromItem(item: Winreg.RegistryItem) {