FIX : Resolve reverse tunnel and WiFi proxy leaks on session end or crash

src/plugin.ts

@@ -18,6 +18,7 @@ import {
   isRealDevice,
   getAdbReverseTunnels,
   getCurrentWifiProxyConfig,
+  removeReverseTunnel,
   ADBInstance,
   UDID,
 } from './utils/adb';
@@ -89,6 +90,36 @@ export class AppiumInterceptorPlugin extends BasePlugin {
     super(name, cliArgs);
     log.debug(`📱 Initializing plugin with CLI args: ${JSON.stringify(cliArgs)}`);
     this.pluginArgs = Object.assign({}, DefaultPluginArgs, cliArgs as unknown as IPluginArgs);
+    this.registerProcessExitHandlers();
+  }
+
+  private registerProcessExitHandlers() {
+    let isCleaningUp = false;
+
+    const cleanupAllProxies = async (signal: string) => {
+      if (isCleaningUp) return;
+      isCleaningUp = true;
+
+      const sessionIds = proxyCache.getAllSessionIds();
+      if (sessionIds.length > 0) {
+        log.info(`[Cleanup] Process received ${signal}. Cleaning up ${sessionIds.length} active proxy sessions...`);
+        for (const sessionId of sessionIds) {
+          try {
+            await this.clearProxy(undefined, sessionId);
+          } catch (err: any) {
+            log.error(`[Cleanup] Error during process exit cleanup for session ${sessionId}: ${err.message}`);
+          }
+        }
+      }
+
+      if (signal === 'SIGINT' || signal === 'SIGTERM') {
+        // Send the signal to ourselves again so default or other handlers can run
+        process.kill(process.pid, signal);
+      }
+    };
+
+    process.once('SIGINT', () => cleanupAllProxies('SIGINT'));
+    process.once('SIGTERM', () => cleanupAllProxies('SIGTERM'));
   }
 
   /**
@@ -163,6 +194,12 @@ export class AppiumInterceptorPlugin extends BasePlugin {
       const adb = driver.sessions[sessionId]?.adb;
       await this.clearProxy(adb, sessionId);
     }
+
+    const remainingSessions = proxyCache.getAllSessionIds();
+    for (const sessionId of remainingSessions) {
+      log.warn(`[${sessionId}] Session still in proxyCache after unexpected shutdown. Forcing cleanup...`);
+      await this.clearProxy(undefined, sessionId);
+    }
   }
 
   async addMock(_next: any, driver: any, config: MockConfig) {
@@ -286,6 +323,7 @@ export class AppiumInterceptorPlugin extends BasePlugin {
           : parseJson(this.pluginArgs.blacklisteddomains),
       );
       const proxy = await setupProxyServer(
+        adb,
         sessionId,
         deviceUDID,
         realDevice,
@@ -307,18 +345,38 @@ export class AppiumInterceptorPlugin extends BasePlugin {
     }
   }
 
-  private async clearProxy(adb: ADBInstance, sessionId: string) {
+  private async clearProxy(adb: ADBInstance | undefined, sessionId: string) {
     const proxy = proxyCache.get(sessionId);
     if (!proxy) {
       log.debug(`[${sessionId}] No proxy registered for this session. Nothing to clear.`);
       return;
     }
 
+    const activeAdb = adb || proxy.options.adb;
+    if (!activeAdb) {
+      log.warn(`[${sessionId}] ADB instance is missing. Cannot revert proxy settings or remove reverse tunnels.`);
+    }
+
     log.debug(`[${sessionId}] Reverting device settings and cleaning up proxy resources...`);
 
     try {
-      // Revert WiFi settings to previous state or off
-      await configureWifiProxy(adb, proxy.options.deviceUDID, false, proxy.previousGlobalProxy);
+      const isReal = proxy.options.isRealDevice ?? false;
+
+      if (activeAdb) {
+        // Revert WiFi settings to previous state or off
+        await configureWifiProxy(activeAdb, proxy.options.deviceUDID, isReal, proxy.previousGlobalProxy);
+
+        // Explicitly remove the adb reverse tunnel if this is a real device
+        if (isReal) {
+          log.debug(`[${sessionId}] Removing reverse tunnel for port ${proxy.port}...`);
+          try {
+            await removeReverseTunnel(activeAdb, proxy.options.deviceUDID, proxy.port);
+          } catch (tunnelErr: any) {
+            log.warn(`[${sessionId}] Failed to remove reverse tunnel: ${tunnelErr.message}`);
+          }
+        }
+      }
+
       // Shutdown the local proxy server
       await cleanUpProxyServer(proxy);
       proxyCache.remove(sessionId);

src/proxy-cache.ts

@@ -14,6 +14,14 @@ class ProxyCache {
   get(sessionId: string) {
     return this.cache.get(sessionId);
   }
+
+  getAllSessionIds(): string[] {
+    return Array.from(this.cache.keys());
+  }
+
+  clear() {
+    this.cache.clear();
+  }
 }
 
 export default new ProxyCache();

src/proxy.ts

@@ -3,6 +3,7 @@ import { Proxy as HttpProxy, IContext, IProxyOptions } from 'http-mitm-proxy';
 import * as net from 'net';
 import { ProxyAgent } from 'proxy-agent';
 import { v4 as uuid } from 'uuid';
+import ADB from 'appium-adb';
 import {
   addDefaultMocks,
   compileMockConfig,
@@ -30,6 +31,8 @@ export interface ProxyOptions {
   certificatePath: string;
   port: number;
   ip: string;
+  adb?: ADB;
+  isRealDevice?: boolean;
   previousConfig?: ProxyOptions;
   whitelistedDomains?: string[];
   blacklistedDomains?: string[];

src/scripts/test-connection.ts

@@ -81,7 +81,7 @@ async function addMock(proxy: Proxy) {
 
 async function verifyDeviceConnection(adb: ADBInstance, udid: UDID, certDirectory: string) {
   const realDevice = await isRealDevice(adb, udid);
-  const proxy = await setupProxyServer(uuid(), udid, realDevice, certDirectory);
+  const proxy = await setupProxyServer(adb, uuid(), udid, realDevice, certDirectory);
   addMock(proxy);
   await configureWifiProxy(adb, udid, realDevice, proxy.options);
   await openUrl(adb, udid, MOCK_BACKEND_URL);

src/utils/proxy.ts

@@ -18,6 +18,7 @@ import { minimatch } from 'minimatch';
 import http from 'http';
 import jsonpath from 'jsonpath';
 import regexParser from 'regex-parser';
+import ADB from 'appium-adb';
 import { validateMockConfig } from '../schema';
 import log from '../logger';
 
@@ -109,6 +110,7 @@ export function modifyResponseBody(ctx: IContext, mockConfig: MockConfig) {
 }
 
 export async function setupProxyServer(
+  adb: ADB,
   sessionId: string,
   deviceUDID: string,
   isRealDevice: boolean,
@@ -127,7 +129,7 @@ export async function setupProxyServer(
   const port = interceptionPort ? Number(interceptionPort) : await getPort();
   log.info(`Selected port: ${port}`);
   const _ip = isRealDevice ? 'localhost' : ip.address('public', 'ipv4');
-  const proxy = new Proxy({ deviceUDID, sessionId, certificatePath, port, ip: _ip, previousConfig: currentWifiProxyConfig, whitelistedDomains, blacklistedDomains});
+  const proxy = new Proxy({ adb, deviceUDID, sessionId, certificatePath, port, ip: _ip, isRealDevice, previousConfig: currentWifiProxyConfig, whitelistedDomains, blacklistedDomains});
   await proxy.start();
   if (!proxy.isStarted()) {
     throw new Error('Unable to start the proxy server');