HiveGrid-Wireframes/03 Quellen und Referenzen/hivegrid-cloud/apps/web/app/modules/chat/services/chatGatewayRateLimit.server.ts

import { MemoryAdapter, type SharedStore } from "@hivegrid/shared-store";

import { getSharedStore } from "~/shared-store.server";

export const CHAT_GATEWAY_TOKEN_SESSION_LIMIT = 20;
export const CHAT_GATEWAY_TOKEN_USER_LIMIT = 60;
export const CHAT_GATEWAY_TOKEN_WINDOW_SECONDS = 60;

const CHAT_GATEWAY_TOKEN_WINDOW_MS = CHAT_GATEWAY_TOKEN_WINDOW_SECONDS * 1_000;

const CHAT_GATEWAY_TOKEN_RATE_LIMIT_SCRIPT = `
local time = redis.call("TIME")
local now = (tonumber(time[1]) * 1000) + math.floor(tonumber(time[2]) / 1000)
local window_ms = tonumber(ARGV[1])
local session_limit = tonumber(ARGV[2])
local user_limit = tonumber(ARGV[3])

if not window_ms or window_ms <= 0 or not session_limit or session_limit <= 0 or
    not user_limit or user_limit < session_limit then
  return redis.error_reply("invalid chat token rate limit arguments")
end

local function read_bucket(key)
  local raw = redis.call("GET", key)
  if not raw then
    return { count = 0, resetAt = now + window_ms }
  end
  local ok, bucket = pcall(cjson.decode, raw)
  if not ok or type(bucket) ~= "table" or type(bucket.count) ~= "number" or
      bucket.count < 0 or bucket.count % 1 ~= 0 or type(bucket.resetAt) ~= "number" then
    return nil
  end
  if bucket.resetAt <= now then
    return { count = 0, resetAt = now + window_ms }
  end
  return bucket
end

local session_bucket = read_bucket(KEYS[1])
local user_bucket = read_bucket(KEYS[2])
if not session_bucket or not user_bucket then
  return redis.error_reply("invalid chat token rate limit bucket")
end

local session_blocked = session_bucket.count >= session_limit
local user_blocked = user_bucket.count >= user_limit
if session_blocked or user_blocked then
  local retry_after_ms = 1
  if session_blocked then retry_after_ms = math.max(retry_after_ms, session_bucket.resetAt - now) end
  if user_blocked then retry_after_ms = math.max(retry_after_ms, user_bucket.resetAt - now) end
  return { 0, 0, retry_after_ms }
end

session_bucket.count = session_bucket.count + 1
user_bucket.count = user_bucket.count + 1
redis.call("SET", KEYS[1], cjson.encode(session_bucket), "PX", math.max(1, session_bucket.resetAt - now))
redis.call("SET", KEYS[2], cjson.encode(user_bucket), "PX", math.max(1, user_bucket.resetAt - now))
local remaining = math.min(session_limit - session_bucket.count, user_limit - user_bucket.count)
return { 1, remaining, 0 }
`;

export type ChatGatewayTokenRateLimitInput = {
  tenantId: string;
  userId: string;
  sessionId: string;
};

export type ChatGatewayTokenRateLimitResult = { allowed: true; remaining: number } | { allowed: false; retryAfterSeconds: number };

let storeFactory: () => SharedStore = () => getSharedStore();

export async function checkChatGatewayTokenRateLimit(input: ChatGatewayTokenRateLimitInput): Promise<ChatGatewayTokenRateLimitResult> {
  const scope = `${encodeKeyPart(input.tenantId)}:${encodeKeyPart(input.userId)}`;
  const result = await storeFactory().eval(
    CHAT_GATEWAY_TOKEN_RATE_LIMIT_SCRIPT,
    [`token-session:${scope}:${encodeKeyPart(input.sessionId)}`, `token-user:${scope}`],
    [String(CHAT_GATEWAY_TOKEN_WINDOW_MS), String(CHAT_GATEWAY_TOKEN_SESSION_LIMIT), String(CHAT_GATEWAY_TOKEN_USER_LIMIT), String(Date.now())]
  );
  const decision = parseDecision(result);
  if (decision.allowed) return { allowed: true, remaining: decision.remaining };
  return { allowed: false, retryAfterSeconds: Math.max(1, Math.ceil(decision.retryAfterMs / 1_000)) };
}

function parseDecision(result: unknown): { allowed: true; remaining: number } | { allowed: false; retryAfterMs: number } {
  if (!Array.isArray(result) || result.length !== 3) throw new Error("chat token rate limit returned an invalid result");
  const allowed = Number(result[0]);
  const remaining = Number(result[1]);
  const retryAfterMs = Number(result[2]);
  if ((allowed !== 0 && allowed !== 1) || !Number.isSafeInteger(remaining) || remaining < 0 || !Number.isFinite(retryAfterMs) || retryAfterMs < 0) {
    throw new Error("chat token rate limit returned an invalid result");
  }
  return allowed === 1 ? { allowed: true, remaining } : { allowed: false, retryAfterMs };
}

function encodeKeyPart(value: string): string {
  return Buffer.from(value, "utf8").toString("base64url");
}

type TokenRateBucket = { count: number; resetAt: number };

async function runTokenRateLimitInMemory(adapter: MemoryAdapter, keys: string[], args: string[]): Promise<[number, number, number]> {
  if (keys.length !== 2 || args.length !== 4) throw new Error("invalid chat token rate limit arguments");
  const windowMs = requirePositiveInteger(args[0], "window");
  const sessionLimit = requirePositiveInteger(args[1], "session limit");
  const userLimit = requirePositiveInteger(args[2], "user limit");
  const now = Number(args[3]);
  if (userLimit < sessionLimit || !Number.isFinite(now)) throw new Error("invalid chat token rate limit arguments");

  const sessionBucket = readBucket(await adapter.get(keys[0] ?? ""), now, windowMs);
  const userBucket = readBucket(await adapter.get(keys[1] ?? ""), now, windowMs);
  const sessionBlocked = sessionBucket.count >= sessionLimit;
  const userBlocked = userBucket.count >= userLimit;
  if (sessionBlocked || userBlocked) {
    return [0, 0, Math.max(1, sessionBlocked ? sessionBucket.resetAt - now : 0, userBlocked ? userBucket.resetAt - now : 0)];
  }

  sessionBucket.count += 1;
  userBucket.count += 1;
  await Promise.all([persistBucket(adapter, keys[0], sessionBucket, now), persistBucket(adapter, keys[1], userBucket, now)]);
  return [1, Math.min(sessionLimit - sessionBucket.count, userLimit - userBucket.count), 0];
}

function readBucket(raw: string | null, now: number, windowMs: number): TokenRateBucket {
  if (raw === null) return { count: 0, resetAt: now + windowMs };
  let parsed: unknown;
  try {
    parsed = JSON.parse(raw) as unknown;
  } catch {
    throw new Error("invalid chat token rate limit bucket");
  }
  if (
    typeof parsed !== "object" ||
    parsed === null ||
    Array.isArray(parsed) ||
    !Number.isSafeInteger((parsed as Partial<TokenRateBucket>).count) ||
    ((parsed as Partial<TokenRateBucket>).count ?? -1) < 0 ||
    !Number.isFinite((parsed as Partial<TokenRateBucket>).resetAt)
  ) {
    throw new Error("invalid chat token rate limit bucket");
  }
  const bucket = parsed as TokenRateBucket;
  return bucket.resetAt <= now ? { count: 0, resetAt: now + windowMs } : bucket;
}

async function persistBucket(adapter: MemoryAdapter, key: string | undefined, bucket: TokenRateBucket, now: number): Promise<void> {
  if (!key) throw new Error("invalid chat token rate limit key");
  const stored = await adapter.set(key, JSON.stringify(bucket), { ttlSeconds: Math.max(1, Math.ceil((bucket.resetAt - now) / 1_000)) });
  if (!stored) throw new Error("chat token rate limit bucket could not be stored");
}

function requirePositiveInteger(value: string | undefined, label: string): number {
  const parsed = Number(value);
  if (!Number.isSafeInteger(parsed) || parsed <= 0) throw new Error(`invalid chat token rate limit ${label}`);
  return parsed;
}

MemoryAdapter.registerScript(CHAT_GATEWAY_TOKEN_RATE_LIMIT_SCRIPT, runTokenRateLimitInMemory);

export function __setChatGatewayTokenRateLimitStoreForTests(store: SharedStore): void {
  storeFactory = () => store;
}

export function __resetChatGatewayTokenRateLimitStoreForTests(): void {
  storeFactory = () => getSharedStore();
}