Multi-Hoster-Upload/gateway/agent-client.js

import WebSocket from 'ws';
import { randomUUID } from 'node:crypto';

const AUTH_TIMEOUT_MS = 6000;
const REQUEST_TIMEOUT_MS = 30000;

const CLOSE_CODE_GUIDANCE = {
  4001: 'auth timeout, retry',
  4002: 'stale or rotated code — re-copy the current code from the server',
  4003: 'brute-force lockout, wait 60s',
};

function normalizeFingerprint(fp) {
  if (typeof fp !== 'string') return '';
  return fp.replace(/:/g, '').toLowerCase();
}

function mapSocketError(err) {
  const code = err && err.code;
  if (code === 'ECONNREFUSED') {
    return 'app not running / wrong port / inbound firewall closed';
  }
  if (code === 'ETIMEDOUT') {
    return 'firewall DROP / NAT not forwarded / tunnel down';
  }
  return (err && err.message) ? err.message : String(err);
}

function mapCloseBeforeAuth(code, sawOpen) {
  if (CLOSE_CODE_GUIDANCE[code]) return CLOSE_CODE_GUIDANCE[code];
  if (sawOpen) {
    return 'old app version without the diagnostic agent — update that server';
  }
  return `connection closed before auth (code ${code})`;
}

export class AgentClient {
  constructor({ host, port, token, fp }) {
    this.host = host;
    this.port = port;
    this.token = token;
    this.fp = fp;
    this.connected = false;
    this.clientId = null;
    this.ws = null;
    this.authed = false;
    this.sawOpen = false;
    this._pending = new Map();
    this._authResolve = null;
    this._authReject = null;
    this._authTimer = null;
  }

  connect() {
    return new Promise((resolve, reject) => {
      this._authResolve = resolve;
      this._authReject = reject;

      const secure = typeof this.fp === 'string' && this.fp.length > 0;
      const scheme = secure ? 'wss' : 'ws';
      const url = `${scheme}://${this.host}:${this.port}`;

      const pinned = secure ? normalizeFingerprint(this.fp) : '';
      const options = secure ? { rejectUnauthorized: false } : undefined;

      let ws;
      try {
        ws = options ? new WebSocket(url, options) : new WebSocket(url);
      } catch (err) {
        this._failAuth(mapSocketError(err));
        return;
      }
      this.ws = ws;

      this._authTimer = setTimeout(() => {
        this._failAuth('auth timeout, retry');
        try { ws.close(); } catch {}
      }, AUTH_TIMEOUT_MS);

      ws.on('open', () => {
        this.sawOpen = true;
        if (secure) {
          const sock = ws._socket;
          const cert = sock && typeof sock.getPeerCertificate === 'function'
            ? sock.getPeerCertificate()
            : null;
          const actual = normalizeFingerprint(cert && cert.fingerprint256);
          if (!actual || actual !== pinned) {
            this._failAuth('server cert changed (reinstalled?) — re-copy the code');
            try { ws.close(); } catch {}
            return;
          }
        }
        this._send({ type: 'auth', token: this.token, role: 'diagnostic' });
      });

      ws.on('message', (raw) => this._onMessage(raw));

      ws.on('error', (err) => {
        if (!this.authed) {
          this._failAuth(mapSocketError(err));
        } else {
          this._rejectAllPending(mapSocketError(err));
        }
      });

      ws.on('close', (code) => {
        this.connected = false;
        if (!this.authed) {
          this._failAuth(mapCloseBeforeAuth(code, this.sawOpen));
        } else {
          this._rejectAllPending(`connection closed (code ${code})`);
        }
      });
    });
  }

  _onMessage(raw) {
    let msg;
    try {
      msg = JSON.parse(raw.toString());
    } catch {
      return;
    }

    if (msg.type === 'auth-ok') {
      this.authed = true;
      this.connected = true;
      this.clientId = msg.clientId ?? null;
      if (this._authTimer) {
        clearTimeout(this._authTimer);
        this._authTimer = null;
      }
      if (this._authResolve) {
        const r = this._authResolve;
        this._authResolve = null;
        this._authReject = null;
        r({ clientId: this.clientId });
      }
      return;
    }

    if (msg.type === 'diag-response' && msg.reqId) {
      const entry = this._pending.get(msg.reqId);
      if (!entry) return;
      this._pending.delete(msg.reqId);
      clearTimeout(entry.timer);
      if (msg.ok) {
        entry.resolve({ ok: true, data: msg.data });
      } else {
        entry.resolve({ ok: false, error: msg.error ?? 'unknown agent error' });
      }
    }
  }

  request(op, args) {
    return new Promise((resolve) => {
      if (!this.connected || !this.ws || this.ws.readyState !== WebSocket.OPEN) {
        resolve({ ok: false, error: 'not connected to a diagnostic agent' });
        return;
      }
      const reqId = randomUUID();
      const timer = setTimeout(() => {
        if (this._pending.has(reqId)) {
          this._pending.delete(reqId);
          resolve({ ok: false, error: `request timed out after ${REQUEST_TIMEOUT_MS}ms (op: ${op})` });
        }
      }, REQUEST_TIMEOUT_MS);
      this._pending.set(reqId, { resolve, timer });
      try {
        this._send({ type: 'diag-request', reqId, op, args: args ?? {} });
      } catch (err) {
        this._pending.delete(reqId);
        clearTimeout(timer);
        resolve({ ok: false, error: mapSocketError(err) });
      }
    });
  }

  close() {
    this.connected = false;
    this.authed = false;
    if (this._authTimer) {
      clearTimeout(this._authTimer);
      this._authTimer = null;
    }
    this._rejectAllPending('connection closed by client');
    if (this.ws) {
      try { this.ws.removeAllListeners(); } catch {}
      try { this.ws.close(); } catch {}
      this.ws = null;
    }
  }

  _send(obj) {
    this.ws.send(JSON.stringify(obj));
  }

  _failAuth(error) {
    if (this._authTimer) {
      clearTimeout(this._authTimer);
      this._authTimer = null;
    }
    if (this._authReject) {
      const rej = this._authReject;
      this._authResolve = null;
      this._authReject = null;
      rej(new Error(error));
    }
  }

  _rejectAllPending(error) {
    for (const [reqId, entry] of this._pending) {
      clearTimeout(entry.timer);
      entry.resolve({ ok: false, error });
      this._pending.delete(reqId);
    }
  }
}