feat(proxy): improve user connection tracking with active and historical stats and encryption

2025-12-01 23:14:04 +01:00 · 2025-10-09 17:05:41 +01:00
parent 078b41caaf
commit e3dfa23ccc
2 changed files with 272 additions and 74 deletions
@@ -5,6 +5,7 @@ import {
  Env,
  makeRequest,
  encryptString,
+  decryptString,
  Cache,
 } from '../utils/index.js';
 import path from 'path';
@@ -90,39 +91,244 @@ export class BuiltinProxy extends BaseProxy {
  }
 }

+interface ConnectionRecord {
+  ip: string;
+  url: string;
+  filename?: string;
+  timestamp: number;
+  lastSeen: number;
+  count: number;
+}
+
+interface UserStats {
+  active: ConnectionRecord[];
+  history: ConnectionRecord[];
+}
+
 export class BuiltinProxyStats {
-  private activeConnections = Cache.getInstance<
-    string,
-    { ip: string; url: string; filename?: string; timestamp: number }[]
-  >('bproxy:stats', 10000, 'sql');
+  private activeConnections = Cache.getInstance<string, string>(
+    'bproxy:active',
+    10000,
+    'sql'
+  );
+  private connectionHistory = Cache.getInstance<string, string>(
+    'bproxy:history',
+    10000,
+    'sql'
+  );
+
+  private readonly ACTIVE_THRESHOLD = 60 * 60 * 1000; // 1 hour in milliseconds

  constructor() {}

+  private encryptConnectionRecords(connections: ConnectionRecord[]): string {
+    const result = encryptString(JSON.stringify(connections));
+    if (!result.success || !result.data) {
+      throw new Error(`Failed to encrypt connection records: ${result.error}`);
+    }
+    return result.data;
+  }
+
+  private decryptConnectionRecords(encryptedData: string): ConnectionRecord[] {
+    const result = decryptString(encryptedData);
+    if (!result.success || !result.data) {
+      logger.warn('Failed to decrypt connection records', {
+        error: result.error,
+      });
+      return [];
+    }
+    try {
+      return JSON.parse(result.data);
+    } catch (error) {
+      logger.warn('Failed to parse decrypted connection records', { error });
+      return [];
+    }
+  }
+
+  public async getAllUserStats(): Promise<Map<string, UserStats>> {
+    const users = Env.BUILTIN_PROXY_AUTH?.keys();
+    const userStats = new Map<string, UserStats>();
+
+    for (const user of users ?? []) {
+      userStats.set(user, await this.getUserStats(user));
+    }
+    return userStats;
+  }
+
+  public async getUserStats(user: string): Promise<UserStats> {
+    const [active, history] = await Promise.all([
+      this.getActiveConnections(user),
+      this.getConnectionHistory(user),
+    ]);
+
+    return { active, history };
+  }
+
+  public async getActiveConnections(user: string): Promise<ConnectionRecord[]> {
+    const encryptedData = await this.activeConnections.get(user);
+    const connections = encryptedData
+      ? this.decryptConnectionRecords(encryptedData)
+      : [];
+    const now = Date.now();
+
+    // Filter out connections older than 1 hour and move them to history
+    const activeConnections: ConnectionRecord[] = [];
+    const expiredConnections: ConnectionRecord[] = [];
+
+    for (const conn of connections) {
+      if (now - conn.lastSeen <= this.ACTIVE_THRESHOLD) {
+        activeConnections.push(conn);
+      } else {
+        expiredConnections.push(conn);
+      }
+    }
+
+    // Move expired connections to history
+    if (expiredConnections.length > 0) {
+      await this.moveToHistory(user, expiredConnections);
+      await this.activeConnections.set(
+        user,
+        this.encryptConnectionRecords(activeConnections),
+        24 * 60 * 60
+      );
+    }
+
+    return activeConnections;
+  }
+
+  public async getConnectionHistory(user: string): Promise<ConnectionRecord[]> {
+    const encryptedData = await this.connectionHistory.get(user);
+    return encryptedData ? this.decryptConnectionRecords(encryptedData) : [];
+  }
+
+  public async addConnection(
+    user: string,
+    ip: string,
+    url: string,
+    timestamp: number,
+    filename?: string
+  ) {
+    logger.debug(`[${user}] Adding connection`, {
+      ip,
+      url,
+      filename,
+      timestamp,
+    });
+
+    const connectionKey = `${ip}:${url}`;
+    const now = Date.now();
+
+    // Get current active connections
+    const activeConnections = await this.getActiveConnections(user);
+
+    // Check if this connection already exists in active connections
+    const existingIndex = activeConnections.findIndex(
+      (conn) => `${conn.ip}:${conn.url}` === connectionKey
+    );
+
+    if (existingIndex >= 0) {
+      const existing = activeConnections[existingIndex];
+      activeConnections[existingIndex] = {
+        ...existing,
+        lastSeen: now,
+        count: existing.count + 1,
+      };
+    } else {
+      // Add new connection
+      activeConnections.push({
+        ip,
+        url,
+        filename,
+        timestamp,
+        lastSeen: now,
+        count: 1,
+      });
+    }
+
+    // Sort by lastSeen (most recent first)
+    activeConnections.sort((a, b) => b.lastSeen - a.lastSeen);
+
+    await this.activeConnections.set(
+      user,
+      this.encryptConnectionRecords(activeConnections),
+      24 * 60 * 60
+    );
+  }
+
+  public async removeConnection(user: string, ip: string, url: string) {
+    const activeConnections = await this.getActiveConnections(user);
+    const connectionKey = `${ip}:${url}`;
+
+    const filteredConnections = activeConnections.filter(
+      (conn) => `${conn.ip}:${conn.url}` !== connectionKey
+    );
+
+    await this.activeConnections.set(
+      user,
+      this.encryptConnectionRecords(filteredConnections),
+      24 * 60 * 60
+    );
+  }
+
+  private async moveToHistory(user: string, connections: ConnectionRecord[]) {
+    const existingHistory = await this.getConnectionHistory(user);
+
+    // Merge with existing history, keeping the most recent record for each connection
+    const historyMap = new Map<string, ConnectionRecord>();
+
+    // Add existing history
+    for (const conn of existingHistory) {
+      const key = `${conn.ip}:${conn.url}`;
+      historyMap.set(key, conn);
+    }
+
+    // Add/update with new connections
+    for (const conn of connections) {
+      const key = `${conn.ip}:${conn.url}`;
+      const existing = historyMap.get(key);
+
+      if (!existing || conn.lastSeen > existing.lastSeen) {
+        historyMap.set(key, conn);
+      } else if (existing) {
+        // Merge counts if the existing record is more recent
+        existing.count += conn.count;
+      }
+    }
+
+    const updatedHistory = Array.from(historyMap.values()).sort(
+      (a, b) => b.lastSeen - a.lastSeen
+    );
+
+    await this.connectionHistory.set(
+      user,
+      this.encryptConnectionRecords(updatedHistory),
+      7 * 24 * 60 * 60
+    ); // Keep history for 7 days
+  }
+
+  // Legacy methods for backward compatibility
  public async getAllActiveConnections(): Promise<
    Map<
      string,
      { ip: string; url: string; filename?: string; timestamp: number }[]
    >
  > {
-    const users = Env.BUILTIN_PROXY_AUTH?.keys();
+    const userStats = await this.getAllUserStats();
+    const result = new Map();

-    // create a map of users and their active connections
-    const connections = new Map<
-      string,
-      { ip: string; url: string; filename?: string; timestamp: number }[]
-    >();
-    for (const user of users ?? []) {
-      connections.set(user, await this.getActiveConnections(user));
+    for (const [user, stats] of userStats) {
+      result.set(
+        user,
+        stats.active.map((conn) => ({
+          ip: conn.ip,
+          url: conn.url,
+          filename: conn.filename,
+          timestamp: conn.timestamp,
+        }))
+      );
    }
-    return connections;
-  }

-  public async getActiveConnections(
-    user: string
-  ): Promise<
-    { ip: string; url: string; filename?: string; timestamp: number }[]
-  > {
-    return (await this.activeConnections.get(user)) ?? [];
+    return result;
  }

  public async addActiveConnection(
@@ -132,37 +338,10 @@ export class BuiltinProxyStats {
    timestamp: number,
    filename?: string
  ) {
-    logger.debug(`[${user}] Adding active connection`, {
-      ip,
-      url,
-      filename,
-      timestamp,
-    });
-
-    const existingConnections = (await this.activeConnections.get(user)) ?? [];
-    const connectionKey = `${ip}:${url}`;
-
-    // Filter out any existing connections with the same IP+filename combination
-    const filteredConnections = existingConnections.filter((conn) => {
-      return `${conn.ip}:${conn.url}` !== connectionKey;
-    });
-
-    // Add the new connection (which will be the most recent for this IP+filename)
-    const updatedConnections = [
-      ...filteredConnections,
-      { ip, url, filename, timestamp },
-    ];
-
-    await this.activeConnections.set(user, updatedConnections, 1 * 60 * 60);
+    return this.addConnection(user, ip, url, timestamp, filename);
  }

  public async removeActiveConnection(user: string, ip: string, url: string) {
-    await this.activeConnections.set(
-      user,
-      ((await this.activeConnections.get(user)) ?? []).filter(
-        (connection) => connection.ip !== ip && connection.url !== url
-      ),
-      24 * 60 * 60
-    );
+    return this.removeConnection(user, ip, url);
  }
 }
@@ -93,30 +93,48 @@ router.get(
    }

    try {
-      const allConnections = await proxyStats.getAllActiveConnections();
+      const allUserStats = await proxyStats.getAllUserStats();

      // Convert Map to a more JSON-friendly format
      const stats = {
        timestamp: new Date().toISOString(),
-        totalUsers: allConnections.size,
-        activeConnections: Object.fromEntries(
-          Array.from(allConnections.entries()).map(([user, connections]) => [
+        totalUsers: allUserStats.size,
+        users: Object.fromEntries(
+          Array.from(allUserStats.entries()).map(([user, userStats]) => [
            user,
-            connections.map((conn) => ({
-              ...conn,
-              timestamp: new Date(conn.timestamp).toISOString(),
-              relativeTimestamp: `${getTimeTakenSincePoint(conn.timestamp)} ago`,
-            })),
+            {
+              active: userStats.active.map((conn) => ({
+                ...conn,
+                timestamp: new Date(conn.timestamp).toISOString(),
+                lastSeen: new Date(conn.lastSeen).toISOString(),
+                relativeTimestamp: `${getTimeTakenSincePoint(conn.timestamp)} ago`,
+                relativeLastSeen: `${getTimeTakenSincePoint(conn.lastSeen)} ago`,
+              })),
+              history: userStats.history.map((conn) => ({
+                ...conn,
+                timestamp: new Date(conn.timestamp).toISOString(),
+                lastSeen: new Date(conn.lastSeen).toISOString(),
+                relativeTimestamp: `${getTimeTakenSincePoint(conn.timestamp)} ago`,
+                relativeLastSeen: `${getTimeTakenSincePoint(conn.lastSeen)} ago`,
+              })),
+            },
          ])
        ),
        summary: {
-          totalActiveConnections: Array.from(allConnections.values()).reduce(
-            (total, connections) => total + connections.length,
+          totalActiveConnections: Array.from(allUserStats.values()).reduce(
+            (total, userStats) => total + userStats.active.length,
            0
          ),
-          usersWithActiveConnections: Array.from(
-            allConnections.entries()
-          ).filter(([_, connections]) => connections.length > 0).length,
+          totalHistoryConnections: Array.from(allUserStats.values()).reduce(
+            (total, userStats) => total + userStats.history.length,
+            0
+          ),
+          usersWithActiveConnections: Array.from(allUserStats.entries()).filter(
+            ([_, userStats]) => userStats.active.length > 0
+          ).length,
+          usersWithHistory: Array.from(allUserStats.entries()).filter(
+            ([_, userStats]) => userStats.history.length > 0
+          ).length,
        },
      };

@@ -175,16 +193,17 @@ router.all(
        );
      }

-      // Track the active connection
-      clientIp = req.ip || req.connection.remoteAddress || 'unknown';
+      // Track the connection
+      clientIp =
+        req.requestIp || req.ip || req.socket.remoteAddress || 'unknown';
      const timestamp = Date.now();
-      proxyStats.addActiveConnection(
-        auth.username,
-        clientIp,
-        data.url,
-        timestamp,
-        filename
-      );
+      proxyStats
+        .addConnection(auth.username, clientIp, data.url, timestamp, filename)
+        .catch((error) =>
+          logger.warn(`[${requestId}] Failed to add connection to stats`, {
+            error: error instanceof Error ? error.message : String(error),
+          })
+        );

      // prepare and execute upstream request
      const { host, ...clientHeaders } = req.headers;
@@ -251,10 +270,10 @@ router.all(
    } catch (error) {
      const totalDuration = Date.now() - startTime;

-      // Remove the active connection tracking on error
+      // Remove the connection tracking on error
      if (auth && clientIp && data) {
        proxyStats
-          .removeActiveConnection(auth.username, clientIp, data.url)
+          .removeConnection(auth.username, clientIp, data.url)
          .catch((statsError) =>
            logger.warn(
              `[${requestId}] Failed to remove connection from stats on error`,