On shutdown, explicitly try to close all interceptors

Author: pimterry

src/index.ts

@@ -11,6 +11,7 @@ import { HttpToolkitServer } from './httptoolkit-server';
 import { checkBrowserConfig } from './browsers';
 import { reportError } from './error-tracking';
 import { delay, ALLOWED_ORIGINS } from './util';
+import { registerShutdownHandler } from './shutdown';
 
 const canAccess = util.promisify(fs.access);
 const mkDir = util.promisify(fs.mkdir);
@@ -48,6 +49,7 @@ export async function runHTK(options: {
     configPath?: string
 } = {}) {
     const startTime = Date.now();
+    registerShutdownHandler();
 
     const configPath = options.configPath || envPaths('httptoolkit', { suffix: '' }).config;
 

src/interceptors/fresh-chrome.ts

@@ -31,7 +31,7 @@ export class FreshChrome implements Interceptor {
 
     constructor(private config: HtkConfig) { }
 
-    isActive(proxyPort: number) {
+    isActive(proxyPort: number | string) {
         return browsers[proxyPort] != null && !!browsers[proxyPort].pid;
     }
 
@@ -78,12 +78,18 @@ export class FreshChrome implements Interceptor {
         await delay(500);
     }
 
-    async deactivate(proxyPort: number) {
+    async deactivate(proxyPort: number | string) {
         if (this.isActive(proxyPort)) {
             const browser = browsers[proxyPort];
             const exitPromise = new Promise((resolve) => browser!.process.once('exit', resolve));
             browser!.stop();
             await exitPromise;
         }
     }
+
+    async deactivateAll(): Promise<void> {
+        await Promise.all(
+            Object.keys(browsers).map((proxyPort) => this.deactivate(proxyPort))
+        );
+    }
 };

src/interceptors/fresh-firefox.ts

@@ -28,7 +28,7 @@ export class FreshFirefox implements Interceptor {
 
     constructor(private config: HtkConfig) { }
 
-    isActive(proxyPort: number) {
+    isActive(proxyPort: number | string) {
         return browsers[proxyPort] != null && !!browsers[proxyPort].pid;
     }
 
@@ -133,12 +133,18 @@ export class FreshFirefox implements Interceptor {
         await delay(1000);
     }
 
-    async deactivate(proxyPort: number) {
+    async deactivate(proxyPort: number | string) {
         if (this.isActive(proxyPort)) {
             const browser = browsers[proxyPort];
             const exitPromise = new Promise((resolve) => browser!.process.once('exit', resolve));
             browser!.stop();
             await exitPromise;
         }
     }
+
+    async deactivateAll(): Promise<void> {
+        await Promise.all(
+            Object.keys(browsers).map((proxyPort) => this.deactivate(proxyPort))
+        );
+    }
 };

src/interceptors/fresh-terminal.ts

@@ -223,7 +223,7 @@ export class TerminalInterceptor implements Interceptor {
         return !!(await getTerminalCommand());
     }
 
-    isActive(proxyPort: number): boolean {
+    isActive(proxyPort: number | string): boolean {
         return !!(terminals[proxyPort] && terminals[proxyPort]!.length);
     }
 
@@ -288,7 +288,7 @@ export class TerminalInterceptor implements Interceptor {
         });
     }
 
-    async deactivate(proxyPort: number): Promise<void> {
+    async deactivate(proxyPort: number | string): Promise<void> {
         if (!this.isActive(proxyPort)) return;
 
         await Promise.all((terminals[proxyPort] || []).map((proc) => {
@@ -299,4 +299,10 @@ export class TerminalInterceptor implements Interceptor {
         }));
     }
 
+    async deactivateAll(): Promise<void> {
+        await Promise.all(
+            Object.keys(terminals).map((proxyPort) => this.deactivate(proxyPort))
+        );
+    }
+
 }

src/interceptors/index.ts

@@ -5,6 +5,7 @@ import { HtkConfig } from '../config';
 import { FreshChrome } from './fresh-chrome';
 import { FreshFirefox } from './fresh-firefox';
 import { TerminalInterceptor } from './fresh-terminal';
+import { addShutdownHandler } from '../shutdown';
 
 export interface Interceptor {
     id: string;
@@ -15,12 +16,22 @@ export interface Interceptor {
 
     activate(proxyPort: number, options?: any): Promise<void>;
     deactivate(proxyPort: number, options?: any): Promise<void>;
+    deactivateAll(): Promise<void>;
 }
 
 export function buildInterceptors(config: HtkConfig): _.Dictionary<Interceptor> {
-    return _.keyBy([
+    const interceptors = [
         new FreshChrome(config),
         new FreshFirefox(config),
         new TerminalInterceptor(config)
-    ], (interceptor) => interceptor.id);
+    ];
+
+    // When the server exits, try to shut down the interceptors too
+    addShutdownHandler(() => shutdownInterceptors(interceptors));
+
+    return _.keyBy(interceptors, (interceptor) => interceptor.id);
+}
+
+function shutdownInterceptors(interceptors: Interceptor[]) {
+    return Promise.all(interceptors.map(i => i.deactivateAll()));
 }

src/shutdown.ts

@@ -0,0 +1,24 @@
+import { reportError } from './error-tracking';
+
+type ShutdownHandler = () => Promise<unknown>;
+const shutdownHandlers: Array<ShutdownHandler> = [];
+
+export function registerShutdownHandler() {
+    process.on('SIGTERM', shutdown);
+    process.on('SIGINT', shutdown);
+}
+
+export function addShutdownHandler(handler: ShutdownHandler) {
+    shutdownHandlers.push(handler);
+}
+
+async function shutdown() {
+    console.log('Shutting down...');
+
+    await Promise.all(shutdownHandlers.map(
+        (handler) => handler().catch(reportError)
+    ));
+
+    process.exit(0);
+}
+

test/interceptors/interceptor-test-utils.ts

@@ -59,4 +59,17 @@ export function itCanBeActivated(interceptorSetup: InterceptorSetup) {
         await interceptor.deactivate(server.port);
         expect(interceptor.isActive(server.port)).to.equal(false);
     });
+
+    it('can deactivate all', async () => {
+        const { interceptor, server } = await interceptorSetup;
+
+        expect(interceptor.isActive(server.port)).to.equal(false);
+
+        await interceptor.activate(server.port);
+        expect(interceptor.isActive(server.port)).to.equal(true);
+        expect(interceptor.isActive(server.port + 1)).to.equal(false);
+
+        await interceptor.deactivateAll();
+        expect(interceptor.isActive(server.port)).to.equal(false);
+    });
 }