test-websocket-security.js

/**
 * Test script to verify the improved WebSocket security
 * Tests token validation and agentId verification
 */

import WebSocket from 'ws';

const SERVER_URL = 'ws://localhost:5000';
const WS_ENDPOINT = '/api/ws/agents';

console.log('Testing WebSocket security improvements...\n');

// Test 1: Connection without token (should be rejected)
function testNoToken() {
  return new Promise((resolve) => {
    console.log('Test 1: Connection without token');
    const wsUrl = `${SERVER_URL}${WS_ENDPOINT}`;
    const ws = new WebSocket(wsUrl);
    
    ws.on('open', () => {
      console.log('❌ Connection should have been rejected');
      ws.close();
      resolve(false);
    });
    
    ws.on('close', (code, reason) => {
      if (code === 1008 && reason.toString() === 'Authentication token required') {
        console.log('✓ Correctly rejected connection without token\n');
        resolve(true);
      } else {
        console.log(`❌ Unexpected close: ${code} ${reason}\n`);
        resolve(false);
      }
    });
    
    ws.on('error', (error) => {
      console.log('✓ Connection rejected (error expected)\n');
      resolve(true);
    });
  });
}

// Test 2: Connection with invalid token (should be rejected)
function testInvalidToken() {
  return new Promise((resolve) => {
    console.log('Test 2: Connection with invalid token');
    const wsUrl = `${SERVER_URL}${WS_ENDPOINT}?token=invalid-token-123`;
    const ws = new WebSocket(wsUrl);
    
    ws.on('open', () => {
      console.log('❌ Connection should have been rejected');
      ws.close();
      resolve(false);
    });
    
    ws.on('close', (code, reason) => {
      if (code === 1008 && (reason.toString().includes('Invalid') || reason.toString().includes('expired'))) {
        console.log('✓ Correctly rejected connection with invalid token\n');
        resolve(true);
      } else {
        console.log(`❌ Unexpected close: ${code} ${reason}\n`);
        resolve(false);
      }
    });
    
    ws.on('error', (error) => {
      console.log('✓ Connection rejected (error expected)\n');
      resolve(true);
    });
  });
}

// Test 3: AgentId spoofing protection (with valid token if available)
function testAgentIdSpoofing() {
  return new Promise((resolve) => {
    console.log('Test 3: AgentId spoofing protection');
    // This test would require a valid token, but we're just testing the structure
    const wsUrl = `${SERVER_URL}${WS_ENDPOINT}?token=test-agent-token`;
    const ws = new WebSocket(wsUrl);
    
    ws.on('open', () => {
      console.log('✓ Connection established');
      
      // Try to send a message with a different agentId than what the token represents
      const spoofMessage = {
        type: 'heartbeat',
        agentId: 'malicious-agent-456', // Different from token
        timestamp: new Date().toISOString(),
        status: 'active'
      };
      
      ws.send(JSON.stringify(spoofMessage));
      
      setTimeout(() => {
        ws.close();
        resolve(true);
      }, 1000);
    });
    
    ws.on('message', (data) => {
      try {
        const message = JSON.parse(data.toString());
        console.log('Received:', message);
        
        if (message.type === 'error' && message.message.includes('AgentId mismatch')) {
          console.log('✓ AgentId spoofing correctly detected and blocked\n');
        } else if (message.type === 'welcome') {
          console.log('✓ Welcome message received');
        }
      } catch (error) {
        console.log('Received non-JSON message:', data.toString());
      }
    });
    
    ws.on('close', (code, reason) => {
      console.log(`Connection closed: ${code} ${reason}\n`);
      resolve(true);
    });
    
    ws.on('error', (error) => {
      console.log('Connection error (might be expected for invalid token)');
      resolve(true);
    });
  });
}

// Run all tests
async function runTests() {
  console.log('=== WebSocket Security Test Suite ===\n');
  
  const results = [];
  
  try {
    results.push(await testNoToken());
    results.push(await testInvalidToken());
    results.push(await testAgentIdSpoofing());
    
    const passed = results.filter(r => r).length;
    console.log(`\n=== Results: ${passed}/${results.length} tests passed ===`);
    
    if (passed === results.length) {
      console.log('✓ All security tests passed!');
    } else {
      console.log('❌ Some tests failed');
    }
  } catch (error) {
    console.error('Test suite error:', error);
  }
}

// Run tests
runTests();