claude-code-limiter-server 1.0.0

package/src/server/services/auth.js

@@ -0,0 +1,174 @@
+'use strict';
+
+const crypto = require('crypto');
+const bcrypt = require('bcryptjs');
+const jwt = require('jsonwebtoken');
+const db = require('../db');
+
+// JWT secret: from env or auto-generated per process lifetime
+let jwtSecret = null;
+function getJWTSecret() {
+  if (jwtSecret) return jwtSecret;
+  jwtSecret = process.env.JWT_SECRET || crypto.randomUUID();
+  if (!process.env.JWT_SECRET) {
+    console.warn('WARNING: JWT_SECRET not set. Generated a random secret. Sessions will not persist across restarts.');
+  }
+  return jwtSecret;
+}
+
+// ---- Password helpers ----
+
+/**
+ * Hash a plain-text password.
+ * @param {string} plain
+ * @returns {string} bcrypt hash
+ */
+function hashPassword(plain) {
+  return bcrypt.hashSync(plain, 10);
+}
+
+/**
+ * Verify a plain-text password against a bcrypt hash.
+ * @param {string} plain
+ * @param {string} hash
+ * @returns {boolean}
+ */
+function verifyPassword(plain, hash) {
+  return bcrypt.compareSync(plain, hash);
+}
+
+// ---- Token helpers ----
+
+/**
+ * Generate a unique token (UUID v4).
+ * @returns {string}
+ */
+function generateToken() {
+  return crypto.randomUUID();
+}
+
+// ---- JWT helpers ----
+
+/**
+ * Create a JWT for the admin session.
+ * @param {object} payload - e.g., { teamId: '...' }
+ * @param {string} [expiresIn] - default '24h'
+ * @returns {string} signed JWT
+ */
+function createJWT(payload, expiresIn) {
+  return jwt.sign(payload, getJWTSecret(), { expiresIn: expiresIn || '24h' });
+}
+
+/**
+ * Verify and decode a JWT.
+ * @param {string} token
+ * @returns {object|null} decoded payload, or null if invalid
+ */
+function verifyJWT(token) {
+  try {
+    return jwt.verify(token, getJWTSecret());
+  } catch {
+    return null;
+  }
+}
+
+// ---- Middleware ----
+
+/**
+ * Express middleware: Admin authentication.
+ * Checks JWT from Authorization header (Bearer <token>) or cookie (jwt=<token>).
+ * Sets req.team on success.
+ */
+function adminAuth(req, res, next) {
+  let token = null;
+
+  // Check Authorization header
+  const authHeader = req.headers.authorization;
+  if (authHeader && authHeader.startsWith('Bearer ')) {
+    token = authHeader.slice(7);
+  }
+
+  // Fallback: check cookie
+  if (!token && req.headers.cookie) {
+    const cookies = parseCookies(req.headers.cookie);
+    token = cookies.jwt || cookies.token;
+  }
+
+  if (!token) {
+    return res.status(401).json({ error: 'Authentication required' });
+  }
+
+  const decoded = verifyJWT(token);
+  if (!decoded || !decoded.teamId) {
+    return res.status(401).json({ error: 'Invalid or expired token' });
+  }
+
+  const team = db.getTeam(decoded.teamId);
+  if (!team) {
+    return res.status(401).json({ error: 'Team not found' });
+  }
+
+  req.team = team;
+  req.teamId = decoded.teamId;
+  next();
+}
+
+/**
+ * Express middleware: Hook authentication.
+ * Checks auth_token from Authorization: Bearer header or request body.
+ * Sets req.user on success.
+ */
+function hookAuth(req, res, next) {
+  let authToken = null;
+
+  // Check Authorization header
+  const authHeader = req.headers.authorization;
+  if (authHeader && authHeader.startsWith('Bearer ')) {
+    authToken = authHeader.slice(7);
+  }
+
+  // Fallback: check request body
+  if (!authToken && req.body && req.body.auth_token) {
+    authToken = req.body.auth_token;
+  }
+
+  if (!authToken) {
+    return res.status(401).json({ error: 'auth_token required' });
+  }
+
+  const user = db.getUserByToken(authToken);
+  if (!user) {
+    return res.status(401).json({ error: 'Invalid auth_token' });
+  }
+
+  // Attach user and their team
+  req.user = user;
+  req.team = db.getTeam(user.team_id);
+  next();
+}
+
+/**
+ * Parse a Cookie header string into an object.
+ */
+function parseCookies(cookieHeader) {
+  const cookies = {};
+  if (!cookieHeader) return cookies;
+  cookieHeader.split(';').forEach(pair => {
+    const [name, ...rest] = pair.trim().split('=');
+    if (name) {
+      cookies[name.trim()] = decodeURIComponent(rest.join('=').trim());
+    }
+  });
+  return cookies;
+}
+
+module.exports = {
+  hashPassword,
+  verifyPassword,
+  generateToken,
+  createJWT,
+  verifyJWT,
+  adminAuth,
+  hookAuth,
+  getJWTSecret,
+};

package/src/server/services/limiter.js

@@ -0,0 +1,226 @@
+'use strict';
+
+const db = require('../db');
+
+/**
+ * Evaluate all limit rules for a user/model combination.
+ *
+ * Evaluation order (first deny wins):
+ *   1. Is user killed/paused?           -> BLOCK
+ *   2. Is model in allowed time window? -> if no, BLOCK
+ *   3. Is per-model cap exceeded?       -> if yes, BLOCK
+ *   4. Is credit budget exceeded?       -> if yes, BLOCK
+ *   5. ALLOW
+ *
+ * @param {object} user - The user row from the database.
+ * @param {string} model - The model being used (opus, sonnet, haiku, default).
+ * @param {object} creditWeights - { opus: 10, sonnet: 3, haiku: 1 }
+ * @returns {{ allowed: boolean, reason?: string, usage: object, credit_balance?: number, credit_budget?: number }}
+ */
+function evaluateLimits(user, model, creditWeights) {
+  // 1. Check user status
+  if (user.status === 'killed') {
+    return {
+      allowed: false,
+      reason: 'Your Claude Code access has been revoked by the admin. Contact your admin to restore access.',
+      usage: {},
+    };
+  }
+  if (user.status === 'paused') {
+    return {
+      allowed: false,
+      reason: 'Your Claude Code access has been paused by the admin. Contact your admin to resume.',
+      usage: {},
+    };
+  }
+
+  const rules = db.getLimitRules(user.id);
+  if (rules.length === 0) {
+    // No rules = unlimited
+    return { allowed: true, usage: {} };
+  }
+
+  // Gather usage data we will need
+  const usageCache = {}; // key: `${model}:${window}` -> count
+  const creditCache = {}; // key: window -> totalCredits
+
+  function getModelUsage(mdl, windowType) {
+    const key = `${mdl}:${windowType}`;
+    if (usageCache[key] !== undefined) return usageCache[key];
+    const windowStart = db.calculateWindowStart(windowType);
+    const rows = db.getDb().prepare(
+      'SELECT COUNT(*) AS count FROM usage_event WHERE user_id = ? AND model = ? AND timestamp >= ?'
+    ).get(user.id, mdl, windowStart);
+    usageCache[key] = rows.count;
+    return rows.count;
+  }
+
+  function getCreditUsage(windowType) {
+    if (creditCache[windowType] !== undefined) return creditCache[windowType];
+    const windowStart = db.calculateWindowStart(windowType);
+    const row = db.getDb().prepare(
+      'SELECT COALESCE(SUM(credit_cost), 0) AS total FROM usage_event WHERE user_id = ? AND timestamp >= ?'
+    ).get(user.id, windowStart);
+    creditCache[windowType] = row.total;
+    return row.total;
+  }
+
+  // Collect all usage for the response
+  const usageSummary = {};
+  let creditBalance = null;
+  let creditBudget = null;
+
+  // 2. Check time_of_day rules
+  const timeRules = rules.filter(r => r.type === 'time_of_day');
+  for (const rule of timeRules) {
+    // Only applies to the specified model, or all if model is null
+    if (rule.model && rule.model !== model) continue;
+
+    const inWindow = isInTimeWindow(rule.schedule_start, rule.schedule_end, rule.schedule_tz);
+    if (!inWindow) {
+      const tz = rule.schedule_tz || 'UTC';
+      const currentTime = getCurrentTimeInTZ(tz);
+      const modelName = rule.model || 'This model';
+      return {
+        allowed: false,
+        reason: `${capitalize(modelName)} is only available ${rule.schedule_start} - ${rule.schedule_end} (${tz}).\nCurrent time: ${currentTime}. Try another model instead.`,
+        usage: usageSummary,
+      };
+    }
+  }
+
+  // 3. Check per_model rules
+  const perModelRules = rules.filter(r => r.type === 'per_model');
+  for (const rule of perModelRules) {
+    // Skip rules that don't apply to this model
+    if (rule.model && rule.model !== model) continue;
+
+    const targetModel = rule.model || model;
+    const count = getModelUsage(targetModel, rule.window);
+    const limit = rule.value;
+
+    // Track in usage summary
+    if (!usageSummary[rule.window]) usageSummary[rule.window] = {};
+    usageSummary[rule.window][targetModel] = { used: count, limit };
+
+    if (limit === -1) continue; // unlimited
+    if (limit === 0 || count >= limit) {
+      const windowLabel = windowToLabel(rule.window);
+      return {
+        allowed: false,
+        reason: `${capitalize(windowLabel)} ${targetModel} limit reached for ${user.name}.\nUsed ${count}/${limit} prompts.\n\nSwitch to another model or try again later.`,
+        usage: usageSummary,
+      };
+    }
+  }
+
+  // 4. Check credit budget rules
+  const creditRules = rules.filter(r => r.type === 'credits');
+  for (const rule of creditRules) {
+    const budget = rule.value;
+    if (budget === -1) continue; // unlimited
+
+    const usedCredits = getCreditUsage(rule.window);
+    const balance = Math.max(0, budget - usedCredits);
+
+    creditBalance = balance;
+    creditBudget = budget;
+
+    if (!usageSummary[rule.window]) usageSummary[rule.window] = {};
+    usageSummary[rule.window]._credits = { used: usedCredits, budget, balance };
+
+    // Check if the next prompt would exceed the budget
+    const modelCost = creditWeights[model] || creditWeights['default'] || 1;
+    if (balance < modelCost) {
+      const windowLabel = windowToLabel(rule.window);
+      return {
+        allowed: false,
+        reason: `${capitalize(windowLabel)} credit budget exhausted for ${user.name}.\nUsed ${usedCredits}/${budget} credits (${balance} remaining).\n${capitalize(model)} costs ${modelCost} credits per prompt.\n\nTry a cheaper model or wait for the window to reset.`,
+        usage: usageSummary,
+        credit_balance: balance,
+        credit_budget: budget,
+      };
+    }
+  }
+
+  // 5. ALLOW
+  // Populate usage for the response even on allow
+  const defaultWindow = creditRules.length > 0 ? creditRules[0].window : (perModelRules.length > 0 ? perModelRules[0].window : 'daily');
+  if (Object.keys(usageSummary).length === 0) {
+    const windowStart = db.calculateWindowStart(defaultWindow);
+    const data = db.getUsageWithCredits(user.id, windowStart);
+    usageSummary[defaultWindow] = data.counts;
+  }
+
+  return {
+    allowed: true,
+    usage: usageSummary,
+    credit_balance: creditBalance,
+    credit_budget: creditBudget,
+  };
+}
+
+/**
+ * Check if the current time falls within a schedule window in a given timezone.
+ * @param {string} startTime - "HH:MM" format
+ * @param {string} endTime - "HH:MM" format
+ * @param {string} [tz] - IANA timezone (e.g., "America/New_York")
+ * @returns {boolean}
+ */
+function isInTimeWindow(startTime, endTime, tz) {
+  if (!startTime || !endTime) return true; // No schedule = always allowed
+
+  const timeZone = tz || 'UTC';
+  const now = new Date();
+
+  // Get current time in the target timezone
+  const parts = db.getDatePartsInTZ(now, timeZone);
+  const currentMinutes = parts.hour * 60 + parts.minute;
+
+  const [startH, startM] = startTime.split(':').map(Number);
+  const [endH, endM] = endTime.split(':').map(Number);
+  const startMinutes = startH * 60 + startM;
+  const endMinutes = endH * 60 + endM;
+
+  if (startMinutes <= endMinutes) {
+    // Same-day window (e.g., 09:00 - 18:00)
+    return currentMinutes >= startMinutes && currentMinutes < endMinutes;
+  } else {
+    // Overnight window (e.g., 22:00 - 06:00)
+    return currentMinutes >= startMinutes || currentMinutes < endMinutes;
+  }
+}
+
+/**
+ * Get the current time formatted in a given timezone.
+ */
+function getCurrentTimeInTZ(tz) {
+  const now = new Date();
+  const formatter = new Intl.DateTimeFormat('en-US', {
+    timeZone: tz,
+    hour: 'numeric',
+    minute: '2-digit',
+    hour12: true,
+  });
+  return formatter.format(now);
+}
+
+function windowToLabel(window) {
+  const labels = {
+    daily: 'daily',
+    weekly: 'weekly',
+    monthly: 'monthly',
+    sliding_24h: 'sliding 24-hour',
+  };
+  return labels[window] || window;
+}
+
+function capitalize(str) {
+  if (!str) return '';
+  return str.charAt(0).toUpperCase() + str.slice(1);
+}
+
+module.exports = {
+  evaluateLimits,
+  isInTimeWindow,
+};

package/src/server/services/usage.js

@@ -0,0 +1,87 @@
+'use strict';
+
+const db = require('../db');
+
+/**
+ * Record a usage event (one completed turn).
+ * @param {string} userId
+ * @param {string} model - opus | sonnet | haiku | default
+ * @param {number} creditCost - cost from credit_weights
+ * @param {string} [timestamp] - ISO string, defaults to now
+ * @param {string} [source] - 'hook' | 'server', defaults to 'hook'
+ */
+function recordEvent(userId, model, creditCost, timestamp, source) {
+  return db.recordUsage({
+    userId,
+    model,
+    creditCost,
+    timestamp: timestamp || new Date().toISOString(),
+    source: source || 'hook',
+  });
+}
+
+/**
+ * Get usage counts by model for a user within a window.
+ * @param {string} userId
+ * @param {string} windowType - daily | weekly | monthly | sliding_24h
+ * @param {string} [tz] - IANA timezone
+ * @returns {{ opus: number, sonnet: number, haiku: number, default: number }}
+ */
+function getUsageByWindow(userId, windowType, tz) {
+  const windowStart = db.calculateWindowStart(windowType, tz);
+  return db.getUsage(userId, windowStart);
+}
+
+/**
+ * Get a full usage summary for a user across all windows.
+ * @param {string} userId
+ * @param {string} [tz] - IANA timezone
+ * @returns {{ daily: object, weekly: object, monthly: object, sliding_24h: object }}
+ */
+function getUsageSummary(userId, tz) {
+  const windows = ['daily', 'weekly', 'monthly', 'sliding_24h'];
+  const summary = {};
+  for (const w of windows) {
+    const data = db.getUsageForWindow(userId, w, tz);
+    summary[w] = {
+      counts: data.counts,
+      totalCredits: data.totalCredits,
+      windowStart: data.windowStart,
+    };
+  }
+  return summary;
+}
+
+/**
+ * Get remaining credit balance for a user.
+ * @param {string} userId
+ * @param {object} creditWeights - { opus: 10, sonnet: 3, haiku: 1 }
+ * @param {string} windowType - daily | weekly | monthly | sliding_24h
+ * @param {number} budget - total credit budget
+ * @param {string} [tz] - IANA timezone
+ * @returns {{ balance: number, used: number, budget: number }}
+ */
+function getCreditBalance(userId, creditWeights, windowType, budget, tz) {
+  const windowStart = db.calculateWindowStart(windowType, tz);
+  const data = db.getUsageWithCredits(userId, windowStart);
+  const used = data.totalCredits;
+  const balance = Math.max(0, budget - used);
+  return { balance, used, budget };
+}
+
+/**
+ * Delete usage events older than N days.
+ * @param {number} days
+ * @returns {{ changes: number }}
+ */
+function cleanupOldEvents(days) {
+  return db.cleanupOldEvents(days || 90);
+}
+
+module.exports = {
+  recordEvent,
+  getUsageByWindow,
+  getUsageSummary,
+  getCreditBalance,
+  cleanupOldEvents,
+};

package/src/server/ws.js

@@ -0,0 +1,74 @@
+'use strict';
+
+const { WebSocketServer } = require('ws');
+
+let wss = null;
+
+/**
+ * Set up a WebSocket server on the given HTTP server.
+ * Listens on the /ws path.
+ * @param {http.Server} server
+ */
+function setupWebSocket(server) {
+  wss = new WebSocketServer({ noServer: true });
+
+  // Handle upgrade requests for /ws path
+  server.on('upgrade', (request, socket, head) => {
+    const { pathname } = new URL(request.url, `http://${request.headers.host}`);
+    if (pathname === '/ws') {
+      wss.handleUpgrade(request, socket, head, (ws) => {
+        wss.emit('connection', ws, request);
+      });
+    } else {
+      socket.destroy();
+    }
+  });
+
+  wss.on('connection', (ws) => {
+    // Send a welcome message
+    ws.send(JSON.stringify({
+      type: 'connected',
+      timestamp: new Date().toISOString(),
+    }));
+
+    ws.on('error', (err) => {
+      console.error('WebSocket client error:', err.message);
+    });
+  });
+
+  console.log('WebSocket server ready on /ws');
+}
+
+/**
+ * Broadcast an event to all connected dashboard clients.
+ * @param {object} event - Event object with at least a `type` field.
+ *   Supported types: user_check, user_blocked, user_counted, user_killed, user_status_change
+ */
+function broadcast(event) {
+  if (!wss) return;
+
+  const message = JSON.stringify(event);
+  for (const client of wss.clients) {
+    if (client.readyState === 1) { // WebSocket.OPEN
+      client.send(message);
+    }
+  }
+}
+
+/**
+ * Get the number of connected clients.
+ */
+function getClientCount() {
+  if (!wss) return 0;
+  let count = 0;
+  for (const client of wss.clients) {
+    if (client.readyState === 1) count++;
+  }
+  return count;
+}
+
+module.exports = {
+  setupWebSocket,
+  broadcast,
+  getClientCount,
+};