Improve VPN client connection management

test/server/device-traffic.test.js

@@ -11,6 +11,12 @@ import {
 
 const uploadChain = 'VPN_PROXY_TRAFFIC_UP';
 const downloadChain = 'VPN_PROXY_TRAFFIC_DOWN';
+const routeAccounting = {
+  directChain: 'VPN_PROXY_DIRECT',
+  directMark: '0x40000000',
+  tproxyMark: '1',
+  gatewayClientCidrs: ['10.0.0.0/8', '172.16.0.0/12', '192.168.0.0/16'],
+};
 const dataplaneSource = fs.readFileSync(
   path.resolve(import.meta.dirname, '../../src/server/dataplane.ts'),
   'utf8',
@@ -23,7 +29,7 @@ const observation = (ip, mac = '00:11:22:33:44:55', deviceInterface = 'eth0') =>
 
 test('dataplane exposes cached traffic snapshots without making accounting a readiness dependency', () => {
   assert.match(dataplaneSource, /req\.method === 'GET' && req\.url === '\/device-traffic'[\s\S]*traffic\.snapshot\(\)/);
-  assert.match(dataplaneSource, /ready = true;[\s\S]*setImmediate[\s\S]*traffic\.refresh\(\)/);
+  assert.match(dataplaneSource, /ready = true;[\s\S]*deviceTrafficAccountingEnabled[\s\S]*setImmediate[\s\S]*traffic\.refresh\(\)/);
   assert.match(dataplaneSource, /traffic\.refresh\(\)\.catch/);
 });
 
@@ -56,6 +62,7 @@ test('traffic rules split local proxy traffic from public Gateway traffic in one
     bypassCidrs: ['10.0.0.0/8', '192.168.0.0/16'],
     uploadChain,
     downloadChain,
+    ...routeAccounting,
     slot: 'A',
     proxyPort: 8080,
   });
@@ -76,7 +83,26 @@ test('traffic rules split local proxy traffic from public Gateway traffic in one
     < lines.findIndex((line) => line.endsWith(':upload -j RETURN')));
   assert.ok(lines.indexOf('-A VPN_PROXY_TRAFFIC_DOWN_A -s 10.0.0.0/8 -j RETURN')
     < lines.findIndex((line) => line.endsWith(':download -j RETURN')));
+  assert.match(restore, /-A VPN_PROXY_DIRECT_A -m addrtype --dst-type LOCAL -j RETURN/);
+  assert.match(restore, /-A VPN_PROXY_DIRECT_A -m mark --mark 1\/1 -j RETURN/);
+  assert.match(restore, /-A VPN_PROXY_DIRECT_A -i br-\+ -j RETURN/);
+  assert.match(restore, /-A VPN_PROXY_DIRECT_A -d 10\.0\.0\.0\/8 -j RETURN/);
+  assert.match(restore, /-A VPN_PROXY_DIRECT_A -s 192\.168\.0\.0\/16 .*harbor-traffic:global:direct-upload -j CONNMARK --set-xmark 0x40000000\/0x40000000/);
+  assert.match(restore, /-A VPN_PROXY_TRAFFIC_DOWN_A -d 192\.168\.0\.0\/16 .*harbor-traffic:global:direct-download$/m);
+  assert.doesNotMatch(restore, /-A VPN_PROXY_TRAFFIC_DOWN_A -m connmark .*harbor-traffic:global:direct-download$/m);
+  assert.match(restore, /-A VPN_PROXY_DIRECT_A .*--mark 0x40000000\/0x40000000 .*harbor-traffic:[a-f0-9]{16}:direct-upload$/m);
+  assert.match(restore, /-A VPN_PROXY_TRAFFIC_DOWN_A .*--mark 0x40000000\/0x40000000 .*harbor-traffic:[a-f0-9]{16}:direct-download$/m);
   assert.doesNotMatch(restore, /TPROXY|DNAT|SNAT|REDIRECT/);
+  assert.throws(() => buildTrafficRestore({
+    devices: [],
+    bypassCidrs: [],
+    uploadChain,
+    downloadChain,
+    ...routeAccounting,
+    gatewayClientCidrs: ['10.0.0.0/8', '10.1.0.0/16'],
+    slot: 'A',
+    proxyPort: 8080,
+  }), /traffic accounting/);
 });
 
 test('counter parser preserves exact uint64 strings and sums TCP plus UDP proxy rules', () => {
@@ -86,6 +112,11 @@ test('counter parser preserves exact uint64 strings and sums TCP plus UDP proxy
     'VPN_PROXY_TRAFFIC_UP_A_P',
   );
   assert.equal(counters.get('0123456789abcdef:proxy-upload'), '9007199254741002');
+  const direct = parseTrafficCounters(
+    '[5:77] -A VPN_PROXY_TRAFFIC_DOWN_A -m connmark --mark 0x40000000/0x40000000 -m comment --comment "harbor-traffic:0123456789abcdef:direct-download"\n',
+    'VPN_PROXY_TRAFFIC_DOWN_A',
+  );
+  assert.equal(direct.get('0123456789abcdef:direct-download'), '77');
 });
 
 test('traffic service preserves active rules and snapshot when replacement fails', async () => {
@@ -111,11 +142,18 @@ test('traffic service preserves active rules and snapshot when replacement fails
       return {
         status: 0,
         stdout: `[1:${bytes}] -A ${chain} -m comment --comment "harbor-traffic:${firstDevice.key}:${direction}" -j RETURN\n`
-          + proxyBytes.map((value) => `[1:${value}] -A ${chain}_P -m comment --comment "harbor-traffic:${firstDevice.key}:${proxyDirection}" -j RETURN`).join('\n'),
+          + proxyBytes.map((value) => `[1:${value}] -A ${chain}_P -m comment --comment "harbor-traffic:${firstDevice.key}:${proxyDirection}" -j RETURN`).join('\n')
+          + (direction === 'download'
+            ? `\n[1:300] -A ${routeAccounting.directChain}_A -m comment --comment "harbor-traffic:${firstDevice.key}:direct-upload" -j CONNMARK\n`
+              + `[1:500] -A ${routeAccounting.directChain}_A -m comment --comment "harbor-traffic:global:direct-upload" -j CONNMARK\n`
+              + `[1:400] -A ${chain} -m comment --comment "harbor-traffic:${firstDevice.key}:direct-download"\n`
+              + `[1:600] -A ${chain} -m comment --comment "harbor-traffic:global:direct-download"\n`
+            : ''),
         stderr: '',
       };
     }
-    if (failReplacement && args.includes('-R') && args.includes(downloadChain)) {
+    if (failReplacement && command === 'iptables-restore'
+      && options.input.includes(`-R ${downloadChain}`)) {
       return { status: 1, stdout: '', stderr: 'cannot switch download rules' };
     }
     return { status: 0, stdout: '', stderr: '' };
@@ -126,6 +164,7 @@ test('traffic service preserves active rules and snapshot when replacement fails
     downloadChain,
     bypassCidrs: ['10.0.0.0/8'],
     proxyPort: 8080,
+    ...routeAccounting,
     run,
     nextGeneration: () => generations.shift(),
   });
@@ -133,6 +172,7 @@ test('traffic service preserves active rules and snapshot when replacement fails
   const first = await service.refresh();
   assert.equal(first.epoch, 'boot');
   assert.equal(first.generation, 'rules-a');
+  assert.deepEqual(first.direct, { uploadBytes: '500', downloadBytes: '600' });
   assert.deepEqual(first.devices, [{
     ip: '192.168.50.7',
     mac: '00:11:22:33:44:55',
@@ -141,6 +181,8 @@ test('traffic service preserves active rules and snapshot when replacement fails
     downloadBytes: '3400',
     proxyUploadBytes: '110',
     proxyDownloadBytes: '220',
+    directUploadBytes: '300',
+    directDownloadBytes: '400',
   }]);
 
   const switchCallsBefore = calls.filter(([, args]) => args.includes('-R') || args.includes('-A')).length;
@@ -162,6 +204,11 @@ test('traffic service preserves active rules and snapshot when replacement fails
     .filter(([, args]) => args.includes('-R') && args.includes(uploadChain))
     .map(([, args]) => args.at(-1));
   assert.deepEqual(uploadSwitches, [`${uploadChain}_B`, `${uploadChain}_A`]);
+  const failedMangleSwitch = calls.find(([command, , options]) => (
+    command === 'iptables-restore' && options.input.includes(`-R ${downloadChain}`)
+  ));
+  assert.match(failedMangleSwitch[2].input, new RegExp(`-R ${downloadChain} 1 -j ${downloadChain}_B`));
+  assert.match(failedMangleSwitch[2].input, new RegExp(`-R ${routeAccounting.directChain} 1 -j ${routeAccounting.directChain}_B`));
 
   failReplacement = false;
   observed = {
@@ -187,6 +234,32 @@ test('traffic service preserves active rules and snapshot when replacement fails
   assert.ok(calls.every(([, , options]) => options.timeout === 2_000));
 });
 
+test('an initial accounting failure does not publish fresh zero counters', async () => {
+  const service = createDeviceTrafficService({
+    observe: () => ({
+      observedAt: '2026-08-07T12:00:00.000Z',
+      observations: [observation('192.168.50.7')],
+      error: null,
+    }),
+    uploadChain,
+    downloadChain,
+    bypassCidrs: [],
+    proxyPort: 8080,
+    ...routeAccounting,
+    run: (command) => ({
+      status: command === 'iptables-restore' ? 1 : 0,
+      stdout: '',
+      stderr: command === 'iptables-restore' ? 'prepare failed' : '',
+    }),
+  });
+
+  const snapshot = await service.refresh();
+  assert.equal(snapshot.observedAt, null);
+  assert.deepEqual(snapshot.direct, { uploadBytes: '0', downloadBytes: '0' });
+  assert.deepEqual(snapshot.devices, []);
+  assert.match(snapshot.source.error, /prepare failed/);
+});
+
 test('a 512-device refresh keeps a fixed subprocess count and a cached snapshot', async () => {
   const observations = Array.from({ length: 512 }, (_, index) => observation(
     `10.${Math.floor(index / 254)}.${Math.floor((index % 254) / 254)}.${(index % 254) + 1}`,
@@ -194,9 +267,11 @@ test('a 512-device refresh keeps a fixed subprocess count and a cached snapshot'
   ));
   const calls = [];
   let releaseRestore;
+  let blockedRestore = false;
   const run = (command, args, options) => {
     calls.push([command, args, options]);
-    if (command === 'iptables-restore') {
+    if (command === 'iptables-restore' && !blockedRestore) {
+      blockedRestore = true;
       return new Promise((resolve) => {
         releaseRestore = () => resolve({ status: 0, stdout: '', stderr: '' });
       });
@@ -210,6 +285,7 @@ test('a 512-device refresh keeps a fixed subprocess count and a cached snapshot'
     downloadChain,
     bypassCidrs: [],
     proxyPort: 8080,
+    ...routeAccounting,
     run,
     nextGeneration: () => generations.shift(),
   });
@@ -225,8 +301,10 @@ test('a 512-device refresh keeps a fixed subprocess count and a cached snapshot'
   assert.equal(snapshot.devices.length, 512);
   assert.deepEqual(
     calls.map(([command]) => command),
-    ['iptables-restore', 'iptables', 'iptables', 'iptables-save', 'iptables-save'],
+    ['iptables-restore', 'iptables', 'iptables-restore', 'iptables-save', 'iptables-save'],
   );
+  assert.match(calls[2][2].input, new RegExp(`-A ${downloadChain} -j ${downloadChain}_A`));
+  assert.match(calls[2][2].input, new RegExp(`-A ${routeAccounting.directChain} -j ${routeAccounting.directChain}_A`));
 });
 
 test('traffic service finalizes a detached slot once and keeps epoch totals monotonic', async () => {
@@ -240,8 +318,14 @@ test('traffic service finalizes a detached slot once and keeps epoch totals mono
     error: null,
   };
   const values = {
-    A: { upload: '100', download: '200', proxyUpload: '30', proxyDownload: '40' },
-    B: { upload: '5', download: '7', proxyUpload: '2', proxyDownload: '3' },
+    A: {
+      upload: '100', download: '200', proxyUpload: '30', proxyDownload: '40',
+      directUpload: '11', directDownload: '22', globalDirectUpload: '33', globalDirectDownload: '44',
+    },
+    B: {
+      upload: '5', download: '7', proxyUpload: '2', proxyDownload: '3',
+      directUpload: '1', directDownload: '2', globalDirectUpload: '3', globalDirectDownload: '4',
+    },
   };
   const keys = { A: firstDevice.key, B: secondDevice.key };
   let failNextCounterRead = false;
@@ -260,6 +344,12 @@ test('traffic service finalizes a detached slot once and keeps epoch totals mono
       stdout: ['A', 'B'].flatMap((slot) => [
         `[1:${values[slot][direction]}] -A ${tableChain}_${slot} -m comment --comment "harbor-traffic:${keys[slot]}:${direction}" -j RETURN`,
         `[1:${values[slot][proxyDirection]}] -A ${tableChain}_${slot}_P -m comment --comment "harbor-traffic:${keys[slot]}:${proxyKind}" -j RETURN`,
+        ...(direction === 'download' ? [
+          `[1:${values[slot].directUpload}] -A ${routeAccounting.directChain}_${slot} -m comment --comment "harbor-traffic:${keys[slot]}:direct-upload"`,
+          `[1:${values[slot].directDownload}] -A ${tableChain}_${slot} -m comment --comment "harbor-traffic:${keys[slot]}:direct-download"`,
+          `[1:${values[slot].globalDirectUpload}] -A ${routeAccounting.directChain}_${slot} -m comment --comment "harbor-traffic:global:direct-upload"`,
+          `[1:${values[slot].globalDirectDownload}] -A ${tableChain}_${slot} -m comment --comment "harbor-traffic:global:direct-download"`,
+        ] : []),
       ]).join('\n'),
       stderr: '',
     };
@@ -271,6 +361,7 @@ test('traffic service finalizes a detached slot once and keeps epoch totals mono
     downloadChain,
     bypassCidrs: [],
     proxyPort: 8080,
+    ...routeAccounting,
     run,
     nextGeneration: () => generations.shift(),
   });
@@ -278,17 +369,23 @@ test('traffic service finalizes a detached slot once and keeps epoch totals mono
   const first = await service.refresh();
   assert.equal(first.epoch, 'epoch-1');
   assert.equal(first.generation, 'rules-a');
-  assert.deepEqual(first.devices.map(({ mac, uploadBytes, downloadBytes, proxyUploadBytes, proxyDownloadBytes }) => ({
-    mac, uploadBytes, downloadBytes, proxyUploadBytes, proxyDownloadBytes,
+  assert.deepEqual(first.direct, { uploadBytes: '33', downloadBytes: '44' });
+  assert.deepEqual(first.devices.map(({ mac, uploadBytes, downloadBytes, proxyUploadBytes, proxyDownloadBytes, directUploadBytes, directDownloadBytes }) => ({
+    mac, uploadBytes, downloadBytes, proxyUploadBytes, proxyDownloadBytes, directUploadBytes, directDownloadBytes,
   })), [{
     mac: firstObservation.mac,
     uploadBytes: '100',
     downloadBytes: '200',
     proxyUploadBytes: '30',
     proxyDownloadBytes: '40',
+    directUploadBytes: '11',
+    directDownloadBytes: '22',
   }]);
 
-  values.A = { upload: '130', download: '240', proxyUpload: '35', proxyDownload: '48' };
+  values.A = {
+    upload: '130', download: '240', proxyUpload: '35', proxyDownload: '48',
+    directUpload: '15', directDownload: '28', globalDirectUpload: '39', globalDirectDownload: '52',
+  };
   observed = {
     observedAt: '2026-08-07T12:01:00.000Z',
     observations: [secondObservation],
@@ -299,29 +396,35 @@ test('traffic service finalizes a detached slot once and keeps epoch totals mono
   assert.equal(pending.epoch, 'epoch-1');
   assert.equal(pending.generation, 'rules-b');
   assert.match(pending.source.error, /retired slot read failed/);
-  assert.deepEqual(pending.devices.map(({ mac, uploadBytes, downloadBytes, proxyUploadBytes, proxyDownloadBytes }) => ({
-    mac, uploadBytes, downloadBytes, proxyUploadBytes, proxyDownloadBytes,
+  assert.deepEqual(pending.direct, { uploadBytes: '36', downloadBytes: '48' });
+  assert.deepEqual(pending.devices.map(({ mac, uploadBytes, downloadBytes, proxyUploadBytes, proxyDownloadBytes, directUploadBytes, directDownloadBytes }) => ({
+    mac, uploadBytes, downloadBytes, proxyUploadBytes, proxyDownloadBytes, directUploadBytes, directDownloadBytes,
   })), [
-    { mac: firstObservation.mac, uploadBytes: '100', downloadBytes: '200', proxyUploadBytes: '30', proxyDownloadBytes: '40' },
-    { mac: secondObservation.mac, uploadBytes: '5', downloadBytes: '7', proxyUploadBytes: '2', proxyDownloadBytes: '3' },
+    { mac: firstObservation.mac, uploadBytes: '100', downloadBytes: '200', proxyUploadBytes: '30', proxyDownloadBytes: '40', directUploadBytes: '11', directDownloadBytes: '22' },
+    { mac: secondObservation.mac, uploadBytes: '5', downloadBytes: '7', proxyUploadBytes: '2', proxyDownloadBytes: '3', directUploadBytes: '1', directDownloadBytes: '2' },
   ]);
 
   const finalized = await service.refresh();
   assert.equal(finalized.generation, 'rules-b');
   assert.equal(finalized.source.error, null);
-  assert.deepEqual(finalized.devices.map(({ mac, uploadBytes, downloadBytes, proxyUploadBytes, proxyDownloadBytes }) => ({
-    mac, uploadBytes, downloadBytes, proxyUploadBytes, proxyDownloadBytes,
+  assert.deepEqual(finalized.direct, { uploadBytes: '42', downloadBytes: '56' });
+  assert.deepEqual(finalized.devices.map(({ mac, uploadBytes, downloadBytes, proxyUploadBytes, proxyDownloadBytes, directUploadBytes, directDownloadBytes }) => ({
+    mac, uploadBytes, downloadBytes, proxyUploadBytes, proxyDownloadBytes, directUploadBytes, directDownloadBytes,
   })), [
-    { mac: firstObservation.mac, uploadBytes: '130', downloadBytes: '240', proxyUploadBytes: '35', proxyDownloadBytes: '48' },
-    { mac: secondObservation.mac, uploadBytes: '5', downloadBytes: '7', proxyUploadBytes: '2', proxyDownloadBytes: '3' },
+    { mac: firstObservation.mac, uploadBytes: '130', downloadBytes: '240', proxyUploadBytes: '35', proxyDownloadBytes: '48', directUploadBytes: '15', directDownloadBytes: '28' },
+    { mac: secondObservation.mac, uploadBytes: '5', downloadBytes: '7', proxyUploadBytes: '2', proxyDownloadBytes: '3', directUploadBytes: '1', directDownloadBytes: '2' },
   ]);
 
-  values.B = { upload: '15', download: '17', proxyUpload: '4', proxyDownload: '6' };
+  values.B = {
+    upload: '15', download: '17', proxyUpload: '4', proxyDownload: '6',
+    directUpload: '4', directDownload: '6', globalDirectUpload: '8', globalDirectDownload: '10',
+  };
   const polled = await service.refresh();
-  assert.deepEqual(polled.devices.map(({ mac, uploadBytes, downloadBytes, proxyUploadBytes, proxyDownloadBytes }) => ({
-    mac, uploadBytes, downloadBytes, proxyUploadBytes, proxyDownloadBytes,
+  assert.deepEqual(polled.direct, { uploadBytes: '47', downloadBytes: '62' });
+  assert.deepEqual(polled.devices.map(({ mac, uploadBytes, downloadBytes, proxyUploadBytes, proxyDownloadBytes, directUploadBytes, directDownloadBytes }) => ({
+    mac, uploadBytes, downloadBytes, proxyUploadBytes, proxyDownloadBytes, directUploadBytes, directDownloadBytes,
   })), [
-    { mac: firstObservation.mac, uploadBytes: '130', downloadBytes: '240', proxyUploadBytes: '35', proxyDownloadBytes: '48' },
-    { mac: secondObservation.mac, uploadBytes: '15', downloadBytes: '17', proxyUploadBytes: '4', proxyDownloadBytes: '6' },
+    { mac: firstObservation.mac, uploadBytes: '130', downloadBytes: '240', proxyUploadBytes: '35', proxyDownloadBytes: '48', directUploadBytes: '15', directDownloadBytes: '28' },
+    { mac: secondObservation.mac, uploadBytes: '15', downloadBytes: '17', proxyUploadBytes: '4', proxyDownloadBytes: '6', directUploadBytes: '4', directDownloadBytes: '6' },
   ]);
 });