index.js

const ALLOWED_UA_PREFIX = 'GatekeeperRelay/';
const MAX_BROADCAST_BYTES = 64 * 1024;

export default {
  async fetch(request, env, ctx) {
    const url = new URL(request.url);
    // Simple health check end point
    if (url.pathname === '/health') return new Response('OK', { status: 200 });

    // UA missing means no connection for you.
    if (url.pathname === '/v1/kills') {
      if (request.method !== 'GET' || request.headers.get('Upgrade') !== 'websocket') {
        return new Response('Expected WebSocket upgrade', { status: 426 });
      }
      const ua = request.headers.get('User-Agent') || '';
      if (!ua.startsWith(ALLOWED_UA_PREFIX)) return new Response('Forbidden', { status: 403 });
    }

    const id = env.KILLSTREAM.idFromName('global');
    return env.KILLSTREAM.get(id).fetch(request);
  }
};

const MAX_CONNECTIONS = 500;   
const MAX_PER_IP = 3;          

export class KillStreamRoom {
  constructor(state, env) {
    this.state = state;
    this.env = env;
    this.stats = { broadcastsReceived: 0, broadcastsForwarded: 0,
      connectionsAccepted: 0, connectionsRejected: 0, authFailures: 0, startedAt: Date.now() };
  }

  async fetch(request) {
    const url = new URL(request.url);

    if (url.pathname === '/v1/kills') {
      if (request.headers.get('Upgrade') !== 'websocket') {
        return new Response('Expected WebSocket upgrade', { status: 426 });
      }
      const ip = request.headers.get('CF-Connecting-IP') || 'unknown';

      if (this.state.getWebSockets().length >= MAX_CONNECTIONS) {
        this.stats.connectionsRejected++;
        return new Response('At capacity', { status: 503 });
      }
      if (this.state.getWebSockets(ip).length >= MAX_PER_IP) {
        this.stats.connectionsRejected++;
        return new Response('Too many connections', { status: 429 });
      }

      const [client, server] = Object.values(new WebSocketPair());
      this.state.acceptWebSocket(server, [ip]);   
      this.stats.connectionsAccepted++;
      return new Response(null, { status: 101, webSocket: client });
    }

    if (url.pathname === '/broadcast' && request.method === 'POST') {
      const auth = request.headers.get('X-Backend-Auth') || '';
      if (!this.env.BACKEND_SECRET || auth !== this.env.BACKEND_SECRET) {
        this.stats.authFailures++;
        return new Response('Unauthorized', { status: 401 });
      }
      const body = await request.text();
      if (body.length > MAX_BROADCAST_BYTES) return new Response('Payload too large', { status: 413 });

      this.stats.broadcastsReceived++;
      let forwarded = 0;
      for (const ws of this.state.getWebSockets()) {
        try { ws.send(body); forwarded++; }
        catch { try { ws.close(1011); } catch {} }   
      }
      this.stats.broadcastsForwarded += forwarded;
      return Response.json({ ok: true, forwardedTo: forwarded });
    }

    if (url.pathname === '/stats') {
      if ((request.headers.get('X-Backend-Auth') || '') !== this.env.BACKEND_SECRET) {
        return new Response('Unauthorized', { status: 401 });
      }
      return Response.json({ ...this.stats,
        currentConnections: this.state.getWebSockets().length,
        uptimeSeconds: Math.floor((Date.now() - this.stats.startedAt) / 1000) });
    }

    return new Response('Not found', { status: 404 });
  }

  async webSocketMessage(ws, message) { try { ws.close(1008, 'read-only'); } catch {} }
  async webSocketClose() {}
  async webSocketError() {}
}