From 45c30a335415ed32edfb02d03392f50de8e2d566 Mon Sep 17 00:00:00 2001 From: Vincent DANTI Date: Tue, 16 Jun 2026 16:03:01 +0200 Subject: [PATCH] fix: resolve reverse tunnel and wifi proxy config leak on session end or crash --- src/plugin.ts | 110 +++++++++++++++++++++++++++++++-- src/proxy-cache.ts | 8 +++ src/proxy.ts | 3 + src/scripts/test-connection.ts | 2 +- src/utils/proxy.ts | 15 ++++- 5 files changed, 131 insertions(+), 7 deletions(-) diff --git a/src/plugin.ts b/src/plugin.ts index cfcfc7d..7555fa6 100644 --- a/src/plugin.ts +++ b/src/plugin.ts @@ -18,6 +18,7 @@ import { isRealDevice, getAdbReverseTunnels, getCurrentWifiProxyConfig, + removeReverseTunnel, ADBInstance, UDID, } from './utils/adb'; @@ -89,6 +90,62 @@ 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); + } + }; + + 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 +220,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 +323,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 +358,7 @@ export class AppiumInterceptorPlugin extends BasePlugin { : parseJson(this.pluginArgs.blacklisteddomains), ); const proxy = await setupProxyServer( + adb, sessionId, deviceUDID, realDevice, @@ -307,18 +380,45 @@ 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); diff --git a/src/proxy-cache.ts b/src/proxy-cache.ts index 057d2a5..84d7419 100644 --- a/src/proxy-cache.ts +++ b/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(); diff --git a/src/proxy.ts b/src/proxy.ts index 9613991..97b02d3 100644 --- a/src/proxy.ts +++ b/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[]; diff --git a/src/scripts/test-connection.ts b/src/scripts/test-connection.ts index 760d816..691dbbc 100644 --- a/src/scripts/test-connection.ts +++ b/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); diff --git a/src/utils/proxy.ts b/src/utils/proxy.ts index dd56883..c9d7c95 100644 --- a/src/utils/proxy.ts +++ b/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');