Track proxy traffic separately in device inventory

src/server/services/deviceInventoryService.js

@@ -15,6 +15,7 @@ const FINGERPRINT_PATTERN = /^[a-f0-9]{64}$/;
 const INTERFACE_PATTERN = /^[a-z0-9_.:-]{1,15}$/i;
 const POLICY_MODES = new Set(['vpn', 'direct']);
 const POLICY_STATUSES = new Set(['applied', 'applying', 'pending', 'failed']);
+const PROXY_RECOVERY_ERROR = 'Повреждённый proxy traffic checkpoint восстановлен из корректных данных';
 
 const DEFAULT_DEVICE_POLICY = Object.freeze({
   desired: 'vpn',
@@ -36,6 +37,15 @@ const DEFAULT_POLICY_STATE = {
   byMac: {},
 };
 
+const DEFAULT_PROXY_TRAFFIC = {
+  schemaVersion: 1,
+  lastObservedAt: null,
+  lastError: null,
+  baselinesByMac: {},
+  totalsByMac: {},
+  rebaselineMacs: [],
+};
+
 const DEFAULT_STATE = {
   schemaVersion: DEVICE_INVENTORY_SCHEMA_VERSION,
   revision: 0,
@@ -50,6 +60,7 @@ const DEFAULT_STATE = {
     baselinesByMac: {},
     totalsByMac: {},
     rebaselineMacs: [],
+    proxy: DEFAULT_PROXY_TRAFFIC,
   },
   devices: [],
 };
@@ -65,6 +76,68 @@ const parseStoredCounter = (value) => {
   return COUNTER_PATTERN.test(counter) ? BigInt(counter).toString() : null;
 };
 
+function normalizeProxyTraffic(value, devices) {
+  const proxy = value && typeof value === 'object' && !Array.isArray(value) ? value : {};
+  if (Number.isSafeInteger(proxy.schemaVersion) && proxy.schemaVersion > 1) {
+    throw new Error(`Unsupported proxy traffic schemaVersion: ${proxy.schemaVersion}`);
+  }
+  const rebaselineMacs = new Set((Array.isArray(proxy.rebaselineMacs) ? proxy.rebaselineMacs : [])
+    .map(normalizeMac).filter((mac) => MAC_PATTERN.test(mac)));
+  let recovered = value !== undefined && (
+    proxy !== value || proxy.schemaVersion !== 1
+    || !proxy.baselinesByMac || typeof proxy.baselinesByMac !== 'object' || Array.isArray(proxy.baselinesByMac)
+    || !proxy.totalsByMac || typeof proxy.totalsByMac !== 'object' || Array.isArray(proxy.totalsByMac)
+  );
+  if (recovered) {
+    for (const device of devices) rebaselineMacs.add(normalizeMac(device.mac));
+  }
+  const baselinesByMac = {};
+  for (const [rawMac, baseline] of recordEntries(proxy.baselinesByMac)) {
+    const mac = normalizeMac(rawMac);
+    const uploadBytes = parseStoredCounter(baseline?.uploadBytes);
+    const downloadBytes = parseStoredCounter(baseline?.downloadBytes);
+    if (!MAC_PATTERN.test(mac) || typeof baseline?.epoch !== 'string' || !baseline.epoch
+      || uploadBytes == null || downloadBytes == null) {
+      if (MAC_PATTERN.test(mac)) rebaselineMacs.add(mac);
+      recovered = true;
+      continue;
+    }
+    baselinesByMac[mac] = { epoch: baseline.epoch, uploadBytes, downloadBytes };
+  }
+  const totalsByMac = {};
+  for (const [rawMac, total] of recordEntries(proxy.totalsByMac)) {
+    const mac = normalizeMac(rawMac);
+    const uploadBytes = parseStoredCounter(total?.uploadBytes);
+    const downloadBytes = parseStoredCounter(total?.downloadBytes);
+    if (!MAC_PATTERN.test(mac) || uploadBytes == null || downloadBytes == null) {
+      if (MAC_PATTERN.test(mac)) rebaselineMacs.add(mac);
+      recovered = true;
+      continue;
+    }
+    totalsByMac[mac] = {
+      uploadBytes,
+      downloadBytes,
+      observedAt: typeof total?.observedAt === 'string' ? total.observedAt : null,
+    };
+  }
+  for (const mac of new Set([...Object.keys(baselinesByMac), ...Object.keys(totalsByMac)])) {
+    if (!Object.hasOwn(baselinesByMac, mac) || !Object.hasOwn(totalsByMac, mac)) {
+      rebaselineMacs.add(mac);
+      recovered = true;
+    }
+  }
+  return {
+    schemaVersion: 1,
+    lastObservedAt: typeof proxy.lastObservedAt === 'string' ? proxy.lastObservedAt : null,
+    lastError: recovered
+      ? PROXY_RECOVERY_ERROR
+      : typeof proxy.lastError === 'string' ? proxy.lastError : null,
+    baselinesByMac,
+    totalsByMac,
+    rebaselineMacs: [...rebaselineMacs],
+  };
+}
+
 function normalizePolicyState(value) {
   const policy = value && typeof value === 'object' && !Array.isArray(value) ? value : {};
   const byMac = {};
@@ -142,6 +215,7 @@ export function migrateDeviceInventoryState(value) {
     ? state.traffic
     : {};
   const devices = Array.isArray(state.devices) ? state.devices : [];
+  const proxyTraffic = normalizeProxyTraffic(traffic.proxy, devices);
   const rebaselineMacs = new Set((Array.isArray(traffic.rebaselineMacs) ? traffic.rebaselineMacs : [])
     .map(normalizeMac).filter((mac) => MAC_PATTERN.test(mac)));
   let recoveredTraffic = version >= 2 && (
@@ -202,6 +276,7 @@ export function migrateDeviceInventoryState(value) {
       baselinesByMac,
       totalsByMac,
       rebaselineMacs: [...rebaselineMacs].filter((mac) => MAC_PATTERN.test(mac)),
+      proxy: proxyTraffic,
     },
     devices,
   };
@@ -276,6 +351,7 @@ export function createDeviceInventoryService({
     const rank = { online: 0, recent: 1, offline: 2 };
     const devices = state.devices.map((device) => {
       const traffic = state.traffic.totalsByMac[device.mac];
+      const proxyTraffic = state.traffic.proxy.totalsByMac[device.mac];
       const policy = policyFor(state, device.mac);
       return {
         ...device,
@@ -283,6 +359,9 @@ export function createDeviceInventoryService({
         uploadBytes: traffic?.uploadBytes || '0',
         downloadBytes: traffic?.downloadBytes || '0',
         trafficObservedAt: traffic?.observedAt || null,
+        proxyUploadBytes: proxyTraffic?.uploadBytes || '0',
+        proxyDownloadBytes: proxyTraffic?.downloadBytes || '0',
+        proxyTrafficObservedAt: proxyTraffic?.observedAt || null,
         desiredPolicy: policy.desired,
         appliedPolicy: policy.applied,
         policyStatus: policy.status,
@@ -303,6 +382,10 @@ export function createDeviceInventoryService({
         traffic: {
           lastObservedAt: state.traffic.lastObservedAt,
           error: state.traffic.lastError,
+          proxy: {
+            lastObservedAt: state.traffic.proxy.lastObservedAt,
+            error: state.traffic.proxy.lastError,
+          },
         },
         policy: {
           lastAppliedAt: state.policy.lastAppliedAt,
@@ -489,8 +572,13 @@ export function createDeviceInventoryService({
             if (typeof trafficResult.epoch !== 'string' || !trafficResult.epoch) {
               throw new Error('Dataplane не вернул traffic epoch');
             }
+            const rows = Array.isArray(trafficResult.devices) ? trafficResult.devices : [];
             const processByMac = new Map();
-            for (const row of Array.isArray(trafficResult.devices) ? trafficResult.devices : []) {
+            const proxyByMac = new Map();
+            let proxyRows = 0;
+            let legacyRows = 0;
+            let proxySampleError = null;
+            for (const row of rows) {
               const mac = normalizeMac(row?.mac);
               const upload = String(row?.uploadBytes ?? '');
               const download = String(row?.downloadBytes ?? '');
@@ -502,7 +590,30 @@ export function createDeviceInventoryService({
                 upload: previous.upload + BigInt(upload),
                 download: previous.download + BigInt(download),
               });
+              const hasProxyUpload = Object.hasOwn(row || {}, 'proxyUploadBytes');
+              const hasProxyDownload = Object.hasOwn(row || {}, 'proxyDownloadBytes');
+              if (!hasProxyUpload && !hasProxyDownload) {
+                legacyRows += 1;
+                continue;
+              }
+              if (!hasProxyUpload || !hasProxyDownload) {
+                proxySampleError = 'Dataplane вернул неполный proxy traffic counter';
+                continue;
+              }
+              const proxyUpload = String(row.proxyUploadBytes);
+              const proxyDownload = String(row.proxyDownloadBytes);
+              if (!COUNTER_PATTERN.test(proxyUpload) || !COUNTER_PATTERN.test(proxyDownload)) {
+                proxySampleError = 'Dataplane вернул невалидный proxy traffic counter';
+                continue;
+              }
+              proxyRows += 1;
+              const previousProxy = proxyByMac.get(mac) || { upload: 0n, download: 0n };
+              proxyByMac.set(mac, {
+                upload: previousProxy.upload + BigInt(proxyUpload),
+                download: previousProxy.download + BigInt(proxyDownload),
+              });
             }
+            if (proxyRows && legacyRows) proxySampleError = 'Dataplane смешал старый и новый proxy traffic contract';
             const knownMacs = new Set(devices.map((device) => device.mac));
             const baselinesByMac = { ...traffic.baselinesByMac };
             const totalsByMac = { ...traffic.totalsByMac };
@@ -546,6 +657,73 @@ export function createDeviceInventoryService({
                 rebaselineMacs.delete(mac);
               }
             }
+            const proxyBaselines = { ...traffic.proxy.baselinesByMac };
+            const proxyTotals = { ...traffic.proxy.totalsByMac };
+            const proxyRebaseline = new Set(traffic.proxy.rebaselineMacs);
+            for (const mac of new Set([
+              ...Object.keys(proxyTotals),
+              ...Object.keys(proxyBaselines),
+              ...proxyRebaseline,
+            ])) {
+              if (knownMacs.has(mac)) continue;
+              delete proxyTotals[mac];
+              delete proxyBaselines[mac];
+              proxyRebaseline.delete(mac);
+            }
+            let proxy = {
+              ...traffic.proxy,
+              lastError: !proxyRebaseline.size && traffic.proxy.lastError === PROXY_RECOVERY_ERROR
+                ? null
+                : traffic.proxy.lastError,
+              baselinesByMac: proxyBaselines,
+              totalsByMac: proxyTotals,
+              rebaselineMacs: [...proxyRebaseline],
+            };
+            if (proxySampleError) {
+              proxy = { ...proxy, lastError: proxySampleError };
+            } else if (proxyRows) {
+              try {
+                const nextProxyBaselines = { ...proxyBaselines };
+                const nextProxyTotals = { ...proxyTotals };
+                const nextProxyRebaseline = new Set(proxyRebaseline);
+                for (const [mac, processTotal] of proxyByMac) {
+                  if (!knownMacs.has(mac)) continue;
+                  const baseline = nextProxyBaselines[mac];
+                  const recovering = nextProxyRebaseline.has(mac);
+                  const sameEpoch = !recovering && baseline?.epoch === trafficResult.epoch;
+                  const baselineUpload = sameEpoch ? BigInt(baseline.uploadBytes) : 0n;
+                  const baselineDownload = sameEpoch ? BigInt(baseline.downloadBytes) : 0n;
+                  if (processTotal.upload < baselineUpload || processTotal.download < baselineDownload) {
+                    throw new Error('Dataplane proxy traffic counter уменьшился внутри одного epoch');
+                  }
+                  const total = nextProxyTotals[mac] || { uploadBytes: '0', downloadBytes: '0' };
+                  nextProxyTotals[mac] = {
+                    uploadBytes: (BigInt(total.uploadBytes)
+                      + (recovering ? 0n : processTotal.upload - baselineUpload)).toString(),
+                    downloadBytes: (BigInt(total.downloadBytes)
+                      + (recovering ? 0n : processTotal.download - baselineDownload)).toString(),
+                    observedAt: trafficResult.observedAt || proxy.lastObservedAt,
+                  };
+                  nextProxyBaselines[mac] = {
+                    epoch: trafficResult.epoch,
+                    uploadBytes: processTotal.upload.toString(),
+                    downloadBytes: processTotal.download.toString(),
+                  };
+                  nextProxyRebaseline.delete(mac);
+                }
+                proxy = {
+                  ...proxy,
+                  lastObservedAt: trafficResult.observedAt || proxy.lastObservedAt,
+                  lastError: trafficResult.source?.error
+                    || (nextProxyRebaseline.size ? proxy.lastError : null),
+                  baselinesByMac: nextProxyBaselines,
+                  totalsByMac: nextProxyTotals,
+                  rebaselineMacs: [...nextProxyRebaseline],
+                };
+              } catch (error) {
+                proxy = { ...proxy, lastError: error.message || String(error) };
+              }
+            }
             traffic = {
               ...traffic,
               epoch: trafficResult.epoch,
@@ -556,6 +734,7 @@ export function createDeviceInventoryService({
               baselinesByMac,
               totalsByMac,
               rebaselineMacs: [...rebaselineMacs],
+              proxy,
             };
           } catch (error) {
             traffic = { ...traffic, lastError: error.message || String(error) };