init

2026-03-13 11:00:01 +08:00
commit fa4c458eda
51 changed files with 8843 additions and 0 deletions
--- a/lib/auth/api-key.ts
+++ b/lib/auth/api-key.ts
@@ -0,0 +1,47 @@
+import crypto from "crypto";
+import { db } from "@/lib/db";
+import { lobsters } from "@/lib/db/schema";
+import { eq } from "drizzle-orm";
+import { redis } from "@/lib/redis";
+
+const API_KEY_CACHE_TTL = 3600; // 1 hour in seconds
+
+export function generateApiKey(): string {
+  return crypto.randomBytes(32).toString("hex");
+}
+
+export async function validateApiKey(apiKey: string) {
+  try {
+    const cacheKey = `lobster:key:${apiKey}`;
+    const cachedLobsterId = await redis.get(cacheKey);
+
+    if (cachedLobsterId) {
+      const lobster = await db
+        .select()
+        .from(lobsters)
+        .where(eq(lobsters.id, cachedLobsterId))
+        .limit(1);
+
+      if (lobster.length > 0) {
+        return lobster[0];
+      }
+    }
+
+    const lobster = await db
+      .select()
+      .from(lobsters)
+      .where(eq(lobsters.apiKey, apiKey))
+      .limit(1);
+
+    if (lobster.length === 0) {
+      return null;
+    }
+
+    await redis.set(cacheKey, lobster[0].id, "EX", API_KEY_CACHE_TTL);
+
+    return lobster[0];
+  } catch (error) {
+    console.error("Failed to validate API key:", error);
+    return null;
+  }
+}
--- a/lib/db/index.ts
+++ b/lib/db/index.ts
@@ -0,0 +1,19 @@
+import { drizzle } from "drizzle-orm/mysql2";
+import mysql from "mysql2/promise";
+import * as schema from "./schema";
+
+const globalForDb = globalThis as unknown as {
+  connection: mysql.Pool | undefined;
+};
+
+const connection =
+  globalForDb.connection ??
+  mysql.createPool({
+    uri: process.env.DATABASE_URL!,
+    waitForConnections: true,
+    connectionLimit: 10,
+  });
+
+if (process.env.NODE_ENV !== "production") globalForDb.connection = connection;
+
+export const db = drizzle(connection, { schema, mode: "default" });
--- a/lib/db/schema.ts
+++ b/lib/db/schema.ts
@@ -0,0 +1,72 @@
+import {
+  mysqlTable,
+  varchar,
+  int,
+  bigint,
+  decimal,
+  datetime,
+  json,
+  index,
+} from "drizzle-orm/mysql-core";
+import { sql } from "drizzle-orm";
+
+export const lobsters = mysqlTable("lobsters", {
+  id: varchar("id", { length: 21 }).primaryKey(),
+  apiKey: varchar("api_key", { length: 64 }).notNull().unique(),
+  name: varchar("name", { length: 100 }).notNull(),
+  platform: varchar("platform", { length: 20 }),
+  model: varchar("model", { length: 50 }),
+  ip: varchar("ip", { length: 45 }),
+  latitude: decimal("latitude", { precision: 10, scale: 7 }),
+  longitude: decimal("longitude", { precision: 10, scale: 7 }),
+  city: varchar("city", { length: 100 }),
+  country: varchar("country", { length: 100 }),
+  countryCode: varchar("country_code", { length: 5 }),
+  region: varchar("region", { length: 50 }),
+  lastHeartbeat: datetime("last_heartbeat"),
+  totalTasks: int("total_tasks").default(0),
+  createdAt: datetime("created_at").default(sql`NOW()`),
+  updatedAt: datetime("updated_at").default(sql`NOW()`),
+});
+
+export const heartbeats = mysqlTable(
+  "heartbeats",
+  {
+    id: bigint("id", { mode: "number" }).primaryKey().autoincrement(),
+    lobsterId: varchar("lobster_id", { length: 21 }).notNull(),
+    ip: varchar("ip", { length: 45 }),
+    timestamp: datetime("timestamp").default(sql`NOW()`),
+  },
+  (table) => [
+    index("heartbeats_lobster_id_idx").on(table.lobsterId),
+    index("heartbeats_timestamp_idx").on(table.timestamp),
+  ]
+);
+
+export const tasks = mysqlTable(
+  "tasks",
+  {
+    id: bigint("id", { mode: "number" }).primaryKey().autoincrement(),
+    lobsterId: varchar("lobster_id", { length: 21 }).notNull(),
+    summary: varchar("summary", { length: 500 }),
+    durationMs: int("duration_ms"),
+    model: varchar("model", { length: 50 }),
+    toolsUsed: json("tools_used").$type<string[]>(),
+    timestamp: datetime("timestamp").default(sql`NOW()`),
+  },
+  (table) => [
+    index("tasks_lobster_id_idx").on(table.lobsterId),
+    index("tasks_timestamp_idx").on(table.timestamp),
+  ]
+);
+
+export const geoCache = mysqlTable("geo_cache", {
+  ip: varchar("ip", { length: 45 }).primaryKey(),
+  latitude: decimal("latitude", { precision: 10, scale: 7 }),
+  longitude: decimal("longitude", { precision: 10, scale: 7 }),
+  city: varchar("city", { length: 100 }),
+  country: varchar("country", { length: 100 }),
+  countryCode: varchar("country_code", { length: 5 }),
+  region: varchar("region", { length: 50 }),
+  updatedAt: datetime("updated_at").default(sql`NOW()`),
+});
--- a/lib/geo/ip-location.ts
+++ b/lib/geo/ip-location.ts
@@ -0,0 +1,245 @@
+import { db } from "@/lib/db";
+import { geoCache } from "@/lib/db/schema";
+import { eq } from "drizzle-orm";
+
+interface GeoLocation {
+  latitude: number;
+  longitude: number;
+  city: string;
+  country: string;
+  countryCode: string;
+  region: string;
+}
+
+interface IpApiResponse {
+  status: string;
+  lat: number;
+  lon: number;
+  city: string;
+  country: string;
+  countryCode: string;
+}
+
+const CONTINENT_MAP: Record<string, string> = {
+  // Asia
+  CN: "Asia",
+  JP: "Asia",
+  KR: "Asia",
+  IN: "Asia",
+  ID: "Asia",
+  TH: "Asia",
+  VN: "Asia",
+  PH: "Asia",
+  MY: "Asia",
+  SG: "Asia",
+  TW: "Asia",
+  HK: "Asia",
+  BD: "Asia",
+  PK: "Asia",
+  LK: "Asia",
+  NP: "Asia",
+  MM: "Asia",
+  KH: "Asia",
+  LA: "Asia",
+  MN: "Asia",
+  KZ: "Asia",
+  UZ: "Asia",
+  AE: "Asia",
+  SA: "Asia",
+  IL: "Asia",
+  TR: "Asia",
+  IQ: "Asia",
+  IR: "Asia",
+  QA: "Asia",
+  KW: "Asia",
+  BH: "Asia",
+  OM: "Asia",
+  JO: "Asia",
+  LB: "Asia",
+  AF: "Asia",
+
+  // Europe
+  GB: "Europe",
+  DE: "Europe",
+  FR: "Europe",
+  IT: "Europe",
+  ES: "Europe",
+  PT: "Europe",
+  NL: "Europe",
+  BE: "Europe",
+  SE: "Europe",
+  NO: "Europe",
+  DK: "Europe",
+  FI: "Europe",
+  PL: "Europe",
+  CZ: "Europe",
+  AT: "Europe",
+  CH: "Europe",
+  IE: "Europe",
+  RO: "Europe",
+  HU: "Europe",
+  GR: "Europe",
+  UA: "Europe",
+  RU: "Europe",
+  BG: "Europe",
+  HR: "Europe",
+  SK: "Europe",
+  SI: "Europe",
+  RS: "Europe",
+  LT: "Europe",
+  LV: "Europe",
+  EE: "Europe",
+  IS: "Europe",
+  LU: "Europe",
+  MT: "Europe",
+  CY: "Europe",
+  AL: "Europe",
+  MK: "Europe",
+  BA: "Europe",
+  ME: "Europe",
+  MD: "Europe",
+  BY: "Europe",
+
+  // Americas
+  US: "Americas",
+  CA: "Americas",
+  MX: "Americas",
+  BR: "Americas",
+  AR: "Americas",
+  CO: "Americas",
+  CL: "Americas",
+  PE: "Americas",
+  VE: "Americas",
+  EC: "Americas",
+  BO: "Americas",
+  PY: "Americas",
+  UY: "Americas",
+  CR: "Americas",
+  PA: "Americas",
+  CU: "Americas",
+  DO: "Americas",
+  GT: "Americas",
+  HN: "Americas",
+  SV: "Americas",
+  NI: "Americas",
+  JM: "Americas",
+  TT: "Americas",
+  HT: "Americas",
+  PR: "Americas",
+  GY: "Americas",
+  SR: "Americas",
+  BZ: "Americas",
+
+  // Africa
+  ZA: "Africa",
+  NG: "Africa",
+  KE: "Africa",
+  EG: "Africa",
+  GH: "Africa",
+  ET: "Africa",
+  TZ: "Africa",
+  MA: "Africa",
+  DZ: "Africa",
+  TN: "Africa",
+  UG: "Africa",
+  SN: "Africa",
+  CI: "Africa",
+  CM: "Africa",
+  MZ: "Africa",
+  MG: "Africa",
+  AO: "Africa",
+  ZW: "Africa",
+  RW: "Africa",
+  LY: "Africa",
+  SD: "Africa",
+  CD: "Africa",
+  ML: "Africa",
+  BF: "Africa",
+  NE: "Africa",
+  MW: "Africa",
+  ZM: "Africa",
+  BW: "Africa",
+  NA: "Africa",
+  MU: "Africa",
+
+  // Oceania
+  AU: "Oceania",
+  NZ: "Oceania",
+  FJ: "Oceania",
+  PG: "Oceania",
+  WS: "Oceania",
+  TO: "Oceania",
+  VU: "Oceania",
+  SB: "Oceania",
+  GU: "Oceania",
+  NC: "Oceania",
+  PF: "Oceania",
+};
+
+function getRegion(countryCode: string): string {
+  return CONTINENT_MAP[countryCode] ?? "Unknown";
+}
+
+export async function getGeoLocation(
+  ip: string
+): Promise<GeoLocation | null> {
+  try {
+    const cached = await db
+      .select()
+      .from(geoCache)
+      .where(eq(geoCache.ip, ip))
+      .limit(1);
+
+    if (cached.length > 0) {
+      const row = cached[0];
+      return {
+        latitude: Number(row.latitude),
+        longitude: Number(row.longitude),
+        city: row.city ?? "",
+        country: row.country ?? "",
+        countryCode: row.countryCode ?? "",
+        region: row.region ?? "",
+      };
+    }
+
+    const response = await fetch(`http://ip-api.com/json/${ip}`, {
+      signal: AbortSignal.timeout(5000),
+    });
+
+    if (!response.ok) {
+      return null;
+    }
+
+    const data: IpApiResponse = await response.json();
+
+    if (data.status !== "success") {
+      return null;
+    }
+
+    const region = getRegion(data.countryCode);
+
+    const geoData: GeoLocation = {
+      latitude: data.lat,
+      longitude: data.lon,
+      city: data.city,
+      country: data.country,
+      countryCode: data.countryCode,
+      region,
+    };
+
+    await db.insert(geoCache).values({
+      ip,
+      latitude: String(geoData.latitude),
+      longitude: String(geoData.longitude),
+      city: geoData.city,
+      country: geoData.country,
+      countryCode: geoData.countryCode,
+      region: geoData.region,
+    });
+
+    return geoData;
+  } catch (error) {
+    console.error("Failed to get geo location for IP:", ip, error);
+    return null;
+  }
+}
--- a/lib/redis/index.ts
+++ b/lib/redis/index.ts
@@ -0,0 +1,107 @@
+import Redis from "ioredis";
+
+const globalForRedis = globalThis as unknown as {
+  redis: Redis | undefined;
+  redisSub: Redis | undefined;
+};
+
+function createRedisClient(): Redis {
+  return new Redis(process.env.REDIS_URL!, {
+    maxRetriesPerRequest: 3,
+    retryStrategy(times) {
+      const delay = Math.min(times * 50, 2000);
+      return delay;
+    },
+  });
+}
+
+export const redis = globalForRedis.redis ?? createRedisClient();
+export const redisSub = globalForRedis.redisSub ?? createRedisClient();
+
+if (process.env.NODE_ENV !== "production") {
+  globalForRedis.redis = redis;
+  globalForRedis.redisSub = redisSub;
+}
+
+const CHANNEL_REALTIME = "channel:realtime";
+const ACTIVE_LOBSTERS_KEY = "active:lobsters";
+const STATS_GLOBAL_KEY = "stats:global";
+const STATS_REGION_KEY = "stats:region";
+const HEATMAP_CACHE_KEY = "cache:heatmap";
+const HOURLY_ACTIVITY_KEY = "stats:hourly";
+
+export async function setLobsterOnline(
+  lobsterId: string,
+  ip: string
+): Promise<void> {
+  await redis.set(`lobster:online:${lobsterId}`, ip, "EX", 300);
+}
+
+export async function isLobsterOnline(lobsterId: string): Promise<boolean> {
+  const result = await redis.exists(`lobster:online:${lobsterId}`);
+  return result === 1;
+}
+
+export async function updateActiveLobsters(lobsterId: string): Promise<void> {
+  const now = Date.now();
+  await redis.zadd(ACTIVE_LOBSTERS_KEY, now, lobsterId);
+}
+
+export async function getActiveLobsterIds(
+  limit: number = 100
+): Promise<string[]> {
+  return redis.zrevrange(ACTIVE_LOBSTERS_KEY, 0, limit - 1);
+}
+
+export async function incrementRegionCount(region: string): Promise<void> {
+  await redis.hincrby(STATS_REGION_KEY, region, 1);
+}
+
+export async function incrementGlobalStat(
+  field: string,
+  amount: number = 1
+): Promise<void> {
+  await redis.hincrby(STATS_GLOBAL_KEY, field, amount);
+}
+
+export async function getGlobalStats(): Promise<Record<string, string>> {
+  return redis.hgetall(STATS_GLOBAL_KEY);
+}
+
+export async function getRegionStats(): Promise<Record<string, string>> {
+  return redis.hgetall(STATS_REGION_KEY);
+}
+
+export async function incrementHourlyActivity(): Promise<void> {
+  const now = new Date();
+  const hourKey = `${now.getUTCFullYear()}-${String(now.getUTCMonth() + 1).padStart(2, "0")}-${String(now.getUTCDate()).padStart(2, "0")}T${String(now.getUTCHours()).padStart(2, "0")}`;
+
+  await redis.hincrby(HOURLY_ACTIVITY_KEY, hourKey, 1);
+  await redis.expire(HOURLY_ACTIVITY_KEY, 48 * 60 * 60);
+}
+
+export async function getHourlyActivity(): Promise<Record<string, number>> {
+  const allData = await redis.hgetall(HOURLY_ACTIVITY_KEY);
+  const now = new Date();
+  const result: Record<string, number> = {};
+
+  for (let i = 23; i >= 0; i--) {
+    const date = new Date(now.getTime() - i * 60 * 60 * 1000);
+    const hourKey = `${date.getUTCFullYear()}-${String(date.getUTCMonth() + 1).padStart(2, "0")}-${String(date.getUTCDate()).padStart(2, "0")}T${String(date.getUTCHours()).padStart(2, "0")}`;
+    result[hourKey] = parseInt(allData[hourKey] || "0", 10);
+  }
+
+  return result;
+}
+
+export async function setCacheHeatmap(data: string): Promise<void> {
+  await redis.set(HEATMAP_CACHE_KEY, data, "EX", 30);
+}
+
+export async function getCacheHeatmap(): Promise<string | null> {
+  return redis.get(HEATMAP_CACHE_KEY);
+}
+
+export async function publishEvent(event: object): Promise<void> {
+  await redis.publish(CHANNEL_REALTIME, JSON.stringify(event));
+}
--- a/lib/utils.ts
+++ b/lib/utils.ts
@@ -0,0 +1,6 @@
+import { type ClassValue, clsx } from "clsx";
+import { twMerge } from "tailwind-merge";
+
+export function cn(...inputs: ClassValue[]) {
+  return twMerge(clsx(inputs));
+}
--- a/lib/validators/schemas.ts
+++ b/lib/validators/schemas.ts
@@ -0,0 +1,24 @@
+import { z } from "zod";
+
+export const registerSchema = z.object({
+  name: z.string().min(1).max(100),
+  platform: z.string().optional(),
+  model: z.string().optional(),
+});
+
+export const heartbeatSchema = z.object({
+  name: z.string().optional(),
+  model: z.string().optional(),
+  platform: z.string().optional(),
+});
+
+export const taskSchema = z.object({
+  summary: z.string().max(500),
+  durationMs: z.number().positive(),
+  model: z.string().optional(),
+  toolsUsed: z.array(z.string()).optional(),
+});
+
+export type RegisterInput = z.infer<typeof registerSchema>;
+export type HeartbeatInput = z.infer<typeof heartbeatSchema>;
+export type TaskInput = z.infer<typeof taskSchema>;