Add device traffic reset with outbound baselines

2026-08-12 23:55:34 +03:00
parent 08cc013def
commit b86812d02b
18 changed files with 369 additions and 33 deletions
@@ -89,6 +89,7 @@ interface InventoryTrafficState {
  baselinesByMac: Record<string, CounterBaseline>;
  totalsByMac: Record<string, TrafficTotal>;
  rebaselineMacs: string[];
+  outboundBaselinesByDeviceId: Record<string, OutboundTrafficBaseline>;
  proxy: ProxyTrafficState;
  global: { gateway: GlobalTrafficSource; proxy: GlobalTrafficSource };
  [key: string]: unknown;
@@ -164,6 +165,15 @@ interface OutboundTrafficCursor {
  unknown: bigint;
 }

+interface OutboundTrafficBaseline {
+  routeEpoch: string;
+  directEpoch: string;
+  vpnBytes: string;
+  directTrackedBytes: string;
+  directIpv4Bytes: string;
+  unknownBytes: string;
+}
+
 interface DeviceObservation {
  mac: string;
  ip: string;
@@ -250,6 +260,7 @@ const DEFAULT_STATE: InventoryState = {
    baselinesByMac: {},
    totalsByMac: {},
    rebaselineMacs: [],
+    outboundBaselinesByDeviceId: {},
    proxy: DEFAULT_PROXY_TRAFFIC,
    global: {
      gateway: DEFAULT_GLOBAL_TRAFFIC_SOURCE,
@@ -599,6 +610,32 @@ export function migrateDeviceInventoryState(value: unknown): InventoryState {
      recoveredTraffic = true;
    }
  }
+  const knownDeviceIds = new Set(devices.map(({ id }) => id));
+  const outboundBaselinesByDeviceId: Record<string, OutboundTrafficBaseline> = {};
+  if (traffic.outboundBaselinesByDeviceId !== undefined
+    && record(traffic.outboundBaselinesByDeviceId) !== traffic.outboundBaselinesByDeviceId) {
+    recoveredTraffic = true;
+  }
+  for (const [id, baseline] of recordEntries(traffic.outboundBaselinesByDeviceId)) {
+    const vpnBytes = parseStoredCounter(baseline.vpnBytes);
+    const directTrackedBytes = parseStoredCounter(baseline.directTrackedBytes);
+    const directIpv4Bytes = parseStoredCounter(baseline.directIpv4Bytes);
+    const unknownBytes = parseStoredCounter(baseline.unknownBytes);
+    if (!DEVICE_ID_PATTERN.test(id) || !knownDeviceIds.has(id)
+      || typeof baseline.routeEpoch !== 'string' || typeof baseline.directEpoch !== 'string'
+      || vpnBytes == null || directTrackedBytes == null || directIpv4Bytes == null || unknownBytes == null) {
+      recoveredTraffic = true;
+      continue;
+    }
+    outboundBaselinesByDeviceId[id] = {
+      routeEpoch: baseline.routeEpoch,
+      directEpoch: baseline.directEpoch,
+      vpnBytes,
+      directTrackedBytes,
+      directIpv4Bytes,
+      unknownBytes,
+    };
+  }
  const global = {
    gateway: normalizeGlobalTrafficSource(record(traffic.global).gateway, {
      epoch: typeof traffic.epoch === 'string' ? traffic.epoch : null,
@@ -630,6 +667,7 @@ export function migrateDeviceInventoryState(value: unknown): InventoryState {
      baselinesByMac,
      totalsByMac,
      rebaselineMacs: [...rebaselineMacs].filter((mac) => MAC_PATTERN.test(mac)),
+      outboundBaselinesByDeviceId,
      proxy: proxyTraffic,
      global,
    },
@@ -837,15 +875,28 @@ export function createDeviceInventoryService({
      const total = totalFor(device.id);
      const current: OutboundTrafficCursor = { signature, routeEpoch, directEpoch, ...total };
      const previous = outboundTrafficCursorByDeviceId.get(device.id);
+      const baseline = state.traffic.outboundBaselinesByDeviceId[device.id];
+      const routeBaseline = baseline?.routeEpoch === routeEpoch ? baseline : null;
+      const directBaseline = baseline?.directEpoch === directEpoch ? baseline : null;
+      const routeVpn = BigInt(routeBaseline?.vpnBytes || '0');
+      const routeDirect = BigInt(routeBaseline?.directTrackedBytes || '0');
+      const routeUnknown = BigInt(routeBaseline?.unknownBytes || '0');
+      const directIpv4Baseline = BigInt(directBaseline?.directIpv4Bytes || '0');
+      const visible = {
+        vpn: total.vpn >= routeVpn ? total.vpn - routeVpn : 0n,
+        directTracked: total.directTracked >= routeDirect ? total.directTracked - routeDirect : 0n,
+        directIpv4: total.directIpv4 >= directIpv4Baseline ? total.directIpv4 - directIpv4Baseline : 0n,
+        unknown: total.unknown >= routeUnknown ? total.unknown - routeUnknown : 0n,
+      };
      outboundTrafficCursorByDeviceId.set(device.id, current);
      if (observedAt) outboundTrafficByDeviceId.set(device.id, {
        observedAt,
        singboxObservedAt: routeObservedAt || null,
        directIpv4ObservedAt: directObservedAt || null,
-        vpnBytes: total.vpn.toString(),
-        directTrackedBytes: total.directTracked.toString(),
-        directIpv4Bytes: total.directIpv4.toString(),
-        unknownBytes: total.unknown.toString(),
+        vpnBytes: visible.vpn.toString(),
+        directTrackedBytes: visible.directTracked.toString(),
+        directIpv4Bytes: visible.directIpv4.toString(),
+        unknownBytes: visible.unknown.toString(),
      });
      if (!previous || previous.signature === signature || !observedAt) continue;
      const routeDelta = (value: bigint, before: bigint) => (
@@ -1542,6 +1593,70 @@ export function createDeviceInventoryService({
    return snapshot();
  }

+  async function resetTraffic(expectedRevision: unknown) {
+    if (typeof expectedRevision !== 'number' || !Number.isSafeInteger(expectedRevision)
+      || expectedRevision < 0) {
+      throw new HarborError('REQUEST_INVALID');
+    }
+    if (refreshPromise) await refreshPromise;
+    const nextState = store.update((stored) => {
+      const state = migrateDeviceInventoryState(stored);
+      if (state.revision !== expectedRevision) throw new HarborError('STATE_CONFLICT');
+      const totalsByMac: Record<string, TrafficTotal> = {};
+      const proxyTotalsByMac: Record<string, TrafficTotal> = {};
+      const outboundBaselinesByDeviceId: Record<string, OutboundTrafficBaseline> = {};
+      const rebaselineMacs = new Set(state.traffic.rebaselineMacs);
+      const proxyRebaselineMacs = new Set(state.traffic.proxy.rebaselineMacs);
+      for (const device of state.devices) {
+        const traffic = state.traffic.totalsByMac[device.mac];
+        const proxy = state.traffic.proxy.totalsByMac[device.mac];
+        totalsByMac[device.mac] = {
+          uploadBytes: '0',
+          downloadBytes: '0',
+          observedAt: traffic?.observedAt || state.traffic.lastObservedAt,
+        };
+        proxyTotalsByMac[device.mac] = {
+          uploadBytes: '0',
+          downloadBytes: '0',
+          observedAt: proxy?.observedAt || state.traffic.proxy.lastObservedAt,
+        };
+        if (!state.traffic.baselinesByMac[device.mac]) rebaselineMacs.add(device.mac);
+        if (!state.traffic.proxy.baselinesByMac[device.mac]) proxyRebaselineMacs.add(device.mac);
+        const outbound = outboundTrafficCursorByDeviceId.get(device.id);
+        if (outbound) outboundBaselinesByDeviceId[device.id] = {
+          routeEpoch: outbound.routeEpoch,
+          directEpoch: outbound.directEpoch,
+          vpnBytes: outbound.vpn.toString(),
+          directTrackedBytes: outbound.directTracked.toString(),
+          directIpv4Bytes: outbound.directIpv4.toString(),
+          unknownBytes: outbound.unknown.toString(),
+        };
+      }
+      return {
+        ...state,
+        revision: state.revision + 1,
+        traffic: {
+          ...state.traffic,
+          totalsByMac,
+          rebaselineMacs: [...rebaselineMacs],
+          outboundBaselinesByDeviceId,
+          proxy: {
+            ...state.traffic.proxy,
+            totalsByMac: proxyTotalsByMac,
+            rebaselineMacs: [...proxyRebaselineMacs],
+          },
+        },
+      };
+    });
+
+    trafficHistoryByMac.clear();
+    trafficCursorByMac.clear();
+    captureTrafficHistory(nextState);
+    captureOutboundTrafficHistory(nextState);
+    outboundTrafficHistoryByDeviceId.clear();
+    return snapshot();
+  }
+
  function setPolicy(id: string, mode: unknown, expectedRevision: unknown) {
    if (typeof expectedRevision !== 'number' || !Number.isSafeInteger(expectedRevision)
      || expectedRevision < 0 || !POLICY_MODES.has(mode)) {
@@ -1585,5 +1700,5 @@ export function createDeviceInventoryService({
    return serializePolicy(() => reconcileLocked(observed, true));
  }

-  return { snapshot, metricsSnapshot, refresh, update, setPolicy, reconcilePolicies };
+  return { snapshot, metricsSnapshot, refresh, update, resetTraffic, setPolicy, reconcilePolicies };
 }