Improve VPN client connection management

src/server/services/deviceTrafficService.ts

@@ -24,9 +24,9 @@ interface TrafficDevice {
   key: string;
 }
 
-type CounterKind = 'upload' | 'download' | 'proxy-upload' | 'proxy-download';
-type CounterField = 'upload' | 'download' | 'proxyUpload' | 'proxyDownload';
-type CounterOutput = 'uploadBytes' | 'downloadBytes' | 'proxyUploadBytes' | 'proxyDownloadBytes';
+type CounterKind = 'upload' | 'download' | 'proxy-upload' | 'proxy-download' | 'direct-upload' | 'direct-download';
+type CounterField = 'upload' | 'download' | 'proxyUpload' | 'proxyDownload' | 'directUpload' | 'directDownload';
+type CounterOutput = 'uploadBytes' | 'downloadBytes' | 'proxyUploadBytes' | 'proxyDownloadBytes' | 'directUploadBytes' | 'directDownloadBytes';
 type CounterValues = Record<CounterField, bigint>;
 
 interface RetiredCounters {
@@ -40,6 +40,7 @@ interface TrafficSnapshot {
   generation: string;
   observedAt: string | null;
   source: { error: string | null };
+  direct: { uploadBytes: string; downloadBytes: string };
   devices: Record<string, unknown>[];
 }
 
@@ -53,6 +54,8 @@ const COUNTERS = [
   ['download', 'download', 'downloadBytes'],
   ['proxy-upload', 'proxyUpload', 'proxyUploadBytes'],
   ['proxy-download', 'proxyDownload', 'proxyDownloadBytes'],
+  ['direct-upload', 'directUpload', 'directUploadBytes'],
+  ['direct-download', 'directDownload', 'directDownloadBytes'],
 ] as const satisfies readonly (readonly [CounterKind, CounterField, CounterOutput])[];
 
 const childChain = (chain: string, slot: string) => `${chain}_${slot}`;
@@ -117,7 +120,37 @@ function isIpv4Cidr(value: unknown) {
     && Number.isInteger(size) && size >= 0 && size <= 32;
 }
 
-const zeroCounters = (): CounterValues => ({ upload: 0n, download: 0n, proxyUpload: 0n, proxyDownload: 0n });
+function cidrRange(cidr: string) {
+  const [address, prefix] = cidr.split('/');
+  const value = address.split('.').reduce((result, octet) => result * 256n + BigInt(octet), 0n);
+  const bits = BigInt(Number(prefix));
+  const mask = bits === 0n ? 0n : (0xffff_ffffn << (32n - bits)) & 0xffff_ffffn;
+  const first = value & mask;
+  return [first, first | (0xffff_ffffn ^ mask)] as const;
+}
+
+function hasOverlappingCidrs(cidrs: readonly string[]) {
+  const ranges = cidrs.map(cidrRange).sort(([left], [right]) => (left < right ? -1 : left > right ? 1 : 0));
+  return ranges.some(([start], index) => index > 0 && start <= ranges[index - 1][1]);
+}
+
+const zeroCounters = (): CounterValues => ({
+  upload: 0n,
+  download: 0n,
+  proxyUpload: 0n,
+  proxyDownload: 0n,
+  directUpload: 0n,
+  directDownload: 0n,
+});
+
+function markValue(value: unknown) {
+  try {
+    const parsed = BigInt(String(value));
+    return parsed > 0n && parsed <= 0xffff_ffffn ? parsed : null;
+  } catch {
+    return null;
+  }
+}
 
 function record(value: unknown): Record<string, unknown> {
   return value && typeof value === 'object' && !Array.isArray(value)
@@ -160,6 +193,10 @@ export function buildTrafficRestore({
   bypassCidrs,
   uploadChain,
   downloadChain,
+  directChain,
+  directMark,
+  tproxyMark,
+  gatewayClientCidrs,
   slot,
   proxyPort,
 }: {
@@ -167,17 +204,28 @@ export function buildTrafficRestore({
   bypassCidrs: readonly string[];
   uploadChain: string;
   downloadChain: string;
+  directChain: string;
+  directMark: string;
+  tproxyMark: string;
+  gatewayClientCidrs: readonly string[];
   slot: string;
   proxyPort: number;
 }) {
-  if (!CHAIN_PATTERN.test(uploadChain) || !CHAIN_PATTERN.test(downloadChain)
+  const parsedDirectMark = markValue(directMark);
+  const parsedTproxyMark = markValue(tproxyMark);
+  if (!CHAIN_PATTERN.test(uploadChain) || !CHAIN_PATTERN.test(downloadChain) || !CHAIN_PATTERN.test(directChain)
     || !['A', 'B'].includes(slot) || !Number.isInteger(proxyPort)
     || proxyPort < 1 || proxyPort > 65_535
-    || !Array.isArray(bypassCidrs) || bypassCidrs.some((cidr) => !isIpv4Cidr(cidr))) {
+    || parsedDirectMark == null || (parsedDirectMark & (parsedDirectMark - 1n)) !== 0n
+    || parsedTproxyMark == null || (parsedDirectMark & parsedTproxyMark) !== 0n
+    || !Array.isArray(bypassCidrs) || bypassCidrs.some((cidr) => !isIpv4Cidr(cidr))
+    || !Array.isArray(gatewayClientCidrs) || gatewayClientCidrs.some((cidr) => !isIpv4Cidr(cidr))
+    || hasOverlappingCidrs(gatewayClientCidrs)) {
     throw new Error('Некорректная конфигурация traffic accounting');
   }
   const uploadChild = childChain(uploadChain, slot);
   const downloadChild = childChain(downloadChain, slot);
+  const directChild = childChain(directChain, slot);
   const proxyUploadChild = proxyChildChain(uploadChain, slot);
   const proxyDownloadChild = proxyChildChain(downloadChain, slot);
   const raw = [
@@ -192,8 +240,12 @@ export function buildTrafficRestore({
   ];
   const mangle = [
     '*mangle',
+    `-F ${directChild}`,
     `-F ${downloadChild}`,
     `-F ${proxyDownloadChild}`,
+    `-A ${directChild} -m addrtype --dst-type LOCAL -j RETURN`,
+    `-A ${directChild} -m mark --mark ${tproxyMark}/${tproxyMark} -j RETURN`,
+    `-A ${directChild} -i br-+ -j RETURN`,
     `-A ${downloadChild} -p tcp --sport ${proxyPort} -m addrtype --src-type LOCAL -j ${proxyDownloadChild}`,
     `-A ${downloadChild} -p tcp --sport ${proxyPort} -m addrtype --src-type LOCAL -j RETURN`,
     `-A ${downloadChild} -p udp --sport ${proxyPort} -m addrtype --src-type LOCAL -j ${proxyDownloadChild}`,
@@ -208,9 +260,16 @@ export function buildTrafficRestore({
   }
   for (const cidr of bypassCidrs) {
     raw.push(`-A ${uploadChild} -d ${cidr} -j RETURN`);
+    mangle.push(`-A ${directChild} -d ${cidr} -j RETURN`);
     mangle.push(`-A ${downloadChild} -s ${cidr} -j RETURN`);
   }
+  for (const cidr of gatewayClientCidrs) {
+    mangle.push(`-A ${directChild} -s ${cidr} -m comment --comment harbor-traffic:global:direct-upload -j CONNMARK --set-xmark ${directMark}/${directMark}`);
+    mangle.push(`-A ${downloadChild} -d ${cidr} -m connmark --mark ${directMark}/${directMark} -m comment --comment harbor-traffic:global:direct-download`);
+  }
   for (const device of devices) {
+    mangle.push(`-A ${directChild} -i ${device.interface} -s ${device.ip} -m mac --mac-source ${device.mac} -m connmark --mark ${directMark}/${directMark} -m comment --comment harbor-traffic:${device.key}:direct-upload`);
+    mangle.push(`-A ${downloadChild} -o ${device.interface} -d ${device.ip} -m connmark --mark ${directMark}/${directMark} -m comment --comment harbor-traffic:${device.key}:direct-download`);
     raw.push(`-A ${uploadChild} -i ${device.interface} -s ${device.ip} -m mac --mac-source ${device.mac} -m comment --comment harbor-traffic:${device.key}:upload -j RETURN`);
     mangle.push(`-A ${downloadChild} -o ${device.interface} -d ${device.ip} -m comment --comment harbor-traffic:${device.key}:download -j RETURN`);
   }
@@ -220,7 +279,7 @@ export function buildTrafficRestore({
 export function parseTrafficCounters(text: unknown, chain: string): Map<string, string> {
   const escapedChain = chain.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
   const linePattern = new RegExp(
-    `^\\[(\\d+):(\\d+)\\] -A ${escapedChain} .*--comment "?harbor-traffic:([a-f0-9]{16}):(upload|download|proxy-upload|proxy-download)"?`,
+    `^\\[(\\d+):(\\d+)\\] -A ${escapedChain} .*--comment "?harbor-traffic:([a-f0-9]{16}|global):(upload|download|proxy-upload|proxy-download|direct-upload|direct-download)"?`,
   );
   const counters = new Map<string, string>();
   for (const line of String(text || '').split(/\r?\n/)) {
@@ -236,6 +295,10 @@ export function createDeviceTrafficService({
   observe,
   uploadChain,
   downloadChain,
+  directChain,
+  directMark,
+  tproxyMark,
+  gatewayClientCidrs,
   bypassCidrs,
   proxyPort,
   run = runCommand,
@@ -244,6 +307,10 @@ export function createDeviceTrafficService({
   observe: () => Promise<unknown> | unknown;
   uploadChain: string;
   downloadChain: string;
+  directChain: string;
+  directMark: string;
+  tproxyMark: string;
+  gatewayClientCidrs: string[];
   bypassCidrs: string[];
   proxyPort: number;
   run?: RunCommand;
@@ -257,12 +324,14 @@ export function createDeviceTrafficService({
   let pendingRetired: RetiredCounters | null = null;
   let refreshPromise: Promise<TrafficSnapshot> | null = null;
   const finalized = new Map<string, CounterValues>();
+  const finalizedDirect = { upload: 0n, download: 0n };
   const devicesByKey = new Map<string, TrafficDevice>();
   let current: TrafficSnapshot = {
     epoch,
     generation: epoch,
     observedAt: null,
     source: { error: null },
+    direct: { uploadBytes: '0', downloadBytes: '0' },
     devices: [],
   };
 
@@ -278,6 +347,10 @@ export function createDeviceTrafficService({
       bypassCidrs,
       uploadChain,
       downloadChain,
+      directChain,
+      directMark,
+      tproxyMark,
+      gatewayClientCidrs,
       slot,
       proxyPort,
     });
@@ -287,22 +360,37 @@ export function createDeviceTrafficService({
   async function switchTo(slot: 'A' | 'B') {
     const uploadChild = childChain(uploadChain, slot);
     const downloadChild = childChain(downloadChain, slot);
+    const directChild = childChain(directChain, slot);
     const replace = activeSlot ? '-R' : '-A';
     const uploadArgs = activeSlot
       ? ['-w', '1', '-t', 'raw', replace, uploadChain, '1', '-j', uploadChild]
       : ['-w', '1', '-t', 'raw', replace, uploadChain, '-j', uploadChild];
-    const downloadArgs = activeSlot
-      ? ['-w', '1', '-t', 'mangle', replace, downloadChain, '1', '-j', downloadChild]
-      : ['-w', '1', '-t', 'mangle', replace, downloadChain, '-j', downloadChild];
+    const mangleInput = [
+      '*mangle',
+      activeSlot
+        ? `-R ${downloadChain} 1 -j ${downloadChild}`
+        : `-A ${downloadChain} -j ${downloadChild}`,
+      activeSlot
+        ? `-R ${directChain} 1 -j ${directChild}`
+        : `-A ${directChain} -j ${directChild}`,
+      'COMMIT',
+      '',
+    ].join('\n');
 
     await execute('iptables', uploadArgs);
     try {
-      await execute('iptables', downloadArgs);
+      await execute('iptables-restore', ['-w', '1', '--noflush'], { ...COMMAND_OPTIONS, input: mangleInput });
     } catch (error) {
-      const rollbackArgs = activeSlot
+      const uploadRollback = activeSlot
         ? ['-w', '1', '-t', 'raw', '-R', uploadChain, '1', '-j', childChain(uploadChain, activeSlot)]
         : ['-w', '1', '-t', 'raw', '-F', uploadChain];
-      await execute('iptables', rollbackArgs);
+      try {
+        await execute('iptables', uploadRollback);
+      } catch (rollbackError) {
+        const original = error instanceof Error ? error.message : String(error);
+        const rollback = rollbackError instanceof Error ? rollbackError.message : String(rollbackError);
+        throw new Error(`${original}; rollback: ${rollback}`, { cause: error });
+      }
       throw error;
     }
   }
@@ -318,6 +406,8 @@ export function createDeviceTrafficService({
       download: parseTrafficCounters(mangle, childChain(downloadChain, slot)),
       proxyUpload: parseTrafficCounters(raw, proxyChildChain(uploadChain, slot)),
       proxyDownload: parseTrafficCounters(mangle, proxyChildChain(downloadChain, slot)),
+      directUpload: parseTrafficCounters(mangle, childChain(directChain, slot)),
+      directDownload: parseTrafficCounters(mangle, childChain(downloadChain, slot)),
     };
     const counters = new Map<string, string>();
     for (const { key } of devices) {
@@ -325,6 +415,8 @@ export function createDeviceTrafficService({
         counters.set(`${key}:${kind}`, parsed[field].get(`${key}:${kind}`) || '0');
       }
     }
+    counters.set('global:direct-upload', parsed.directUpload.get('global:direct-upload') || '0');
+    counters.set('global:direct-download', parsed.directDownload.get('global:direct-download') || '0');
     return counters;
   }
 
@@ -345,6 +437,8 @@ export function createDeviceTrafficService({
       for (const [kind, field] of COUNTERS) next[field] += counter(counters, key, kind);
       finalized.set(key, next);
     }
+    finalizedDirect.upload += counter(counters, 'global', 'direct-upload');
+    finalizedDirect.download += counter(counters, 'global', 'direct-download');
     pendingRetired = null;
     return true;
   }
@@ -368,18 +462,41 @@ export function createDeviceTrafficService({
     }
     return [...totalsByMac.values()]
       .map((total) => {
-        const { key: _key, upload, download, proxyUpload, proxyDownload, ...device } = total;
+        const {
+          key: _key,
+          upload,
+          download,
+          proxyUpload,
+          proxyDownload,
+          directUpload,
+          directDownload,
+          ...device
+        } = total;
         return {
           ...device,
           uploadBytes: upload.toString(),
           downloadBytes: download.toString(),
           proxyUploadBytes: proxyUpload.toString(),
           proxyDownloadBytes: proxyDownload.toString(),
+          directUploadBytes: directUpload.toString(),
+          directDownloadBytes: directDownload.toString(),
         };
       })
       .sort((left, right) => left.mac.localeCompare(right.mac));
   }
 
+  function directTotals() {
+    const pending = pendingRetired?.counters || new Map();
+    return {
+      uploadBytes: (finalizedDirect.upload
+        + counter(pending, 'global', 'direct-upload')
+        + counter(activeCounters, 'global', 'direct-upload')).toString(),
+      downloadBytes: (finalizedDirect.download
+        + counter(pending, 'global', 'direct-download')
+        + counter(activeCounters, 'global', 'direct-download')).toString(),
+    };
+  }
+
   async function performRefresh() {
     let observed: Record<string, unknown>;
     try {
@@ -426,11 +543,13 @@ export function createDeviceTrafficService({
       }
     }
 
-    try {
-      activeCounters = await readCounters(activeDevices, activeSlot);
-      countersRead = true;
-    } catch (error) {
-      sourceError = sourceError || (error instanceof Error ? error.message : String(error));
+    if (activeSlot) {
+      try {
+        activeCounters = await readCounters(activeDevices, activeSlot);
+        countersRead = true;
+      } catch (error) {
+        sourceError = sourceError || (error instanceof Error ? error.message : String(error));
+      }
     }
     current = {
       epoch,
@@ -439,6 +558,7 @@ export function createDeviceTrafficService({
         ? observed.observedAt
         : current.observedAt,
       source: { error: sourceError },
+      direct: countersRead ? directTotals() : current.direct,
       devices: countersRead ? processTotals() : current.devices,
     };
     return structuredClone(current);