tunnel/source/server/routes/agent.ts

import { IncomingMessage } from 'http';
import { WebSocket } from 'ws';
import jwt from 'jsonwebtoken';
import { AuthPayload } from '../middleware/requireAuth';
import { getTunnel } from '../db';
import { registry } from '../relay/TunnelRegistry';
import { dispatchAgentMsg } from '../relay/pending';

export function handleAgentUpgrade( req: IncomingMessage, ws: WebSocket ): void
{
  const match = req.url?.match( /^\/api\/agent\/([^/?]+)/ );
  if ( !match ) { ws.close( 1008, 'Bad URL' ); return; }
  const tunnelId = match[ 1 ];

  const url = new URL( req.url!, 'http://localhost' );
  const token = url.searchParams.get( 'token' );
  if ( !token ) { ws.close( 1008, 'Missing token' ); return; }

  let payload: AuthPayload;
  try
  {
    payload = jwt.verify( token, process.env.JWT_SECRET ?? '' ) as AuthPayload;
  }
  catch
  {
    ws.close( 1008, 'Invalid token' );
    return;
  }

  const tunnel = getTunnel( tunnelId );
  if ( !tunnel ) { ws.close( 1008, 'Tunnel not found' ); return; }
  if ( tunnel.ownerId !== payload.userId ) { ws.close( 1008, 'Forbidden' ); return; }

  registry.register( tunnelId, ws );
  console.log( `[agent] connected: "${tunnel.name}" (${tunnelId})` );

  ws.on( 'message', ( data ) =>
  {
    try { dispatchAgentMsg( JSON.parse( data.toString() ) ); }
    catch { /* ignore malformed messages */ }
  } );

  ws.on( 'close', () =>
  {
    registry.unregister( tunnelId );
    console.log( `[agent] disconnected: "${tunnel.name}" (${tunnelId})` );
  } );
}