Better handle connection being stalled/closed (issue #9)

src/manager.ts

@@ -276,11 +276,11 @@ export class Manager implements vscode.FileSystemProvider, vscode.TreeDataProvid
           delete this.creatingFileSystems[name];
           vscode.commands.executeCommand('workbench.files.action.refreshFilesExplorer');
           this.onDidChangeTreeDataEmitter.fire();
+          client.on('close', hadError => hadError ? this.commandReconnect(name) : (!fs.closing && this.promptReconnect(name)));
           return resolve(fs);
         });
       });
       client.on('timeout', () => reject(new Error(`Socket timed out while connecting SSH FS '${name}'`)));
-      client.on('close', hadError => hadError && this.commandReconnect(name));
       client.on('error', (error) => {
         if (error.description) {
           error.message = `${error.description}\n${error.message}`;
@@ -321,6 +321,17 @@ export class Manager implements vscode.FileSystemProvider, vscode.TreeDataProvid
     if (fs) return fs;
     return null;
   }
+  public async promptReconnect(name: string) {
+    const config = this.getConfig(name);
+    console.log('config', name, config);
+    if (!config) return;
+    const choice = await vscode.window.showWarningMessage(`SSH FS ${config.label || config.name} disconnected`, 'Reconnect', 'Disconnect');
+    if (choice === 'Reconnect') {
+      this.commandReconnect(name);
+    } else {
+      this.commandDisconnect(name);
+    }
+  }
   /* FileSystemProvider */
   public watch(uri: vscode.Uri, options: { recursive: boolean; excludes: string[]; }): vscode.Disposable {
     /*let disp = () => {};

src/sshFileSystem.ts

@@ -4,9 +4,11 @@ import * as ssh2 from 'ssh2';
 import * as ssh2s from 'ssh2-streams';
 import * as vscode from 'vscode';
 import { FileSystemConfig } from './manager';
-import { toPromise } from './toPromise';
 
 export class SSHFileSystem implements vscode.FileSystemProvider {
+  public waitForContinue = false;
+  public closed = false;
+  public closing = false;
   public copy = undefined;
   public onDidChangeFile: vscode.Event<vscode.FileChangeEvent[]>;
   protected onDidChangeFileEmitter = new vscode.EventEmitter<vscode.FileChangeEvent[]>();
@@ -14,9 +16,11 @@ export class SSHFileSystem implements vscode.FileSystemProvider {
   constructor(public readonly authority: string, protected sftp: ssh2.SFTPWrapper,
               public readonly root: string, public readonly config: FileSystemConfig) {
     this.onDidChangeFile = this.onDidChangeFileEmitter.event;
+    this.sftp.on('end', () => this.closed = true);
   }
 
   public disconnect() {
+    this.closing = true;
     this.sftp.end();
   }
 
@@ -25,12 +29,32 @@ export class SSHFileSystem implements vscode.FileSystemProvider {
     return path.posix.resolve(this.root, relPath);
   }
 
+  public continuePromise<T>(func: (cb: (err: Error | null, res?: T) => void) => boolean): Promise<T> {
+    return new Promise<T>((resolve, reject) => {
+      const exec = () => {
+        this.waitForContinue = false;
+        if (this.closed) return reject(new Error('Connection closed'));
+        try {
+          const canContinue = func((err, res) => err ? reject(err) : resolve(res));
+          if (!canContinue) this.waitForContinue = true;
+        } catch (e) {
+          reject(e);
+        }
+      };
+      if (this.waitForContinue) {
+        this.sftp.once('continue', exec);
+      } else {
+        exec();
+      }
+    });
+  }
+
   public watch(uri: vscode.Uri, options: { recursive: boolean; excludes: string[]; }): vscode.Disposable {
     // throw new Error('Method not implemented.');
     return new vscode.Disposable(() => { });
   }
   public async stat(uri: vscode.Uri): Promise<vscode.FileStat> {
-    const stat = await toPromise<ssh2s.Stats>(cb => this.sftp.stat(this.relative(uri.path), cb)).catch((e: Error & { code: number }) => {
+    const stat = await this.continuePromise<ssh2s.Stats>(cb => this.sftp.stat(this.relative(uri.path), cb)).catch((e: Error & { code: number }) => {
       throw e.code === 2 ? vscode.FileSystemError.FileNotFound(uri) : e;
     });
     const { mtime, size } = stat;
@@ -46,7 +70,7 @@ export class SSHFileSystem implements vscode.FileSystemProvider {
     };
   }
   public async readDirectory(uri: vscode.Uri): Promise<[string, vscode.FileType][]> {
-    const entries = await toPromise<ssh2s.FileEntry[]>(cb => this.sftp.readdir(this.relative(uri.path), cb)).catch((e) => {
+    const entries = await this.continuePromise<ssh2s.FileEntry[]>(cb => this.sftp.readdir(this.relative(uri.path), cb)).catch((e) => {
       throw e === 2 ? vscode.FileSystemError.FileNotFound(uri) : e;
     });
     return Promise.all(entries.map(async (file) => {
@@ -56,7 +80,7 @@ export class SSHFileSystem implements vscode.FileSystemProvider {
     }));
   }
   public createDirectory(uri: vscode.Uri): void | Promise<void> {
-    return toPromise(cb => this.sftp.mkdir(this.relative(uri.path), cb));
+    return this.continuePromise(cb => this.sftp.mkdir(this.relative(uri.path), cb));
   }
   public readFile(uri: vscode.Uri): Uint8Array | Promise<Uint8Array> {
     return new Promise((resolve, reject) => {
@@ -74,7 +98,7 @@ export class SSHFileSystem implements vscode.FileSystemProvider {
     return new Promise(async (resolve, reject) => {
       let mode: number | undefined;
       try {
-        const stat = await toPromise<ssh2s.Stats>(cb => this.sftp.stat(this.relative(uri.path), cb));
+        const stat = await this.continuePromise<ssh2s.Stats>(cb => this.sftp.stat(this.relative(uri.path), cb));
         mode = stat.mode;
       } catch (e) {
         if (e.message !== 'No such file') {
@@ -92,15 +116,15 @@ export class SSHFileSystem implements vscode.FileSystemProvider {
     const stats = await this.stat(uri);
     // tslint:disable no-bitwise */
     if (stats.type & (vscode.FileType.SymbolicLink | vscode.FileType.File)) {
-      return toPromise(cb => this.sftp.unlink(this.relative(uri.path), cb));
+      return this.continuePromise(cb => this.sftp.unlink(this.relative(uri.path), cb));
     } else if ((stats.type & vscode.FileType.Directory) && options.recursive) {
-      return toPromise(cb => this.sftp.rmdir(this.relative(uri.path), cb));
+      return this.continuePromise(cb => this.sftp.rmdir(this.relative(uri.path), cb));
     }
-    return toPromise(cb => this.sftp.unlink(this.relative(uri.path), cb));
+    return this.continuePromise(cb => this.sftp.unlink(this.relative(uri.path), cb));
     // tslint:enable no-bitwise */
   }
   public rename(oldUri: vscode.Uri, newUri: vscode.Uri, options: { overwrite: boolean; }): void | Promise<void> {
-    return toPromise(cb => this.sftp.rename(this.relative(oldUri.path), this.relative(newUri.path), cb));
+    return this.continuePromise(cb => this.sftp.rename(this.relative(oldUri.path), this.relative(newUri.path), cb));
   }
 }
 

src/toPromise.ts

@@ -2,7 +2,11 @@
 export type toPromiseCallback<T> = (err: Error | null, res?: T) => void;
 export async function toPromise<T>(func: (cb: toPromiseCallback<T>) => void): Promise<T> {
   return new Promise<T>((resolve, reject) => {
+    try {
       func((err, res) => err ? reject(err) : resolve(res));
+    } catch (e) {
+      reject(e);
+    }
   });
 }