tunnel/electron-agent/agent/TunnelAgent.ts

import WebSocket from 'ws';
import http from 'http';

export interface TunnelAgentConfig
{
  tunnelId:  string;
  token:     string;
  localPort: number;
  serverUrl: string; // e.g. https://tunnel.rokojori.com
}

interface RelayRequest
{
  reqId:   string;
  method:  string;
  path:    string;
  headers: Record<string, string>;
  body:    string; // base64
}


export class TunnelAgent
{
  config: TunnelAgentConfig;
  ws: WebSocket | null = null;
  _active       = false;
  _shouldRun    = false;
  _reconnDelay  = 2000;
  _maxDelay     = 30_000;

  onStatus: ( active: boolean ) => void = () => {};

  constructor( config: TunnelAgentConfig ) { this.config = config; }

  start(): void   { this._shouldRun = true;  this._connect(); }
  stop(): void    { this._shouldRun = false; this.ws?.close(); this.ws = null; this._setActive( false ); }
  isActive(): boolean { return this._active; }

  private _setActive( v: boolean ): void
  {
    if ( this._active === v ) return;
    this._active = v;
    this.onStatus( v );
  }

  private _connect(): void
  {
    if ( !this._shouldRun ) return;

    const { serverUrl, tunnelId, token } = this.config;
    const wsBase = serverUrl.replace( /^http/, 'ws' );
    const url    = `${ wsBase }/api/agent/${ tunnelId }?token=${ token }`;

    this.ws = new WebSocket( url );

    this.ws.on( 'open', () =>
    {
      this._reconnDelay = 2000;
      this._setActive( true );
      console.log( `[agent ${ tunnelId.slice( 0, 8 ) }] connected → :${ this.config.localPort }` );
    } );

    this.ws.on( 'message', data => this._forward( data as Buffer ) );

    this.ws.on( 'close', () =>
    {
      this._setActive( false );
      if ( this._shouldRun )
      {
        setTimeout( () => this._connect(), this._reconnDelay );
        this._reconnDelay = Math.min( this._reconnDelay * 2, this._maxDelay );
      }
    } );

    this.ws.on( 'error', err =>
      console.error( `[agent ${ tunnelId.slice( 0, 8 ) }] error: ${ err.message }` )
    );
  }

  private _send( msg: object ): void { this.ws?.send( JSON.stringify( msg ) ); }

  private _forward( data: Buffer ): void
  {
    let req: RelayRequest;
    try { req = JSON.parse( data.toString() ) as RelayRequest; }
    catch { return; }

    const bodyBuf = Buffer.from( req.body ?? '', 'base64' );

    const opts: http.RequestOptions = {
      hostname: 'localhost',
      port:     this.config.localPort,
      method:   req.method,
      path:     req.path,
      headers:  { ...req.headers, 'content-length': String( bodyBuf.length ) },
    };

    const localReq = http.request( opts, localRes =>
    {
      const respHeaders: Record<string, string> = {};
      for ( const [ k, v ] of Object.entries( localRes.headers ) )
      {
        if ( typeof v === 'string' )   respHeaders[ k ] = v;
        else if ( Array.isArray( v ) ) respHeaders[ k ] = v.join( ', ' );
      }

      // Send headers immediately — starts streaming
      this._send( { type: 'res_start', reqId: req.reqId, status: localRes.statusCode ?? 200, headers: respHeaders } );

      localRes.on( 'data', ( chunk: Buffer ) =>
        this._send( { type: 'res_data', reqId: req.reqId, chunk: Buffer.from( chunk ).toString( 'base64' ) } )
      );

      localRes.on( 'end', () =>
        this._send( { type: 'res_end', reqId: req.reqId } )
      );
    } );

    localReq.on( 'error', err =>
    {
      this._send( { type: 'res_start', reqId: req.reqId, status: 502, headers: { 'content-type': 'application/json' } } );
      this._send( { type: 'res_data',  reqId: req.reqId, chunk: Buffer.from( JSON.stringify( { error: err.message } ) ).toString( 'base64' ) } );
      this._send( { type: 'res_end',   reqId: req.reqId } );
    } );

    if ( bodyBuf.length > 0 ) localReq.write( bodyBuf );
    localReq.end();
  }
}