Cloudflare Worker¶

export class WebSocketConnection {
  constructor(controller, env) {
    this.controller = controller;
    this.env = env;
    this.sessions = new Map();
  }

  async fetch(request) {
    if (request.headers.get("Upgrade") === "websocket") {
      return this.handleWebSocket(request);
    }
    return new Response("Expected WebSocket", { status: 400 });
  }

  async handleWebSocket(request) {
    const webSocketPair = new WebSocketPair();
    const [client, server] = Object.values(webSocketPair);

    server.accept();
    this.handleSession(server, request);

    return new Response(null, {
      status: 101,
      webSocket: client,
    });
  }

  handleSession(webSocket, request) {
    const sessionId = this.extractSessionId(request);
    const session = {
      webSocket,
      connectorId: null,
      authenticated: false,
      lastActivity: Date.now()
    };

    this.sessions.set(sessionId, session);

    webSocket.addEventListener('message', event => {
      this.handleMessage(sessionId, JSON.parse(event.data));
    });

    webSocket.addEventListener('close', () => {
      this.sessions.delete(sessionId);
    });
  }

  async handleMessage(sessionId, message) {
    const session = this.sessions.get(sessionId);
    if (!session) return;

    switch (message.type) {
      case 'register':
        await this.registerConnector(sessionId, message);
        break;
      case 'console-log':
        await this.relayToAI(sessionId, message);
        break;
      case 'heartbeat':
        session.lastActivity = Date.now();
        break;
    }
  }
}

// Key: session:{sessionId}
{
  "sessionId": "session123",
  "connectorId": "browser-connector-uuid", 
  "createdAt": 1704067200000,
  "lastActivity": 1704070800000,
  "authenticated": true,
  "capabilities": ["console", "network"],
  "tunnelActive": true
}

// Key: token:{tokenHash}
{
  "connectorId": "browser-connector-uuid",
  "issuedAt": 1704067200000,
  "expiresAt": 1704070800000,
  "scope": ["read", "write"],
  "sessionId": "session123"
}

// Key: analytics:{date}:{connectorId}
{
  "date": "2024-01-01",
  "connectorId": "browser-connector-uuid",
  "requests": 1247,
  "dataTransferred": 2048576,
  "errors": 3,
  "uptime": 0.999
}

async function validateJWT(token, env) {
  try {
    const payload = await jwt.verify(token, env.JWT_SECRET);

    // Validate token claims
    if (payload.exp < Date.now() / 1000) {
      throw new Error('Token expired');
    }

    if (!payload.connectorId) {
      throw new Error('Invalid connector ID');
    }

    return payload;
  } catch (error) {
    throw new Error('Invalid token');
  }
}

function validateOrigin(request) {
  const origin = request.headers.get('Origin');
  const allowedOrigins = [
    'http://localhost:3025',
    'https://rapidtriage.me',
    'chrome-extension://*'
  ];

  return allowedOrigins.some(allowed => 
    allowed === '*' || origin === allowed || 
    (allowed.endsWith('*') && origin.startsWith(allowed.slice(0, -1)))
  );
}

class RateLimiter {
  constructor(env) {
    this.env = env;
  }

  async checkLimit(clientId, limit = 100, window = 60) {
    const key = `ratelimit:${clientId}:${Math.floor(Date.now() / (window * 1000))}`;

    const current = await this.env.SESSIONS.get(key);
    const count = current ? parseInt(current) + 1 : 1;

    if (count > limit) {
      throw new Error('Rate limit exceeded');
    }

    await this.env.SESSIONS.put(key, count.toString(), {
      expirationTtl: window
    });

    return { count, limit, remaining: limit - count };
  }
}

// Track key performance indicators
async function recordMetrics(env, metrics) {
  await env.ANALYTICS.put(`metrics:${Date.now()}`, JSON.stringify({
    timestamp: new Date().toISOString(),
    tunnelsActive: metrics.tunnelsActive,
    requestsPerSecond: metrics.requestsPerSecond,
    dataTransferredMB: metrics.dataTransferredMB,
    errorRate: metrics.errorRate,
    responseTimeP95: metrics.responseTimeP95
  }));
}

function handleError(error, context) {
  console.error('Worker Error:', {
    message: error.message,
    stack: error.stack,
    context: context,
    timestamp: new Date().toISOString()
  });

  // Send to external error tracking (optional)
  // await sendToSentry(error, context);
}

async function healthCheck(env) {
  const checks = {
    kvStorage: await testKVStorage(env),
    durableObjects: await testDurableObjects(env),
    externalAPIs: await testExternalAPIs(env)
  };

  const healthy = Object.values(checks).every(check => check.status === 'ok');

  return {
    status: healthy ? 'healthy' : 'degraded',
    checks,
    timestamp: new Date().toISOString()
  };
}

// Enable detailed logging in development
const DEBUG = env.ENVIRONMENT === 'development';

function debugLog(message, data) {
  if (DEBUG) {
    console.log(`[DEBUG] ${message}`, data);
  }
}

// Cache frequently accessed data
async function getCachedData(env, key, fetchFn, ttl = 300) {
  let cached = await env.CACHE.get(key);

  if (cached) {
    return JSON.parse(cached);
  }

  const data = await fetchFn();
  await env.CACHE.put(key, JSON.stringify(data), {
    expirationTtl: ttl
  });

  return data;
}

// Cleanup inactive sessions
async function cleanupSessions(env) {
  const cutoff = Date.now() - (30 * 60 * 1000); // 30 minutes

  const keys = await env.SESSIONS.list({ prefix: 'session:' });

  for (const key of keys.keys) {
    const session = await env.SESSIONS.get(key.name);
    if (session && JSON.parse(session).lastActivity < cutoff) {
      await env.SESSIONS.delete(key.name);
    }
  }
}

// Test worker health
await fetch('https://rapidtriage.me/health');

// Validate authentication
await fetch('https://rapidtriage.me/auth/validate', {
  headers: { 'Authorization': 'Bearer ' + token }
});

// Check session status
await fetch(`https://rapidtriage.me/tunnel/${sessionId}/status`, {
  headers: { 'Authorization': 'Bearer ' + token }
});