49 lines
1.5 KiB
TypeScript
49 lines
1.5 KiB
TypeScript
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 { resolvePending } 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 { resolvePending( JSON.parse( data.toString() ) ); }
|
|
catch { /* ignore malformed messages */ }
|
|
} );
|
|
|
|
ws.on( 'close', () =>
|
|
{
|
|
registry.unregister( tunnelId );
|
|
console.log( `[agent] disconnected: "${tunnel.name}" (${tunnelId})` );
|
|
} );
|
|
}
|