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,57 @@ 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}`,
+            );
+          }
+        }
+      }
+    };
+
+    const cleanupWithTimeout = async (signal: string, timeoutMs: number = 10000) => {
+      return Promise.race([
+        cleanupAllProxies(signal),
+        new Promise((_, reject) =>
+          setTimeout(() => reject(new Error('Cleanup timeout')), timeoutMs),
+        ),
+      ]);
+    };
+
+    process.once('SIGINT', async () => {
+      try {
+        await cleanupWithTimeout('SIGINT');
+      } catch (err: any) {
+        log.error(`Cleanup failed or timed out: ${err.message}`);
+      }
+    });
+
+    process.once('SIGTERM', async () => {
+      try {
+        await cleanupWithTimeout('SIGTERM');
+      } catch (err: any) {
+        log.error(`Cleanup failed or timed out: ${err.message}`);
+      }
+    });
   }
 
   /**
@@ -163,6 +215,14 @@ 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) {
@@ -258,8 +318,15 @@ export class AppiumInterceptorPlugin extends BasePlugin {
     return proxy;
   }
 
-  private async setupProxy(adb: ADBInstance, sessionId: string, deviceUDID: UDID, interceptionPort?: number) {
-    log.debug(`setupProxy(sessionId=${sessionId}, deviceUDID:${deviceUDID}, interceptionPort:${interceptionPort})`);
+  private async setupProxy(
+    adb: ADBInstance,
+    sessionId: string,
+    deviceUDID: UDID,
+    interceptionPort?: number,
+  ) {
+    log.debug(
+      `setupProxy(sessionId=${sessionId}, deviceUDID:${deviceUDID}, interceptionPort:${interceptionPort})`,
+    );
 
     if (proxyCache.get(sessionId)) {
       log.warn(`[${sessionId}] A proxy is already active for this session. Skipping setup.`);
@@ -286,6 +353,7 @@ export class AppiumInterceptorPlugin extends BasePlugin {
           : parseJson(this.pluginArgs.blacklisteddomains),
       );
       const proxy = await setupProxyServer(
+        adb,
         sessionId,
         deviceUDID,
         realDevice,
@@ -307,18 +375,49 @@ 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
+        try {
+          await configureWifiProxy(
+            activeAdb,
+            proxy.options.deviceUDID,
+            isReal,
+            proxy.previousGlobalProxy,
+          );
+        } catch (err: any) {
+          log.warn(`[${sessionId}] Failed to revert WiFi proxy settings: ${err.message}`);
+        }
+
+        // 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,18 @@ 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({
+    deviceUDID: deviceUDID,
+    sessionId: sessionId,
+    certificatePath: certificatePath,
+    port: port,
+    ip: _ip,
+    adb: adb,
+    isRealDevice: isRealDevice,
+    previousConfig: currentWifiProxyConfig,
+    whitelistedDomains: whitelistedDomains,
+    blacklistedDomains: blacklistedDomains,
+  });
   await proxy.start();
   if (!proxy.isStarted()) {
     throw new Error('Unable to start the proxy server');