smartproxy/test/test.nftables-status.ts

import { SmartProxy } from '../ts/proxies/smart-proxy/index.js';
import { NFTablesManager } from '../ts/proxies/smart-proxy/nftables-manager.js';
import { createNfTablesRoute } from '../ts/proxies/smart-proxy/utils/route-helpers.js';
import { expect, tap } from '@git.zone/tstest/tapbundle';
import * as child_process from 'child_process';
import { promisify } from 'util';

const exec = promisify(child_process.exec);

// Check if we have root privileges
async function checkRootPrivileges(): Promise<boolean> {
  try {
    const { stdout } = await exec('id -u');
    return stdout.trim() === '0';
  } catch (err) {
    return false;
  }
}

// Skip tests if not root
const isRoot = await checkRootPrivileges();
if (!isRoot) {
  console.log('');
  console.log('========================================');
  console.log('NFTables status tests require root privileges');
  console.log('Skipping NFTables status tests');
  console.log('========================================');
  console.log('');
}

// Define the test function based on root privileges
const testFn = isRoot ? tap.test : tap.skip.test;

testFn('NFTablesManager status functionality', async () => {
  const nftablesManager = new NFTablesManager({ routes: [] });
  
  // Create test routes
  const testRoutes = [
    createNfTablesRoute('test-route-1', { host: 'localhost', port: 8080 }, { ports: 9080 }),
    createNfTablesRoute('test-route-2', { host: 'localhost', port: 8081 }, { ports: 9081 }),
    createNfTablesRoute('test-route-3', { host: 'localhost', port: 8082 }, { 
      ports: 9082,
      ipAllowList: ['127.0.0.1', '192.168.1.0/24']
    })
  ];
  
  // Get initial status (should be empty)
  let status = await nftablesManager.getStatus();
  expect(Object.keys(status).length).toEqual(0);
  
  // Provision routes
  for (const route of testRoutes) {
    await nftablesManager.provisionRoute(route);
  }
  
  // Get status after provisioning
  status = await nftablesManager.getStatus();
  expect(Object.keys(status).length).toEqual(3);
  
  // Check status structure
  for (const routeStatus of Object.values(status)) {
    expect(routeStatus).toHaveProperty('active');
    expect(routeStatus).toHaveProperty('ruleCount');
    expect(routeStatus).toHaveProperty('lastUpdate');
    expect(routeStatus.active).toBeTrue();
  }
  
  // Deprovision one route
  await nftablesManager.deprovisionRoute(testRoutes[0]);
  
  // Check status after deprovisioning
  status = await nftablesManager.getStatus();
  expect(Object.keys(status).length).toEqual(2);
  
  // Cleanup remaining routes
  await nftablesManager.stop();
  
  // Final status should be empty
  status = await nftablesManager.getStatus();
  expect(Object.keys(status).length).toEqual(0);
});

testFn('SmartProxy getNfTablesStatus functionality', async () => {
  const smartProxy = new SmartProxy({
    routes: [
      createNfTablesRoute('proxy-test-1', { host: 'localhost', port: 3000 }, { ports: 3001 }),
      createNfTablesRoute('proxy-test-2', { host: 'localhost', port: 3002 }, { ports: 3003 }),
      // Include a non-NFTables route to ensure it's not included in the status
      {
        name: 'non-nftables-route',
        match: { ports: 3004 },
        action: {
          type: 'forward',
          target: { host: 'localhost', port: 3005 }
        }
      }
    ]
  });
  
  // Start the proxy
  await smartProxy.start();
  
  // Get NFTables status
  const status = await smartProxy.getNfTablesStatus();
  
  // Should only have 2 NFTables routes
  const statusKeys = Object.keys(status);
  expect(statusKeys.length).toEqual(2);
  
  // Check that both NFTables routes are in the status
  const routeIds = statusKeys.sort();
  expect(routeIds).toContain('proxy-test-1:3001');
  expect(routeIds).toContain('proxy-test-2:3003');
  
  // Verify status structure
  for (const [routeId, routeStatus] of Object.entries(status)) {
    expect(routeStatus).toHaveProperty('active', true);
    expect(routeStatus).toHaveProperty('ruleCount');
    expect(routeStatus.ruleCount).toHaveProperty('total');
    expect(routeStatus.ruleCount.total).toBeGreaterThan(0);
  }
  
  // Stop the proxy
  await smartProxy.stop();
  
  // After stopping, status should be empty
  const finalStatus = await smartProxy.getNfTablesStatus();
  expect(Object.keys(finalStatus).length).toEqual(0);
});

testFn('NFTables route update status tracking', async () => {
  const smartProxy = new SmartProxy({
    routes: [
      createNfTablesRoute('update-test', { host: 'localhost', port: 4000 }, { ports: 4001 })
    ]
  });
  
  await smartProxy.start();
  
  // Get initial status
  let status = await smartProxy.getNfTablesStatus();
  expect(Object.keys(status).length).toEqual(1);
  const initialUpdate = status['update-test:4001'].lastUpdate;
  
  // Wait a moment
  await new Promise(resolve => setTimeout(resolve, 10));
  
  // Update the route
  await smartProxy.updateRoutes([
    createNfTablesRoute('update-test', { host: 'localhost', port: 4002 }, { ports: 4001 })
  ]);
  
  // Get status after update
  status = await smartProxy.getNfTablesStatus();
  expect(Object.keys(status).length).toEqual(1);
  const updatedTime = status['update-test:4001'].lastUpdate;
  
  // The update time should be different
  expect(updatedTime.getTime()).toBeGreaterThan(initialUpdate.getTime());
  
  await smartProxy.stop();
});

export default tap.start();