fix(double agent): clean up backend

Author: blakebyrnes

double-agent/collect-controller/lib/Server.ts

@@ -255,12 +255,13 @@ export default class Server {
 
     await this.saveMetaFiles(activeScraper, assignment);
     const profilesDir = extractAssignmentProfilesDir(activeScraper, assignment);
-    await pipeDirToStream(profilesDir, res);
-
     const session = this.collect.getSession(assignment.sessionId);
-    await this.collect.deleteSession(session);
-
-    delete this.activeUsersById[userId];
+    try {
+      await pipeDirToStream(profilesDir, res);
+    } finally {
+      await this.collect.deleteSession(session);
+      delete this.activeUsersById[userId];
+    }
   }
 
   private async downloadAll(

double-agent/collect/lib/SessionTracker.ts

@@ -1,6 +1,7 @@
 import { IAssignmentType } from '@double-agent/collect-controller/interfaces/IAssignment';
 import * as http from 'http';
 import * as http2 from 'http2';
+import Config from '@double-agent/config';
 import Session from './Session';
 import PluginDelegate from './PluginDelegate';
 import BaseServer from '../servers/BaseServer';
@@ -10,6 +11,7 @@ let sessionIdCounter = 0;
 export default class SessionTracker {
   private pluginDelegate: PluginDelegate = new PluginDelegate();
   private sessions: { [sessionId: string]: Session } = {};
+  private sessionExpiryById: { [sessionId: string]: NodeJS.Timeout } = {};
 
   public async createSession(
     assignmentType: IAssignmentType,
@@ -21,13 +23,18 @@ export default class SessionTracker {
     await session.startServers();
 
     this.sessions[sessionId] = session;
+    this.scheduleSessionExpiry(sessionId);
     return session;
   }
 
   public getSession(sessionId: string): Session {
     return this.sessions[sessionId];
   }
 
+  public touchSession(sessionId: string): void {
+    this.scheduleSessionExpiry(sessionId);
+  }
+
   public getSessionIdFromServerRequest(
     server: BaseServer,
     req: http.IncomingMessage | http2.Http2ServerRequest,
@@ -48,12 +55,40 @@ export default class SessionTracker {
 
   public async deleteSession(sessionId: string): Promise<void> {
     if (!this.sessions[sessionId]) return;
+    this.clearSessionExpiry(sessionId);
     await this.sessions[sessionId].close();
     delete this.sessions[sessionId];
   }
 
   public async shutdown(): Promise<void> {
     await Promise.allSettled(Object.values(this.sessions).map((x) => x.close()));
     await Promise.allSettled(this.pluginDelegate.plugins.map((x) => x.stop()));
+    this.clearAllSessionExpiry();
+  }
+
+  private scheduleSessionExpiry(sessionId: string): void {
+    const ttlMs = Config.collect.sessionTtlMs;
+    if (!ttlMs || ttlMs <= 0) return;
+    this.clearSessionExpiry(sessionId);
+    const timeout = setTimeout(() => {
+      console.warn('Session expired due to inactivity/ttl', sessionId);
+      void this.deleteSession(sessionId);
+    }, ttlMs);
+    timeout.unref();
+    this.sessionExpiryById[sessionId] = timeout;
+  }
+
+  private clearSessionExpiry(sessionId: string): void {
+    const timeout = this.sessionExpiryById[sessionId];
+    if (timeout) {
+      clearTimeout(timeout);
+      delete this.sessionExpiryById[sessionId];
+    }
+  }
+
+  private clearAllSessionExpiry(): void {
+    for (const sessionId of Object.keys(this.sessionExpiryById)) {
+      this.clearSessionExpiry(sessionId);
+    }
   }
 }

double-agent/collect/lib/createHttpRequestHandler.ts

@@ -55,6 +55,7 @@ export default function createHttpRequestHandler(
       if (!session) {
         throw new Error(`Missing session: ${sessionId}`);
       }
+      sessionTracker.touchSession(sessionId);
       const { requestDetails } = await extractRequestDetails(server, req, session);
       const ctx = new RequestContext(server, req, res, requestUrl, requestDetails, session);
       const userAgentId = session.userAgentId;

double-agent/collect/lib/createTlsRequestHandler.ts

@@ -24,6 +24,7 @@ export default function createTlsRequestHandler(
     }
 
     const session = sessionTracker.getSessionFromServerRequest(server, req);
+    sessionTracker.touchSession(session.id);
     const { requestDetails } = await extractRequestDetails(server, req, session);
     const ctx = new RequestContext(server, req, res, requestUrl, requestDetails, session);
     const handler = server.getHandlerFn(requestUrl.pathname);

double-agent/collect/lib/createWebsocketHandler.ts

@@ -20,6 +20,7 @@ export default function createWebsocketHandler(
   ): Promise<void> {
     const { sessionTracker } = detectionContext;
     const session = sessionTracker.getSessionFromServerRequest(server, req);
+    sessionTracker.touchSession(session.id);
     const { requestDetails, requestUrl } = await extractRequestDetails(
       server,
       req,

double-agent/config/index.ts

@@ -55,6 +55,7 @@ export default class Config {
     shouldGenerateProfiles: parseEnvBool(env.GENERATE_PROFILES),
     pluginStartingPort: parseEnvInt(env.PLUGIN_STARTING_PORT),
     pluginMaxPort: parseEnvInt(env.PLUGIN_MAX_PORT) || 20000,
+    sessionTtlMs: parseEnvInt(env.COLLECT_SESSION_TTL_MS) ?? 60 * 60 * 1000,
 
     // collect plugins
     tcpNetworkDevice: env.TCP_NETWORK_DEVICE,

double-agent/tls-server/child.ts

@@ -10,6 +10,13 @@ let connectionCount = 0;
 let lastConnectionDate: Date;
 let activeConnection: { id: number; req: IncomingMessage; res: ServerResponse };
 const connections: (typeof activeConnection)[] = [];
+let childServer: https.Server;
+let isShuttingDown = false;
+
+process.on('disconnect', () => {
+  shutdown('parent disconnect');
+  process.exit(0);
+});
 
 process.on('message', (message: any) => {
   if (message.start) {
@@ -54,7 +61,7 @@ function start(options: { port: number; key?: string; cert?: string }): void {
       enableTrace: true,
       sessionTimeout: 10,
     });
-    const childServer = https.createServer(options, onConnection);
+    childServer = https.createServer(options, onConnection);
 
     childServer.on('error', err => {
       process.send({ error: err.message });
@@ -65,21 +72,25 @@ function start(options: { port: number; key?: string; cert?: string }): void {
       process.send({ started: true });
     });
 
-    ShutdownHandler.register(() => {
-      if (childServer) childServer.unref().close();
-      while (connections.length) {
-        const connection = connections.pop();
-
-        console.log('Force closing active connection during shutdown.');
-        connection.res.end('Server shutting down');
-        connection.req.destroy();
-      }
-    });
+    ShutdownHandler.register(() => shutdown('shutdown handler'));
   } catch (err) {
     console.log(err);
   }
 }
 
+function shutdown(reason?: string): void {
+  if (isShuttingDown) return;
+  isShuttingDown = true;
+  if (reason) console.log(`TLS child shutdown: ${reason}`);
+  if (childServer) childServer.unref().close();
+  while (connections.length) {
+    const connection = connections.pop();
+    console.log('Force closing active connection during shutdown.');
+    connection.res.end('Server shutting down');
+    connection.req.destroy();
+  }
+}
+
 async function onConnection(req, res): Promise<void> {
   lastConnectionDate = new Date();
   res.connection.setKeepAlive(false);