AI Gateway Go 0.10.0 源码快照 + 旗舰版需求规划报告

M0-M7 已完成:核心网关(身份/RBAC/TOTP/OIDC/SAML/Provider/配额/路由/内容策略/审计/定价)+ 资源市场(MCP/Skills/数字员工)。
含 22 个 PostgreSQL 迁移、管理端/门户端前端源码、OpenAPI 契约、部署 compose。

Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
ben
2026-08-12 11:45:54 +08:00
commit 5759c1862e
807 changed files with 114727 additions and 0 deletions
+9
View File
@@ -0,0 +1,9 @@
.git
.env
deploy/production.env
bin
coverage.out
web/node_modules
web/**/node_modules
web/**/dist
web/.pnpm-store
+75
View File
@@ -0,0 +1,75 @@
APP_ENV=local
HTTP_ADDR=:8080
HTTP_READ_HEADER_TIMEOUT=5s
HTTP_IDLE_TIMEOUT=120s
HTTP_SHUTDOWN_TIMEOUT=20s
HTTP_MAX_BODY_BYTES=33554432
DATABASE_URL=postgres://gateway:gateway@localhost:5432/gateway?sslmode=disable
DATABASE_MAX_CONNS=40
DATABASE_MIN_CONNS=4
REDIS_CRITICAL_URL=redis://localhost:6379/0
REDIS_CACHE_URL=redis://localhost:6380/0
AUTH_SESSION_TTL=12h
AUTH_TOTP_CHALLENGE_TTL=5m
LOGIN_MAX_FAILURES=5
LOGIN_LOCK_DURATION=15m
# 登录防爆破:按 IP 的滑动窗口限流(与上面的账号锁定叠加)。
# 单 IP 在窗口内最多 LOGIN_RATE_LIMIT_MAX 次登录尝试,超出返回 429。
LOGIN_RATE_LIMIT_MAX=30
LOGIN_RATE_LIMIT_WINDOW=5m
# Base64-encoded 32-byte AES-256 key. Generate with: openssl rand -base64 32
CREDENTIAL_MASTER_KEY=
CREDENTIAL_KEK_VERSION=1
ALLOW_PRIVATE_PROVIDER_URLS=false
ALLOW_PRIVATE_TOOL_URLS=false
ALLOW_PRIVATE_WEBHOOK_URLS=false
PROVIDER_REFRESH_INTERVAL=5s
# Transitional bootstrap credential. It will be replaced by database-backed API keys.
# Disabled by default for security; it bypasses per-key rate limits and tenant
# isolation, so only enable it during migration with a strong random value.
GATEWAY_BOOTSTRAP_API_KEY_ENABLED=false
GATEWAY_BOOTSTRAP_API_KEY=
UPSTREAM_BASE_URL=https://api.openai.com
UPSTREAM_API_KEY=
UPSTREAM_FALLBACK_ENABLED=true
AUDIT_QUEUE_SIZE=4096
AUDIT_BATCH_SIZE=200
AUDIT_FLUSH_INTERVAL=1s
AUDIT_RETENTION=2160h
USAGE_RETENTION=17520h
AUDIT_PARTITION_MONTHS_AHEAD=3
AUDIT_MAINTENANCE_INTERVAL=6h
# Keep the {outbox} hash tag so marker and Stream keys share a Redis Cluster slot.
OUTBOX_STREAM=gateway:{outbox}:events
OUTBOX_BATCH_SIZE=100
OUTBOX_POLL_INTERVAL=500ms
OUTBOX_LEASE=30s
OUTBOX_MAX_ATTEMPTS=10
OUTBOX_MAX_BACKOFF=5m
OUTBOX_STREAM_MAX_LENGTH=100000
OUTBOX_MARKER_TTL=720h
# Multi-instance immutable snapshot refresh. Admin writes also reload the local process immediately.
CONTENT_POLICY_REFRESH_INTERVAL=30s
PRICING_REFRESH_INTERVAL=30s
# Opt-in non-streaming JSON shadow traffic. A dedicated credential is mandatory;
# production client Authorization headers are never forwarded to the shadow.
SHADOW_BASE_URL=
SHADOW_API_KEY=
SHADOW_SAMPLE_RATE=0
SHADOW_TIMEOUT=20s
SHADOW_MAX_BODY_BYTES=2097152
SHADOW_MAX_CONCURRENT=16
# Used only by cmd/gateway-bootstrap; never commit the real value.
BOOTSTRAP_ADMIN_USERNAME=admin
BOOTSTRAP_ADMIN_PASSWORD=
+9
View File
@@ -0,0 +1,9 @@
.env
deploy/production.env
.idea/
.vscode/
.DS_Store
bin/
coverage.out
web/node_modules/
web/**/dist/
+32
View File
@@ -0,0 +1,32 @@
.PHONY: fmt test vet build run migrate bootstrap compose-up compose-down
fmt:
gofmt -w $$(find cmd internal -name '*.go' -type f)
test:
go test ./...
vet:
go vet ./...
build:
go build -trimpath -o bin/gateway-api ./cmd/gateway-api
go build -trimpath -o bin/gateway-migrator ./cmd/gateway-migrator
go build -trimpath -o bin/gateway-bootstrap ./cmd/gateway-bootstrap
go build -trimpath -o bin/gateway-outbox-worker ./cmd/gateway-outbox-worker
go build -trimpath -o bin/gateway-maintenance ./cmd/gateway-maintenance
run:
go run ./cmd/gateway-api
migrate:
go run ./cmd/gateway-migrator
bootstrap:
go run ./cmd/gateway-bootstrap
compose-up:
docker compose -f deploy/docker-compose.yml up -d --build
compose-down:
docker compose -f deploy/docker-compose.yml down
+112
View File
@@ -0,0 +1,112 @@
# AI Gateway Go
AI Gateway 的全量 Go 重构工程。M0–M6 工程实现已完成,当前可运行基线包含 PostgreSQL、双 Redis、独立迁移器、Art Design Pro 管理端/门户端、OpenAI 兼容网关和可发布的 AI 资产编排运行时。
## 当前能力
- `GET /healthz`:进程存活检查。
- `GET /readyz`PostgreSQL、critical Redis 和 cache Redis 状态;前两者失败会返回 `503`
- `GET /metrics`:不依赖外部组件的基础 Prometheus 文本指标。
- `POST /v1/chat/completions``POST /v1/responses``POST /v1/embeddings``POST /v1/messages`:透明代理,支持 SSE。
- `GET /v1/models`:透明代理。
- 独立 `gateway-migrator`API 启动时不会修改数据库结构。
- PostgreSQL 权威管理员/门户身份、PBKDF2 旧密码兼容和 Redis 可撤销会话。
- 管理端与门户端 RFC 6238 TOTP 两步验证、5 分钟挑战令牌、时间步防重放和一次性备用码。
- Provider 管理 API;凭据使用 AES-256-GCM 加密,变更与 outbox 同事务提交。
- PostgreSQL Provider 原子运行时快照、默认/显式路由、能力校验和共享 HTTP 连接池;刷新失败保留上一有效版本。
- Provider 连接测试和模型目录同步;上游缺失模型保留历史记录并自动停用,同步事件与目录更新同事务提交。
- 管理员/门户账号管理、内置角色与可扩展 `resource:action` 权限字符串;账号和权限变更即时生效。
- 层级部门管理和门户账号部门绑定;循环层级、停用父部门和停用在用部门均由服务端约束。
- OIDC Authorization Code + PKCE 登录、RS256 ID Token 校验、外部账号映射与自动开户。
- SAML 2.0 SP 发起登录、SP metadata、签名断言验证、RelayState/Request ID 绑定、断言防重放与自动开户。
- 旧系统 38 个持久化实体的数据字典,以及幂等的旧 ID → UUIDv5 兼容映射基础设施。
- PostgreSQL API Key(明文仅显示一次、scope、过期时间、Redis 共享缓存、即时撤销)及 Art 管理页面。
- 按 API Key 配置的每分钟请求上限、UTC 自然月请求配额和 Token 配额;critical Redis Lua 原子执行,Token 采用请求前预留、JSON/SSE usage 回写校准,返回 429、`Retry-After` 和配额响应头。
- Provider 级响应头超时、幂等安全重试与熔断;POST 只有在客户端提供 `Idempotency-Key` 且请求体可重放时才会重试。
- PostgreSQL 模型别名和条件路由,按 endpoint、API Key、tenant 过滤,并在最高优先级组内做确定性加权 Provider 选择;规则通过不可变快照运行并由 Redis 通知多实例刷新。
- 有界异步调用审计和按日 usage 聚合:批量写 PostgreSQL,记录状态、延迟及输入/输出 Token;管理端提供权限隔离的筛选查询,审计存储异常不会阻塞模型调用。
- 独立事务 Outbox Worker:多实例 `SKIP LOCKED` 抢占、Redis Stream 原子去重投递、指数退避与死信;管理端可审阅并人工重试,消费者可用同事务幂等入口。
- PostgreSQL 审计月分区维护:自动预创建、default 分区在线迁移、过期整分区快速丢弃、边界精确清理和更长期的 usage 保留策略。
- 可扩展内容策略:Go RE2 不可变编译快照,按端点、模型/API Key 匹配,支持仅审计、阻断与提示词文本脱敏;默认保护常见 API Key、Token、密码和 secret。
- 带时间版本的模型价格与成本核算:按 Provider/模型选择价格,输入与输出 Token 分别计价,结果进入调用审计和 PostgreSQL 按日聚合。
- Prompt 分类、模板和不可变版本,支持显式变量定义、必填校验、历史版本激活与 API Key 渲染接口。
- PostgreSQL 知识库:2 MiB 有界文本正文、段落感知重叠分块、FTS + 中文二元词片混合检索,以及可替换的 `Retriever` 接口;不依赖对象存储或向量数据库。
- 声明式 HTTP 工具:JSON Schema 基础校验、KEK 加密请求头、注册和拨号双层 SSRF 防护、禁止重定向、1 MiB 响应限制与调用记录。
- AI 应用草稿和不可变发布版本,将模型、Prompt、知识库、工具组合为 `/v1/applications/{code}/chat/completions`;所有模型轮次继续经过鉴权、配额、内容策略、路由、成本和审计。
- 独立通知 Worker 消费可靠 outbox,按精确事件或末尾 `*` 模式投递 HMAC-SHA256 Webhook;内容策略命中由审计批处理异步产生脱敏事件,失败投递可在 Art 管理端重试。
- 门户自助工作台:部门范围资产目录、Prompt 搜索/收藏、个人审计/用量/成本、模型访问申请与管理员审批。
- 门户应用托管会话:服务端加密运行凭证、单会话租约、不可变消息序列和 SHA-256 哈希链,不向浏览器暴露应用 API Key。
- 独立事实核验配置、作用域策略与事件契约,复用 Provider 加密凭据和知识库引用,为同步/异步执行器保留清晰模块边界。
- 旧 Python 源码 201 条路由全部有覆盖、替代或退役决策,未决契约缺口为 0;OpenAPI 0.10.0 覆盖全部 Go 字面量路由。
MinIO/S3 和 ClickHouse 不属于基线部署,也不是启动依赖。审计与统计第一阶段存放在 PostgreSQL;对象存储和分析库仅保留后续 adapter 扩展点。
## 本地启动
1. 复制 `.env.example``.env` 并修改密钥。
2. 启动 PostgreSQL 和两套 Redis`docker compose -f deploy/docker-compose.yml up -d postgres redis-critical redis-cache`
3. 执行迁移:`go run ./cmd/gateway-migrator`
4. 设置 `BOOTSTRAP_ADMIN_PASSWORD` 后执行 `go run ./cmd/gateway-bootstrap` 创建初始管理员。
5. 启动 API`go run ./cmd/gateway-api`
6. 启动可靠事件投递:`go run ./cmd/gateway-outbox-worker`
7. 启动审计分区与保留维护:`go run ./cmd/gateway-maintenance`
8. 启动通知投递:`go run ./cmd/gateway-notification-worker`
完整容器部署可执行 `docker compose -f deploy/docker-compose.yml up -d --build`。Art Design Pro 管理端默认暴露在 `http://127.0.0.1:8081`,门户端在 `http://127.0.0.1:8082`;两者同源代理 `/api/*``/v1/*``/healthz``/readyz` 到 Go APIGo API 仍可从 `8080` 端口直接访问。
生产部署请使用 `deploy/docker-compose.production.yml`
`deploy/production.env.example`,完整步骤见
[`deploy/PRODUCTION.md`](deploy/PRODUCTION.md)。生产编排默认仅绑定回环地址、要求显式提供数据库/Redis/加密密钥,并关闭迁移期 bootstrap API key。
如果主机未安装 Go,可以在 Compose 依赖启动和迁移完成后执行:
```bash
docker compose -f deploy/docker-compose.yml run --rm \
-e BOOTSTRAP_ADMIN_PASSWORD='替换为强口令' \
--entrypoint gateway-bootstrap gateway-api
```
未设置 `HTTP_ADDR` 时只监听 `127.0.0.1:8080`;容器配置会显式监听 `:8080`
生产环境必须设置 `APP_ENV=production`,并使用带 TLS 的数据库和 Redis 地址。启用环境变量上游回退时必须设置 `UPSTREAM_API_KEY`。数据库 API Key 已可用于网关请求;bootstrap key 默认关闭,仅用于迁移兼容,需要时显式设置 `GATEWAY_BOOTSTRAP_API_KEY_ENABLED=true` 并配置强随机值,迁移结束后立即关闭,此时无需再配置 bootstrap key。`gateway_bootstrap_api_key_uses_total` 指标可用于确认旧客户端是否已经清零。bootstrap key 会绕过按 key 的限流与配额,切勿长期开启。
内置管理员角色为 `superadmin``operator``auditor`。超级管理员拥有通配权限;运维管理员默认管理 Provider 和 API Key;审计员只有读取权限。账号还可获得独立的 `resource:action` 权限,服务端每次请求都会从 PostgreSQL 复核账号状态和权限,因此停用或权限回收即时生效。
内容策略和模型价格以 PostgreSQL 为权威源,管理端保存后本实例立即刷新;其他实例分别按 `CONTENT_POLICY_REFRESH_INTERVAL``PRICING_REFRESH_INTERVAL`(默认 30 秒)加载不可变快照。正则采用 Go RE2,不支持回溯、环视等可能导致灾难性耗时的表达式。价格金额以微货币单位整数保存,避免浮点累计误差;例如 USD 2.50/百万 Token 存为 `2500000`
部门采用稳定 UUID 与邻接表层级。部门仍包含启用门户用户或启用子部门时不能停用;移动部门时会在事务内检测自身引用和任意深度循环。门户账号只能绑定启用部门,解绑不会删除任何历史身份数据。
OIDC 登录使用 Redis 保存 5 分钟一次性 `state``nonce` 和 PKCE verifier;回调校验发现文档、JWKS/RS256 签名、issuer、audience/azp、iat/nbf/exp 与 nonce。门户会话令牌不进入回调 URL,只使用 60 秒一次性交换码。OIDC Client Secret 使用 `identity-provider-credentials` 独立 AEAD 用途加密。
SAML 当前只开放 SP-initiated Redirect 登录和 POST ACS,显式拒绝 IdP-initiated 与 Artifact 绑定。IdP metadata 通过 SSRF 安全客户端获取并短时缓存,启用前校验 Entity ID、Redirect SSO 端点、签名证书和过期时间。ACS 依次验证 XML 签名、issuer、audience、destination、recipient、InResponseTo 和时间窗,只接受 SHA-256 及以上的签名/摘要算法,并使用 Redis 防止 RelayState 和 Assertion ID 重放。
登录防爆破采用两层叠加:账号锁定(`LOGIN_MAX_FAILURES` 次失败后锁定 `LOGIN_LOCK_DURATION`,针对单账号)与按 IP 的 Redis 滑动窗口限流(`LOGIN_RATE_LIMIT_MAX` 次/`LOGIN_RATE_LIMIT_WINDOW`,超出返回 429,针对跨账号撞库)。限流在 Redis 不可用时 fail-open,账号锁定仍然生效。来源 IP 取自 `X-Forwarded-For``deploy/nginx-web.conf``/api/` 路径用 `$remote_addr` 覆盖该头,避免客户端伪造头绕过限流;直接访问网关端口绕过 nginx 的请求仍可伪造该头,因此生产建议将管理面代理收敛在受信网络内。
旧 Python 数据结构的完整字段字典见 `docs/legacy-data-dictionary.md`。迁移使用 `gateway.legacy_id_mappings` 保存来源、实体、旧 ID 与新 UUID 的对应关系;新 UUID 由固定 namespace 的 UUIDv5 生成,同一条旧数据可安全重跑而不会生成不同主键。字段字典可用 `scripts/export_legacy_dictionary.py` 从旧工程重新生成。
`CREDENTIAL_MASTER_KEY` 必须是 32 字节密钥的 Base64 编码。Compose 中的默认值只允许本地开发,生产环境必须替换并纳入密钥托管与备份;丢失该密钥会导致 Provider 凭据、TOTP、工具请求头、通知签名密钥和应用运行凭证无法解密。每类密文使用独立用途标签,不能相互替换。
工具端点和通知 Webhook 默认只允许公网 HTTP(S) 地址,并在实际拨号时重新解析与校验地址。确有内网服务时分别显式设置 `ALLOW_PRIVATE_TOOL_URLS=true``ALLOW_PRIVATE_WEBHOOK_URLS=true`;这两个开关与 Provider 私网开关相互独立。
影子流量默认关闭。设置 `SHADOW_BASE_URL`、专用 `SHADOW_API_KEY` 和大于 0 的 `SHADOW_SAMPLE_RATE` 后,只复制确定性采样的非流式 JSON POST;生产客户端凭据不会转发。主响应不等待影子请求,Prometheus 指标只比较 HTTP 状态和 JSON 结构签名,不记录提示词或模型正文。压测使用独立 `gateway-loadtest` 二进制,可用错误率和 p95 阈值直接控制退出码。
旧路由矩阵由 `scripts/compare_route_contracts.py` 从 Python AST 与 Go OpenAPI 生成。旧源码当前实际包含 201 个路由装饰器,并非早期估算的 141 个;矩阵明确区分同契约覆盖、新契约替代、退役和真实缺口。切换步骤见 `docs/cutover-runbook.md`,旧数据暂存与密文边界见 `docs/legacy-import-runbook.md`
KEK 轮换时,将 `CREDENTIAL_KEK_VERSION``CREDENTIAL_MASTER_KEY` 设置为新活动版本与新密钥,并用 `CREDENTIAL_KEK_KEYRING` 加载历史密钥,例如 `{"1":"<old-base64-key>"}`。服务重启并确认可解密后,在供应商页面执行“轮换凭据”;所有旧版本 Provider 凭据和对应 outbox 事件会在一个事务内提交。确认全部 Provider 已切换后才能从 keyring 移除旧密钥。TOTP 同样通过 keyring 保持历史版本可解密,新配置会使用活动版本。
Provider 以 PostgreSQL 为权威源,在进程内构建不可变快照。管理 API 修改后会刷新本实例并通过 critical Redis Pub/Sub 通知其他实例立即刷新;默认每 5 秒轮询仍作为通知丢失或 Redis 短暂异常时的兜底。请求可通过 `X-Gateway-Provider` 指定 Provider;未指定时使用 `config.default=true` 的启用项,否则使用按 code 排序的第一项。`UPSTREAM_FALLBACK_ENABLED=false` 可关闭环境变量上游回退。
供应商连接测试和模型同步访问 OpenAI-compatible `/v1/models`。控制面请求限制为 10 秒、4 MiB 响应,拒绝重定向,并在禁用私网 Provider 时对实际拨号地址再次执行 SSRF 校验。模型同步不会物理删除旧模型,本次未出现的模型会标记为停用,为后续别名和路由规则保留引用稳定性。
## 目录
- `api/openapi`:对外契约。
- `cmd`:可部署二进制。
- `internal/platform`:配置、数据库、缓存、HTTP 运行时。
- `internal/provider`:上游 Provider 扩展接口及实现。
- `internal/gateway`:协议入口和代理编排。
- `internal/workbench`:Prompt、知识、工具、应用编排和通知。
- `internal/portal`:门户目录、申请、个人统计、加密应用凭证与托管会话。
- `internal/factcheck`:事实核验设置、策略和事件契约。
- `migrations`:仅由 migrator 执行的 PostgreSQL 迁移。
- `web`Art Design Pro 管理端和门户端。
+20
View File
@@ -0,0 +1,20 @@
# AI Gateway Go 0.10.0 clean deployment bundle
Generated on 2026-08-11 from `/home/ben/ai-gateway-go`.
Included:
- Go backend and worker source
- 21 PostgreSQL migrations
- Admin and portal Art Design Pro source
- OpenAPI contract, operations documents and legacy migration tools
- Local and production Docker Compose definitions
Excluded:
- `node_modules`, frontend `dist`, Go binaries and caches
- local `.env`, `deploy/production.env`, logs and coverage output
- database volumes, Redis data and all runtime credentials
Start with `deploy/PRODUCTION.md`. Verify the accompanying ZIP checksum before
copying the bundle to a deployment host.
File diff suppressed because it is too large Load Diff
+383
View File
@@ -0,0 +1,383 @@
package main
import (
"context"
"errors"
"log/slog"
"net/http"
"os"
"os/signal"
"syscall"
"time"
"aigateway.local/core/internal/apikey"
"aigateway.local/core/internal/audit"
"aigateway.local/core/internal/contentpolicy"
"aigateway.local/core/internal/factcheck"
"aigateway.local/core/internal/gateway"
"aigateway.local/core/internal/identity"
"aigateway.local/core/internal/operations"
"aigateway.local/core/internal/outbox"
"aigateway.local/core/internal/platform/cache"
"aigateway.local/core/internal/platform/config"
"aigateway.local/core/internal/platform/cryptox"
"aigateway.local/core/internal/platform/database"
"aigateway.local/core/internal/platform/health"
"aigateway.local/core/internal/platform/httpserver"
"aigateway.local/core/internal/portal"
"aigateway.local/core/internal/pricing"
"aigateway.local/core/internal/provider"
providercontrolplane "aigateway.local/core/internal/provider/controlplane"
provideropenai "aigateway.local/core/internal/provider/openai"
providerruntime "aigateway.local/core/internal/provider/runtime"
"aigateway.local/core/internal/shadow"
"aigateway.local/core/internal/workbench"
)
var version = "dev"
func main() {
logger := slog.New(slog.NewJSONHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelInfo}))
cfg, err := config.Load()
if err != nil {
logger.Error("invalid configuration", "error", err)
os.Exit(1)
}
if err := cfg.ValidateRuntime(); err != nil {
logger.Error("invalid runtime configuration", "error", err)
os.Exit(1)
}
ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
defer stop()
db, err := database.Open(ctx, cfg.Database)
if err != nil {
logger.Error("database initialization failed", "error", err)
os.Exit(1)
}
if db != nil {
defer db.Close()
}
criticalRedis, err := cache.Open(cfg.Redis.CriticalURL)
if err != nil {
logger.Error("critical redis initialization failed", "error", err)
os.Exit(1)
}
if criticalRedis != nil {
defer criticalRedis.Close()
}
cacheRedis, err := cache.Open(cfg.Redis.CacheURL)
if err != nil {
logger.Error("cache redis initialization failed", "error", err)
os.Exit(1)
}
if cacheRedis != nil {
defer cacheRedis.Close()
}
checker := health.Checker{Timeout: 1500 * time.Millisecond}
checker.Dependencies = append(checker.Dependencies,
health.Dependency{Name: "postgres", Required: true, Probe: func(probeCtx context.Context) error {
if db == nil {
return errors.New("not configured")
}
return db.Ping(probeCtx)
}},
health.Dependency{Name: "redis_critical", Required: true, Probe: func(probeCtx context.Context) error {
if criticalRedis == nil {
return errors.New("not configured")
}
return criticalRedis.Ping(probeCtx).Err()
}},
health.Dependency{Name: "redis_cache", Required: false, Probe: func(probeCtx context.Context) error {
if cacheRedis == nil {
return errors.New("not configured")
}
return cacheRedis.Ping(probeCtx).Err()
}},
)
adapter, err := provideropenai.New(cfg.Upstream.BaseURL, cfg.Upstream.APIKey)
if err != nil {
logger.Error("provider initialization failed", "error", err)
os.Exit(1)
}
var fallbackAdapter provider.Adapter
if cfg.Upstream.FallbackEnabled {
fallbackAdapter = adapter
}
credentialCipher, err := provider.NewCredentialCipher(
cfg.Credentials.MasterKey, cfg.Credentials.KEKVersion, cfg.Credentials.KEKKeyring,
)
if err != nil {
logger.Error("credential encryption initialization failed", "error", err)
os.Exit(1)
}
providerRepository := provider.NewRepository(db)
providerResolver := providerruntime.NewResolver(
providerRepository, credentialCipher, fallbackAdapter, cfg.Credentials.ProviderRefreshInterval, logger,
)
providerResolver.SetNotificationClient(criticalRedis)
go providerResolver.Run(ctx)
apiKeyRepository := apikey.NewRepository(db)
bootstrapAPIKey := ""
if cfg.Security.BootstrapAPIKeyEnabled {
bootstrapAPIKey = cfg.Security.BootstrapAPIKey
logger.Warn("bootstrap API key compatibility is enabled; monitor usage and disable after migration")
}
apiKeyAuthenticator := apikey.NewAuthenticator(apiKeyRepository, criticalRedis, bootstrapAPIKey)
apiKeyAuthenticator.SetLogger(logger)
proxy := gateway.NewDynamicProxy(providerResolver, apiKeyAuthenticator, cfg.Server.MaxBodyBytes, logger)
proxy.SetAdmissionController(gateway.NewRedisAdmissionController(criticalRedis))
proxy.SetTokenQuotaController(gateway.NewRedisTokenQuotaController(criticalRedis))
proxy.SetResiliencePolicy(gateway.ResiliencePolicy{
ResponseHeaderTimeout: cfg.Upstream.ResponseHeaderTimeout, MaxRetries: cfg.Upstream.MaxRetries,
RetryBackoff: cfg.Upstream.RetryBackoff, CircuitThreshold: cfg.Upstream.CircuitThreshold,
CircuitOpenDuration: cfg.Upstream.CircuitOpenDuration,
})
auditRecorder := audit.NewRecorder(db, logger, cfg.Audit.QueueSize, cfg.Audit.BatchSize, cfg.Audit.FlushInterval)
auditContext, stopAudit := context.WithCancel(context.Background())
auditStopped := make(chan struct{})
go func() {
auditRecorder.Run(auditContext)
close(auditStopped)
}()
proxy.SetAuditRecorder(auditRecorder)
contentPolicyEngine := contentpolicy.NewEngine(db, cfg.RuntimeData.ContentPolicyRefreshInterval, logger)
pricingService := pricing.NewService(db, cfg.RuntimeData.PricingRefreshInterval, logger)
if err := contentPolicyEngine.Reload(ctx); err != nil {
logger.Error("content policy initialization failed", "error", err)
os.Exit(1)
}
if err := pricingService.Reload(ctx); err != nil {
logger.Error("model pricing initialization failed", "error", err)
os.Exit(1)
}
go contentPolicyEngine.Run(ctx)
go pricingService.Run(ctx)
proxy.SetContentPolicyEngine(contentPolicyEngine)
proxy.SetPricingService(pricingService)
identityRepository := identity.NewRepository(db)
sessionStore := identity.NewSessionStore(criticalRedis, cfg.Auth.SessionTTL)
loginLimiter := identity.NewLoginLimiter(criticalRedis, cfg.Auth.LoginRateLimitMax, cfg.Auth.LoginRateLimitWindow)
totpCipher, err := cryptox.NewKeyring(
cfg.Credentials.MasterKey, cfg.Credentials.KEKVersion, cfg.Credentials.KEKKeyring, "totp-secret",
)
if err != nil {
logger.Error("TOTP encryption initialization failed", "error", err)
os.Exit(1)
}
identityService := identity.NewService(identityRepository, sessionStore, loginLimiter, cfg.Auth, totpCipher)
idpCipher, err := cryptox.NewKeyring(
cfg.Credentials.MasterKey, cfg.Credentials.KEKVersion, cfg.Credentials.KEKKeyring, "identity-provider-credentials",
)
if err != nil {
logger.Error("identity provider encryption initialization failed", "error", err)
os.Exit(1)
}
identityService.SetIdentityProviderCipher(idpCipher, cfg.Credentials.AllowPrivateProviderURL)
identityHandler := identity.NewHTTPHandler(identityService)
identityManagementHandler := identity.NewManagementHTTPHandler(identityService)
providerHandler := provider.NewAdminHTTPHandler(
providerRepository, credentialCipher, identityService, cfg.Credentials.AllowPrivateProviderURL,
)
providerOperations := providercontrolplane.NewService(
providerRepository, credentialCipher, cfg.Credentials.AllowPrivateProviderURL,
)
providerHandler.SetOperations(providerOperations)
providerHandler.SetChangeHook(func(changeCtx context.Context) error {
reloadErr := providerResolver.Reload(changeCtx)
notifyErr := providerResolver.Notify(changeCtx)
if reloadErr != nil || notifyErr != nil {
logger.Warn("provider change propagation was incomplete", "reload_error", reloadErr, "notify_error", notifyErr)
}
return errors.Join(reloadErr, notifyErr)
})
apiKeyHandler := apikey.NewAdminHTTPHandler(apiKeyRepository, apiKeyAuthenticator, identityService)
apiKeyHandler.SetUsageStore(apikey.NewUsageStore(criticalRedis))
auditHandler := audit.NewAdminHTTPHandler(audit.NewQueryService(db), identityService)
outboxHandler := outbox.NewAdminHTTPHandler(outbox.NewStore(db), identityService)
contentPolicyHandler := contentpolicy.NewAdminHTTPHandler(contentpolicy.NewStore(db), contentPolicyEngine, identityService)
pricingHandler := pricing.NewAdminHTTPHandler(pricingService, identityService)
factCheckHandler := factcheck.NewAdminHTTPHandler(factcheck.NewService(db), identityService)
workbenchService := workbench.NewService(db)
toolCipher, err := cryptox.NewKeyring(
cfg.Credentials.MasterKey, cfg.Credentials.KEKVersion, cfg.Credentials.KEKKeyring, "tool-request-headers",
)
if err != nil {
logger.Error("tool credential encryption initialization failed", "error", err)
os.Exit(1)
}
notificationCipher, err := cryptox.NewKeyring(
cfg.Credentials.MasterKey, cfg.Credentials.KEKVersion, cfg.Credentials.KEKKeyring, "notification-signing-secret",
)
if err != nil {
logger.Error("notification encryption initialization failed", "error", err)
os.Exit(1)
}
toolService := workbench.NewToolService(workbenchService, toolCipher, cfg.Credentials.AllowPrivateToolURL)
notificationService := workbench.NewNotificationService(workbenchService, notificationCipher, cfg.Credentials.AllowPrivateWebhookURL)
workbenchHandler := workbench.NewAdminHTTPHandler(workbenchService, toolService, notificationService, identityService)
mcpServerCipher, err := cryptox.NewKeyring(
cfg.Credentials.MasterKey, cfg.Credentials.KEKVersion, cfg.Credentials.KEKKeyring, "mcp-server-headers",
)
if err != nil {
logger.Error("MCP server credential encryption initialization failed", "error", err)
os.Exit(1)
}
mcpServerService := workbench.NewMCPServerService(workbenchService, mcpServerCipher, cfg.Credentials.AllowPrivateToolURL)
skillService := workbench.NewSkillService(workbenchService)
digitalEmployeeService := workbench.NewDigitalEmployeeService(workbenchService, skillService, toolService, mcpServerService)
marketplaceService := workbench.NewMarketplaceService(workbenchService, mcpServerService, skillService, digitalEmployeeService)
mcpClient := workbench.NewMCPClient(cfg.Credentials.AllowPrivateToolURL, 60*time.Second)
marketplaceHandler := workbench.NewMarketplaceAdminHTTPHandler(marketplaceService, mcpServerService, skillService, digitalEmployeeService, mcpClient, identityService)
applicationKeyCipher, err := cryptox.NewKeyring(
cfg.Credentials.MasterKey, cfg.Credentials.KEKVersion, cfg.Credentials.KEKKeyring, "application-runtime-key",
)
if err != nil {
logger.Error("application runtime credential initialization failed", "error", err)
os.Exit(1)
}
shadowMiddleware := shadow.New(cfg.Shadow, logger)
governedGateway := shadowMiddleware.Wrap(proxy)
workbenchRuntime := workbench.NewRuntimeHTTPHandler(workbenchService, toolService, workbench.NewPostgreSQLRetriever(workbenchService), apiKeyAuthenticator, governedGateway, workbench.MarketplaceDeps{
MCPServers: mcpServerService,
Skills: skillService,
Employees: digitalEmployeeService,
Market: marketplaceService,
MCPClient: mcpClient,
})
workbenchRuntime.SetLogger(logger)
// Wire the fact-check engine: the admin fact-check settings/policies UI now
// actually governs application answers instead of being inert configuration.
factCheckEngine := factcheck.NewEngine(db, workbench.NewFactCheckRetriever(workbench.NewPostgreSQLRetriever(workbenchService)), logger)
workbenchRuntime.SetFactCheckEngine(factCheckEngine)
portalService := portal.NewService(db, workbenchService, toolService, identityService)
portalService.SetApplicationRuntime(portal.NewRuntimeCredentials(db, apiKeyRepository, applicationKeyCipher), workbenchRuntime)
portalService.SetMarketplace(marketplaceService)
portalHandler := portal.NewHTTPHandler(portalService, identityService)
portalAdminHandler := portal.NewAdminHTTPHandler(portalService, identityService)
startedAt := time.Now()
operationsHandler := operations.NewAdminHTTPHandler(db, identityService, version, startedAt, func(reloadCtx context.Context) error {
return errors.Join(providerResolver.Reload(reloadCtx), contentPolicyEngine.Reload(reloadCtx), pricingService.Reload(reloadCtx))
})
controlMux := http.NewServeMux()
controlMux.Handle("/api/v1/admin/providers", providerHandler)
controlMux.Handle("/api/v1/admin/providers/", providerHandler)
controlMux.Handle("/api/v1/admin/model-routes", providerHandler)
controlMux.Handle("/api/v1/admin/model-routes/", providerHandler)
controlMux.Handle("/api/v1/admin/api-keys", apiKeyHandler)
controlMux.Handle("/api/v1/admin/api-keys/", apiKeyHandler)
controlMux.Handle("/api/v1/admin/audit-events", auditHandler)
controlMux.Handle("/api/v1/admin/usage/", auditHandler)
controlMux.Handle("/api/v1/admin/outbox-events", outboxHandler)
controlMux.Handle("/api/v1/admin/outbox-events/", outboxHandler)
controlMux.Handle("/api/v1/admin/content-policies", contentPolicyHandler)
controlMux.Handle("/api/v1/admin/content-policies/", contentPolicyHandler)
controlMux.Handle("/api/v1/admin/model-prices", pricingHandler)
controlMux.Handle("/api/v1/admin/model-prices/", pricingHandler)
controlMux.Handle("/api/v1/admin/fact-check/", factCheckHandler)
controlMux.Handle("/api/v1/admin/prompt-categories", workbenchHandler)
controlMux.Handle("/api/v1/admin/prompt-categories/", workbenchHandler)
controlMux.Handle("/api/v1/admin/prompts", workbenchHandler)
controlMux.Handle("/api/v1/admin/prompts/", workbenchHandler)
controlMux.Handle("/api/v1/admin/knowledge-bases", workbenchHandler)
controlMux.Handle("/api/v1/admin/knowledge-bases/", workbenchHandler)
controlMux.Handle("/api/v1/admin/tools", workbenchHandler)
controlMux.Handle("/api/v1/admin/tools/", workbenchHandler)
controlMux.Handle("/api/v1/admin/applications", workbenchHandler)
controlMux.Handle("/api/v1/admin/applications/", workbenchHandler)
controlMux.Handle("/api/v1/admin/marketplace-categories", marketplaceHandler)
controlMux.Handle("/api/v1/admin/marketplace-categories/", marketplaceHandler)
controlMux.Handle("/api/v1/admin/mcp-servers", marketplaceHandler)
controlMux.Handle("/api/v1/admin/mcp-servers/", marketplaceHandler)
controlMux.Handle("/api/v1/admin/skills", marketplaceHandler)
controlMux.Handle("/api/v1/admin/skills/", marketplaceHandler)
controlMux.Handle("/api/v1/admin/digital-employees", marketplaceHandler)
controlMux.Handle("/api/v1/admin/digital-employees/", marketplaceHandler)
controlMux.Handle("/api/v1/admin/marketplace/", marketplaceHandler)
controlMux.Handle("/api/v1/portal/marketplace", portalHandler)
controlMux.Handle("/api/v1/portal/marketplace/", portalHandler)
controlMux.Handle("/api/v1/admin/notification-channels", workbenchHandler)
controlMux.Handle("/api/v1/admin/notification-channels/", workbenchHandler)
controlMux.Handle("/api/v1/admin/notification-deliveries", workbenchHandler)
controlMux.Handle("/api/v1/admin/notification-deliveries/", workbenchHandler)
controlMux.Handle("/api/v1/admin/models", portalAdminHandler)
controlMux.Handle("/api/v1/admin/model-requests", portalAdminHandler)
controlMux.Handle("/api/v1/admin/model-requests/", portalAdminHandler)
controlMux.Handle("/api/v1/admin/system-info", operationsHandler)
controlMux.Handle("/api/v1/admin/monitoring/overview", operationsHandler)
controlMux.Handle("/api/v1/admin/reload", operationsHandler)
controlMux.Handle("/api/v1/admin/identities/", identityManagementHandler)
controlMux.Handle("/api/v1/admin/departments", identityManagementHandler)
controlMux.Handle("/api/v1/admin/departments/", identityManagementHandler)
controlMux.Handle("/api/v1/admin/identity-providers", identityManagementHandler)
controlMux.Handle("/api/v1/admin/identity-providers/", identityManagementHandler)
controlMux.Handle("/api/v1/admin/saml-providers", identityManagementHandler)
controlMux.Handle("/api/v1/admin/saml-providers/", identityManagementHandler)
controlMux.Handle("/api/v1/portal/applications", portalHandler)
controlMux.Handle("/api/v1/portal/apps/", portalHandler)
controlMux.Handle("/api/v1/portal/catalog", portalHandler)
controlMux.Handle("/api/v1/portal/cost", portalHandler)
controlMux.Handle("/api/v1/portal/docs-info", portalHandler)
controlMux.Handle("/api/v1/portal/knowledge", portalHandler)
controlMux.Handle("/api/v1/portal/logs", portalHandler)
controlMux.Handle("/api/v1/portal/logs/", portalHandler)
controlMux.Handle("/api/v1/portal/model-requests", portalHandler)
controlMux.Handle("/api/v1/portal/model-requests/", portalHandler)
controlMux.Handle("/api/v1/portal/password", portalHandler)
controlMux.Handle("/api/v1/portal/prompts", portalHandler)
controlMux.Handle("/api/v1/portal/prompts/", portalHandler)
controlMux.Handle("/api/v1/portal/stats", portalHandler)
controlMux.Handle("/api/v1/portal/tools", portalHandler)
controlMux.Handle("/api/v1/", identityHandler)
publicMux := http.NewServeMux()
publicMux.Handle("/v1/prompts", workbenchRuntime)
publicMux.Handle("/v1/prompts/", workbenchRuntime)
publicMux.Handle("/v1/knowledge/", workbenchRuntime)
publicMux.Handle("/v1/tools", workbenchRuntime)
publicMux.Handle("/v1/tools/", workbenchRuntime)
publicMux.Handle("/v1/applications/", workbenchRuntime)
publicMux.Handle("/v1/skills/", workbenchRuntime)
publicMux.Handle("/v1/mcp-servers", workbenchRuntime)
publicMux.Handle("/v1/mcp-servers/", workbenchRuntime)
publicMux.Handle("/v1/digital-employees/", workbenchRuntime)
publicMux.Handle("/v1/", governedGateway)
server := httpserver.New(httpserver.Dependencies{
Config: cfg, Logger: logger, Checker: checker, Gateway: publicMux, Control: controlMux,
Version: version, StartedAt: startedAt,
BootstrapUses: apiKeyAuthenticator.BootstrapUses,
ExtraMetrics: func() string { return auditRecorder.Prometheus() + shadowMiddleware.Prometheus() },
})
serverErrors := make(chan error, 1)
go func() {
logger.Info("gateway API started", "address", server.Addr, "version", version)
serverErrors <- server.ListenAndServe()
}()
select {
case <-ctx.Done():
logger.Info("shutdown requested")
case err := <-serverErrors:
if !errors.Is(err, http.ErrServerClosed) {
logger.Error("gateway API stopped unexpectedly", "error", err)
os.Exit(1)
}
}
shutdownCtx, cancel := context.WithTimeout(context.Background(), cfg.Server.ShutdownTimeout)
defer cancel()
if err := server.Shutdown(shutdownCtx); err != nil {
logger.Error("graceful shutdown failed", "error", err)
_ = server.Close()
os.Exit(1)
}
stopAudit()
select {
case <-auditStopped:
case <-shutdownCtx.Done():
logger.Warn("audit recorder drain timed out")
}
logger.Info("gateway API stopped")
}
+84
View File
@@ -0,0 +1,84 @@
package main
import (
"context"
"flag"
"log/slog"
"os"
"strings"
"aigateway.local/core/internal/identity"
"aigateway.local/core/internal/platform/config"
"aigateway.local/core/internal/platform/database"
)
func main() {
adminUsername := flag.String("admin-username", env("BOOTSTRAP_ADMIN_USERNAME", "admin"), "initial administrator username")
adminDisplayName := flag.String("admin-display-name", "系统管理员", "initial administrator display name")
portalAccount := flag.String("portal-account", strings.TrimSpace(os.Getenv("BOOTSTRAP_PORTAL_ACCOUNT")), "optional initial portal account")
portalName := flag.String("portal-name", "初始用户", "initial portal user display name")
flag.Parse()
logger := slog.New(slog.NewJSONHandler(os.Stdout, nil))
cfg, err := config.Load()
if err != nil {
logger.Error("invalid configuration", "error", err)
os.Exit(1)
}
if cfg.Database.URL == "" {
logger.Error("DATABASE_URL is required")
os.Exit(1)
}
adminPassword := os.Getenv("BOOTSTRAP_ADMIN_PASSWORD")
portalPassword := os.Getenv("BOOTSTRAP_PORTAL_PASSWORD")
if len(adminPassword) < 12 {
logger.Error("BOOTSTRAP_ADMIN_PASSWORD must contain at least 12 characters")
os.Exit(1)
}
if *portalAccount != "" && len(portalPassword) < 12 {
logger.Error("BOOTSTRAP_PORTAL_PASSWORD must contain at least 12 characters when creating a portal user")
os.Exit(1)
}
ctx := context.Background()
pool, err := database.Open(ctx, cfg.Database)
if err != nil {
logger.Error("database initialization failed", "error", err)
os.Exit(1)
}
defer pool.Close()
repository := identity.NewRepository(pool)
hasher := identity.PasswordHasher{}
adminHash, err := hasher.Hash(adminPassword)
if err != nil {
logger.Error("password hashing failed", "error", err)
os.Exit(1)
}
adminID, err := repository.CreateAdmin(ctx, *adminUsername, *adminDisplayName, "superadmin", adminHash)
if err != nil {
logger.Error("administrator bootstrap failed", "error", err)
os.Exit(1)
}
logger.Info("administrator created", "id", adminID, "username", strings.ToLower(strings.TrimSpace(*adminUsername)))
if *portalAccount != "" {
portalHash, err := hasher.Hash(portalPassword)
if err != nil {
logger.Error("portal password hashing failed", "error", err)
os.Exit(1)
}
portalID, err := repository.CreatePortalUser(ctx, *portalAccount, *portalName, portalHash)
if err != nil {
logger.Error("portal user bootstrap failed", "error", err)
os.Exit(1)
}
logger.Info("portal user created", "id", portalID, "account", strings.ToLower(strings.TrimSpace(*portalAccount)))
}
}
func env(key, fallback string) string {
if value := strings.TrimSpace(os.Getenv(key)); value != "" {
return value
}
return fallback
}
+178
View File
@@ -0,0 +1,178 @@
package main
import (
"bufio"
"context"
"crypto/sha256"
"encoding/hex"
"encoding/json"
"errors"
"flag"
"fmt"
"io"
"os"
"strings"
"aigateway.local/core/internal/platform/config"
"aigateway.local/core/internal/platform/database"
platformid "aigateway.local/core/internal/platform/id"
"aigateway.local/core/internal/platform/legacyid"
"github.com/jackc/pgx/v5/pgxpool"
)
type record struct {
SourceSystem string `json:"source_system"`
EntityType string `json:"entity_type"`
LegacyID string `json:"legacy_id"`
Data json.RawMessage `json:"data"`
Checksum string `json:"checksum"`
}
type staged struct {
record
NewID string
}
func main() {
input := flag.String("input", "-", "JSONL file from export_legacy_data.py, or - for stdin")
dryRun := flag.Bool("dry-run", false, "validate without database writes")
flag.Parse()
reader, closeInput, err := openInput(*input)
if err != nil {
fatal(err)
}
defer closeInput()
records, sourceChecksum, counts, err := readRecords(reader)
if err != nil {
fatal(err)
}
summary := map[string]any{"records": len(records), "entities": counts, "source_checksum": sourceChecksum, "dry_run": *dryRun}
if *dryRun {
write(summary)
return
}
cfg, err := config.Load()
if err != nil {
fatal(err)
}
if cfg.Database.URL == "" {
fatal(errors.New("DATABASE_URL is required"))
}
ctx := context.Background()
pool, err := database.Open(ctx, cfg.Database)
if err != nil {
fatal(err)
}
defer pool.Close()
batchID, err := stage(ctx, pool, records, sourceChecksum, counts)
if err != nil {
fatal(err)
}
summary["batch_id"] = batchID
summary["status"] = "staged"
write(summary)
}
func openInput(path string) (io.Reader, func(), error) {
if path == "-" {
return os.Stdin, func() {}, nil
}
file, err := os.Open(path)
if err != nil {
return nil, func() {}, err
}
return file, func() { _ = file.Close() }, nil
}
func readRecords(reader io.Reader) ([]staged, string, map[string]int, error) {
scanner := bufio.NewScanner(reader)
scanner.Buffer(make([]byte, 64<<10), 32<<20)
aggregate := sha256.New()
items := []staged{}
counts := map[string]int{}
source := ""
line := 0
for scanner.Scan() {
line++
raw := bytesTrimSpace(scanner.Bytes())
if len(raw) == 0 {
continue
}
_, _ = aggregate.Write(raw)
_, _ = aggregate.Write([]byte{'\n'})
var item record
if err := json.Unmarshal(raw, &item); err != nil {
return nil, "", nil, fmt.Errorf("line %d: %w", line, err)
}
item.SourceSystem = strings.TrimSpace(item.SourceSystem)
item.EntityType = strings.TrimSpace(item.EntityType)
item.LegacyID = strings.TrimSpace(item.LegacyID)
if source == "" {
source = item.SourceSystem
}
if item.SourceSystem != source {
return nil, "", nil, fmt.Errorf("line %d: mixed source systems", line)
}
canonical, _ := json.Marshal(map[string]any{"source_system": item.SourceSystem, "entity_type": item.EntityType, "legacy_id": item.LegacyID, "data": item.Data})
digest := sha256.Sum256(canonical)
if hex.EncodeToString(digest[:]) != item.Checksum {
return nil, "", nil, fmt.Errorf("line %d: checksum mismatch", line)
}
newID, err := legacyid.UUID(item.SourceSystem, item.EntityType, item.LegacyID)
if err != nil {
return nil, "", nil, fmt.Errorf("line %d: %w", line, err)
}
items = append(items, staged{record: item, NewID: newID})
counts[item.EntityType]++
}
if err := scanner.Err(); err != nil {
return nil, "", nil, err
}
if len(items) == 0 {
return nil, "", nil, errors.New("input contains no records")
}
return items, hex.EncodeToString(aggregate.Sum(nil)), counts, nil
}
func stage(ctx context.Context, pool *pgxpool.Pool, items []staged, sourceChecksum string, counts map[string]int) (string, error) {
batchID, err := platformid.NewUUID()
if err != nil {
return "", err
}
tx, err := pool.Begin(ctx)
if err != nil {
return "", err
}
defer func() { _ = tx.Rollback(ctx) }()
countsJSON, _ := json.Marshal(counts)
source := items[0].SourceSystem
err = tx.QueryRow(ctx, `INSERT INTO gateway.legacy_import_batches(id,source_system,source_checksum,status,record_count,entity_counts) VALUES($1,$2,$3,'staged',$4,$5) ON CONFLICT(source_system,source_checksum) DO UPDATE SET source_checksum=excluded.source_checksum RETURNING id::text`, batchID, source, sourceChecksum, len(items), countsJSON).Scan(&batchID)
if err != nil {
return "", err
}
for _, item := range items {
tag, insertErr := tx.Exec(ctx, `INSERT INTO gateway.legacy_import_records(first_seen_batch_id,source_system,entity_type,legacy_id,new_id,payload,payload_checksum) VALUES($1,$2,$3,$4,$5,$6,$7) ON CONFLICT(source_system,entity_type,legacy_id) DO UPDATE SET payload_checksum=gateway.legacy_import_records.payload_checksum WHERE gateway.legacy_import_records.payload_checksum=excluded.payload_checksum`, batchID, item.SourceSystem, item.EntityType, item.LegacyID, item.NewID, item.Data, item.Checksum)
if insertErr != nil {
return "", insertErr
}
if tag.RowsAffected() == 0 {
return "", fmt.Errorf("legacy record changed since an earlier import: %s/%s", item.EntityType, item.LegacyID)
}
if _, insertErr = tx.Exec(ctx, `INSERT INTO gateway.legacy_import_batch_records(batch_id,source_system,entity_type,legacy_id,payload_checksum) VALUES($1,$2,$3,$4,$5) ON CONFLICT DO NOTHING`, batchID, item.SourceSystem, item.EntityType, item.LegacyID, item.Checksum); insertErr != nil {
return "", insertErr
}
metadata, _ := json.Marshal(map[string]any{"batch_id": batchID, "payload_checksum": item.Checksum})
if _, insertErr = tx.Exec(ctx, `INSERT INTO gateway.legacy_id_mappings(source_system,entity_type,legacy_id,new_id,metadata) VALUES($1,$2,$3,$4,$5) ON CONFLICT(source_system,entity_type,legacy_id) DO UPDATE SET metadata=excluded.metadata WHERE gateway.legacy_id_mappings.new_id=excluded.new_id`, item.SourceSystem, item.EntityType, item.LegacyID, item.NewID, metadata); insertErr != nil {
return "", insertErr
}
}
if err = tx.Commit(ctx); err != nil {
return "", err
}
return batchID, nil
}
func bytesTrimSpace(value []byte) []byte { return []byte(strings.TrimSpace(string(value))) }
func write(value any) {
encoded, _ := json.MarshalIndent(value, "", " ")
fmt.Println(string(encoded))
}
func fatal(err error) { fmt.Fprintln(os.Stderr, err); os.Exit(1) }
+27
View File
@@ -0,0 +1,27 @@
package main
import (
"bytes"
"crypto/sha256"
"encoding/hex"
"encoding/json"
"testing"
)
func TestReadRecordsValidatesChecksumAndDeterministicID(t *testing.T) {
data := json.RawMessage(`{"id":7,"name":"研发"}`)
canonical, _ := json.Marshal(map[string]any{"source_system": "python-gateway", "entity_type": "departments", "legacy_id": "7", "data": data})
digest := sha256.Sum256(canonical)
line, _ := json.Marshal(record{SourceSystem: "python-gateway", EntityType: "departments", LegacyID: "7", Data: data, Checksum: hex.EncodeToString(digest[:])})
items, source, counts, err := readRecords(bytes.NewReader(append(line, '\n')))
if err != nil {
t.Fatal(err)
}
if len(items) != 1 || items[0].NewID == "" || source == "" || counts["departments"] != 1 {
t.Fatalf("unexpected result: %#v %s %#v", items, source, counts)
}
line[len(line)-2] ^= 1
if _, _, _, err = readRecords(bytes.NewReader(append(line, '\n'))); err == nil {
t.Fatal("expected tamper detection")
}
}
+137
View File
@@ -0,0 +1,137 @@
package main
import (
"bytes"
"context"
"encoding/json"
"flag"
"fmt"
"io"
"net/http"
"os"
"sort"
"strings"
"sync"
"sync/atomic"
"time"
)
type sample struct {
latency time.Duration
status int
err string
}
type report struct {
Requests int `json:"requests"`
Success int `json:"success"`
Errors int `json:"errors"`
ErrorRate float64 `json:"error_rate"`
RequestsPerSecond float64 `json:"requests_per_second"`
P50MS float64 `json:"p50_ms"`
P95MS float64 `json:"p95_ms"`
P99MS float64 `json:"p99_ms"`
DurationSeconds float64 `json:"duration_seconds"`
StatusCounts map[int]int `json:"status_counts"`
}
func main() {
base := flag.String("url", "http://127.0.0.1:8080", "gateway base URL")
path := flag.String("path", "/v1/chat/completions", "request path")
key := flag.String("api-key", os.Getenv("GATEWAY_LOADTEST_API_KEY"), "dedicated load-test API key")
model := flag.String("model", "test-model", "model or model alias")
concurrency := flag.Int("concurrency", 16, "concurrent workers")
duration := flag.Duration("duration", 30*time.Second, "test duration")
timeout := flag.Duration("timeout", 60*time.Second, "per-request timeout")
maxError := flag.Float64("max-error-rate", 0.01, "failure threshold")
maxP95 := flag.Duration("max-p95", 2*time.Second, "p95 latency threshold")
flag.Parse()
if *concurrency < 1 || *concurrency > 2000 || *duration < time.Second || *key == "" {
fmt.Fprintln(os.Stderr, "invalid arguments: api-key, positive duration and concurrency 1..2000 are required")
os.Exit(2)
}
target := strings.TrimRight(*base, "/") + *path
payload, _ := json.Marshal(map[string]any{"model": *model, "messages": []map[string]string{{"role": "user", "content": "Reply with OK."}}, "stream": false, "max_tokens": 8})
client := &http.Client{Timeout: *timeout, Transport: &http.Transport{MaxIdleConns: *concurrency * 2, MaxIdleConnsPerHost: *concurrency, IdleConnTimeout: 90 * time.Second, TLSHandshakeTimeout: 5 * time.Second}}
ctx, cancel := context.WithTimeout(context.Background(), *duration)
defer cancel()
started := time.Now()
results := make(chan sample, *concurrency*4)
var sequence atomic.Uint64
var workers sync.WaitGroup
for worker := 0; worker < *concurrency; worker++ {
workers.Add(1)
go func() {
defer workers.Done()
for ctx.Err() == nil {
number := sequence.Add(1)
request, err := http.NewRequestWithContext(ctx, http.MethodPost, target, bytes.NewReader(payload))
if err != nil {
results <- sample{err: err.Error()}
continue
}
request.Header.Set("Authorization", "Bearer "+*key)
request.Header.Set("Content-Type", "application/json")
request.Header.Set("X-Request-ID", fmt.Sprintf("load-%d", number))
request.Header.Set("Idempotency-Key", fmt.Sprintf("load-%d", number))
begin := time.Now()
response, err := client.Do(request)
elapsed := time.Since(begin)
if err != nil {
if ctx.Err() == nil {
results <- sample{latency: elapsed, err: err.Error()}
}
continue
}
_, readErr := io.Copy(io.Discard, io.LimitReader(response.Body, 2<<20))
response.Body.Close()
entry := sample{latency: elapsed, status: response.StatusCode}
if readErr != nil {
entry.err = readErr.Error()
}
results <- entry
}
}()
}
go func() { workers.Wait(); close(results) }()
samples := []sample{}
for result := range results {
samples = append(samples, result)
}
elapsed := time.Since(started)
summary := summarize(samples, elapsed)
encoded, _ := json.MarshalIndent(summary, "", " ")
fmt.Println(string(encoded))
if summary.ErrorRate > *maxError || time.Duration(summary.P95MS*float64(time.Millisecond)) > *maxP95 {
os.Exit(1)
}
}
func summarize(samples []sample, elapsed time.Duration) report {
summary := report{Requests: len(samples), StatusCounts: map[int]int{}, DurationSeconds: elapsed.Seconds()}
latencies := make([]time.Duration, 0, len(samples))
for _, item := range samples {
summary.StatusCounts[item.status]++
latencies = append(latencies, item.latency)
if item.err == "" && item.status >= 200 && item.status < 300 {
summary.Success++
} else {
summary.Errors++
}
}
sort.Slice(latencies, func(i, j int) bool { return latencies[i] < latencies[j] })
if summary.Requests > 0 {
summary.ErrorRate = float64(summary.Errors) / float64(summary.Requests)
summary.RequestsPerSecond = float64(summary.Requests) / elapsed.Seconds()
summary.P50MS = percentile(latencies, .50)
summary.P95MS = percentile(latencies, .95)
summary.P99MS = percentile(latencies, .99)
}
return summary
}
func percentile(values []time.Duration, p float64) float64 {
if len(values) == 0 {
return 0
}
index := int(float64(len(values)-1) * p)
return float64(values[index]) / float64(time.Millisecond)
}
+16
View File
@@ -0,0 +1,16 @@
package main
import (
"testing"
"time"
)
func TestSummarize(t *testing.T) {
summary := summarize([]sample{{latency: 10 * time.Millisecond, status: 200}, {latency: 20 * time.Millisecond, status: 200}, {latency: 100 * time.Millisecond, status: 500}}, time.Second)
if summary.Requests != 3 || summary.Success != 2 || summary.Errors != 1 {
t.Fatalf("unexpected counts: %#v", summary)
}
if summary.P95MS != 20 {
t.Fatalf("unexpected p95: %v", summary.P95MS)
}
}
+52
View File
@@ -0,0 +1,52 @@
package main
import (
"context"
"log/slog"
"os"
"os/signal"
"syscall"
"time"
"aigateway.local/core/internal/audit"
"aigateway.local/core/internal/platform/config"
"aigateway.local/core/internal/platform/database"
)
func main() {
logger := slog.New(slog.NewJSONHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelInfo}))
cfg, err := config.Load()
if err != nil {
logger.Error("invalid configuration", "error", err)
os.Exit(1)
}
if cfg.Database.URL == "" {
logger.Error("DATABASE_URL is required")
os.Exit(1)
}
ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
defer stop()
db, err := database.Open(ctx, cfg.Database)
if err != nil {
logger.Error("database initialization failed", "error", err)
os.Exit(1)
}
defer db.Close()
maintenance := audit.NewMaintenance(db, cfg.Audit.Retention, cfg.Audit.UsageRetention, cfg.Audit.PartitionMonthsAhead)
ticker := time.NewTicker(cfg.Audit.MaintenanceInterval)
defer ticker.Stop()
for {
result, runErr := maintenance.Run(ctx, time.Now())
if runErr != nil && ctx.Err() == nil {
logger.Error("audit maintenance failed", "error", runErr)
} else if runErr == nil {
logger.Info("audit maintenance complete", "created_partitions", result.CreatedPartitions, "dropped_partitions", result.DroppedPartitions, "deleted_audit_rows", result.DeletedAuditRows, "deleted_usage_rows", result.DeletedUsageRows)
}
select {
case <-ctx.Done():
logger.Info("maintenance worker stopped")
return
case <-ticker.C:
}
}
}
+44
View File
@@ -0,0 +1,44 @@
package main
import (
"context"
"flag"
"log/slog"
"os"
"aigateway.local/core/internal/platform/config"
"aigateway.local/core/internal/platform/database"
"aigateway.local/core/internal/platform/migrate"
)
func main() {
directory := flag.String("dir", "migrations", "directory containing ordered SQL migrations")
flag.Parse()
logger := slog.New(slog.NewJSONHandler(os.Stdout, nil))
cfg, err := config.Load()
if err != nil {
logger.Error("invalid configuration", "error", err)
os.Exit(1)
}
if cfg.Database.URL == "" {
logger.Error("DATABASE_URL is required by the migrator")
os.Exit(1)
}
ctx := context.Background()
pool, err := database.Open(ctx, cfg.Database)
if err != nil {
logger.Error("database initialization failed", "error", err)
os.Exit(1)
}
defer pool.Close()
migrations, err := migrate.Load(*directory)
if err != nil {
logger.Error("migration loading failed", "error", err)
os.Exit(1)
}
if err := migrate.Apply(ctx, pool, migrations); err != nil {
logger.Error("migration failed", "error", err)
os.Exit(1)
}
logger.Info("migrations complete", "count", len(migrations))
}
+61
View File
@@ -0,0 +1,61 @@
package main
import (
"context"
"log/slog"
"os"
"os/signal"
"syscall"
"aigateway.local/core/internal/platform/cache"
"aigateway.local/core/internal/platform/config"
"aigateway.local/core/internal/platform/cryptox"
"aigateway.local/core/internal/platform/database"
platformid "aigateway.local/core/internal/platform/id"
"aigateway.local/core/internal/workbench"
)
func main() {
logger := slog.New(slog.NewJSONHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelInfo}))
cfg, err := config.Load()
if err != nil {
logger.Error("invalid configuration", "error", err)
os.Exit(1)
}
if cfg.Database.URL == "" || cfg.Redis.CriticalURL == "" {
logger.Error("DATABASE_URL and REDIS_CRITICAL_URL are required")
os.Exit(1)
}
ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
defer stop()
db, err := database.Open(ctx, cfg.Database)
if err != nil {
logger.Error("database initialization failed", "error", err)
os.Exit(1)
}
defer db.Close()
client, err := cache.Open(cfg.Redis.CriticalURL)
if err != nil {
logger.Error("redis initialization failed", "error", err)
os.Exit(1)
}
defer client.Close()
cipher, err := cryptox.NewKeyring(cfg.Credentials.MasterKey, cfg.Credentials.KEKVersion, cfg.Credentials.KEKKeyring, "notification-signing-secret")
if err != nil {
logger.Error("notification cipher initialization failed", "error", err)
os.Exit(1)
}
consumer, err := platformid.NewUUID()
if err != nil {
logger.Error("consumer ID generation failed", "error", err)
os.Exit(1)
}
service := workbench.NewNotificationService(workbench.NewService(db), cipher, cfg.Credentials.AllowPrivateWebhookURL)
dispatcher := workbench.NewNotificationDispatcher(service, client, cfg.Outbox.Stream, "notification-"+consumer, logger)
logger.Info("notification worker started", "consumer", consumer, "stream", cfg.Outbox.Stream)
if err = dispatcher.Run(ctx); err != nil {
logger.Error("notification worker stopped unexpectedly", "error", err)
os.Exit(1)
}
logger.Info("notification worker stopped", "consumer", consumer)
}
+68
View File
@@ -0,0 +1,68 @@
package main
import (
"context"
"log/slog"
"os"
"os/signal"
"syscall"
"aigateway.local/core/internal/outbox"
"aigateway.local/core/internal/platform/cache"
"aigateway.local/core/internal/platform/config"
"aigateway.local/core/internal/platform/database"
platformid "aigateway.local/core/internal/platform/id"
)
func main() {
logger := slog.New(slog.NewJSONHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelInfo}))
cfg, err := config.Load()
if err != nil {
logger.Error("invalid configuration", "error", err)
os.Exit(1)
}
if cfg.Database.URL == "" || cfg.Redis.CriticalURL == "" {
logger.Error("DATABASE_URL and REDIS_CRITICAL_URL are required")
os.Exit(1)
}
ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
defer stop()
db, err := database.Open(ctx, cfg.Database)
if err != nil {
logger.Error("database initialization failed", "error", err)
os.Exit(1)
}
defer db.Close()
redisClient, err := cache.Open(cfg.Redis.CriticalURL)
if err != nil {
logger.Error("redis initialization failed", "error", err)
os.Exit(1)
}
defer redisClient.Close()
if err := db.Ping(ctx); err != nil {
logger.Error("database is unavailable", "error", err)
os.Exit(1)
}
if err := redisClient.Ping(ctx).Err(); err != nil {
logger.Error("redis is unavailable", "error", err)
os.Exit(1)
}
workerID, err := platformid.NewUUID()
if err != nil {
logger.Error("worker ID generation failed", "error", err)
os.Exit(1)
}
workerID = "outbox-" + workerID
worker := outbox.NewWorker(
outbox.NewStore(db),
outbox.NewRedisPublisher(redisClient, cfg.Outbox.Stream, cfg.Outbox.StreamMaxLength, cfg.Outbox.MarkerTTL),
outbox.WorkerConfig{WorkerID: workerID, BatchSize: cfg.Outbox.BatchSize, PollInterval: cfg.Outbox.PollInterval, Lease: cfg.Outbox.Lease, MaxAttempts: cfg.Outbox.MaxAttempts, MaxBackoff: cfg.Outbox.MaxBackoff},
logger,
)
logger.Info("outbox worker started", "worker_id", workerID, "stream", cfg.Outbox.Stream)
if err := worker.Run(ctx); err != nil {
logger.Error("outbox worker stopped unexpectedly", "error", err)
os.Exit(1)
}
logger.Info("outbox worker stopped", "worker_id", workerID)
}
+33
View File
@@ -0,0 +1,33 @@
FROM golang:1.26.5-alpine AS builder
WORKDIR /src
COPY go.mod go.sum ./
RUN go mod download
COPY cmd ./cmd
COPY internal ./internal
ARG VERSION=dev
RUN CGO_ENABLED=0 go build -trimpath -ldflags="-s -w -X main.version=${VERSION}" -o /out/gateway-api ./cmd/gateway-api
RUN CGO_ENABLED=0 go build -trimpath -ldflags="-s -w" -o /out/gateway-migrator ./cmd/gateway-migrator
RUN CGO_ENABLED=0 go build -trimpath -ldflags="-s -w" -o /out/gateway-bootstrap ./cmd/gateway-bootstrap
RUN CGO_ENABLED=0 go build -trimpath -ldflags="-s -w" -o /out/gateway-outbox-worker ./cmd/gateway-outbox-worker
RUN CGO_ENABLED=0 go build -trimpath -ldflags="-s -w" -o /out/gateway-maintenance ./cmd/gateway-maintenance
RUN CGO_ENABLED=0 go build -trimpath -ldflags="-s -w" -o /out/gateway-notification-worker ./cmd/gateway-notification-worker
RUN CGO_ENABLED=0 go build -trimpath -ldflags="-s -w" -o /out/gateway-loadtest ./cmd/gateway-loadtest
RUN CGO_ENABLED=0 go build -trimpath -ldflags="-s -w" -o /out/gateway-legacy-import ./cmd/gateway-legacy-import
FROM alpine:3.22
RUN apk add --no-cache ca-certificates tzdata \
&& addgroup -S gateway \
&& adduser -S -G gateway gateway
WORKDIR /app
COPY --from=builder /out/gateway-api /usr/local/bin/gateway-api
COPY --from=builder /out/gateway-migrator /usr/local/bin/gateway-migrator
COPY --from=builder /out/gateway-bootstrap /usr/local/bin/gateway-bootstrap
COPY --from=builder /out/gateway-outbox-worker /usr/local/bin/gateway-outbox-worker
COPY --from=builder /out/gateway-maintenance /usr/local/bin/gateway-maintenance
COPY --from=builder /out/gateway-notification-worker /usr/local/bin/gateway-notification-worker
COPY --from=builder /out/gateway-loadtest /usr/local/bin/gateway-loadtest
COPY --from=builder /out/gateway-legacy-import /usr/local/bin/gateway-legacy-import
COPY migrations ./migrations
USER gateway
EXPOSE 8080
ENTRYPOINT ["gateway-api"]
+16
View File
@@ -0,0 +1,16 @@
FROM node:24-alpine AS builder
ARG APP=admin
WORKDIR /src/web
RUN corepack enable
COPY web/package.json web/pnpm-lock.yaml web/pnpm-workspace.yaml ./
COPY web/apps ./apps
COPY web/packages ./packages
RUN CI=1 corepack pnpm@10.15.1 install --frozen-lockfile
RUN corepack pnpm@10.15.1 --filter @aigateway/${APP} build
FROM nginx:1.29-alpine
ARG APP=admin
COPY deploy/nginx-web.conf /etc/nginx/conf.d/default.conf
RUN sed -i "s/__APP__/${APP}/g" /etc/nginx/conf.d/default.conf
COPY --from=builder /src/web/apps/${APP}/dist /usr/share/nginx/html/${APP}
EXPOSE 80
+74
View File
@@ -0,0 +1,74 @@
# Production deployment
This bundle builds the Go services and both Art Design Pro applications from
source. PostgreSQL and two Redis roles are included; MinIO/S3 and ClickHouse are
not required.
## Prerequisites
- Docker Engine with Compose v2
- At least 4 CPU cores, 8 GiB RAM and 30 GiB free disk for an initial build
- An external TLS reverse proxy or load balancer
- A backup destination for the PostgreSQL volume
## First deployment
Run all commands from the repository root:
```bash
cp deploy/production.env.example deploy/production.env
chmod 600 deploy/production.env
# Edit deploy/production.env and replace every CHANGE_ME value.
docker compose \
--env-file deploy/production.env \
-f deploy/docker-compose.production.yml \
config --quiet
docker compose \
--env-file deploy/production.env \
-f deploy/docker-compose.production.yml \
up -d --build
```
Create the initial administrator once:
```bash
docker compose \
--env-file deploy/production.env \
-f deploy/docker-compose.production.yml \
--profile tools run --rm bootstrap-admin
```
Then remove `BOOTSTRAP_ADMIN_PASSWORD` from `deploy/production.env` and use the
admin UI to create database-backed gateway API keys.
## Endpoints
- API and OpenAI-compatible gateway: `127.0.0.1:8080`
- Admin UI: `http://127.0.0.1:8081/admin/`
- Portal UI: `http://127.0.0.1:8082/portal/`
- Liveness/readiness: `/healthz` and `/readyz`
Ports bind to loopback by default. Terminate TLS at a reverse proxy and forward
to these endpoints. Change `*_BIND_IP` only when the host firewall and network
policy are already in place.
## Operations
Check status and logs:
```bash
docker compose --env-file deploy/production.env -f deploy/docker-compose.production.yml ps
docker compose --env-file deploy/production.env -f deploy/docker-compose.production.yml logs --tail=200 gateway-api
curl --fail http://127.0.0.1:8080/readyz
```
For upgrades, back up PostgreSQL first, change `GATEWAY_VERSION`, then run the
same `up -d --build` command. The one-shot migrator applies forward migrations
before the API starts. Do not use `docker compose down -v` in production because
it removes persistent data.
The bundled database URLs use `sslmode=disable` only for the private Compose
network. When using an external PostgreSQL or Redis service, require TLS and use
`sslmode=verify-full` / `rediss://` as supported by that service.
+189
View File
@@ -0,0 +1,189 @@
name: ai-gateway
x-backend-build: &backend-build
context: ..
dockerfile: deploy/Dockerfile
args:
VERSION: ${GATEWAY_VERSION:-0.10.0}
x-gateway-environment: &gateway-environment
APP_ENV: production
HTTP_ADDR: :8080
DATABASE_URL: ${DATABASE_URL:?DATABASE_URL is required}
REDIS_CRITICAL_URL: ${REDIS_CRITICAL_URL:?REDIS_CRITICAL_URL is required}
REDIS_CACHE_URL: ${REDIS_CACHE_URL:-}
GATEWAY_BOOTSTRAP_API_KEY_ENABLED: ${GATEWAY_BOOTSTRAP_API_KEY_ENABLED:-false}
GATEWAY_BOOTSTRAP_API_KEY: ${GATEWAY_BOOTSTRAP_API_KEY:-}
CREDENTIAL_MASTER_KEY: ${CREDENTIAL_MASTER_KEY:?CREDENTIAL_MASTER_KEY is required}
CREDENTIAL_KEK_VERSION: ${CREDENTIAL_KEK_VERSION:-1}
CREDENTIAL_KEK_KEYRING: ${CREDENTIAL_KEK_KEYRING:-}
ALLOW_PRIVATE_PROVIDER_URLS: ${ALLOW_PRIVATE_PROVIDER_URLS:-false}
ALLOW_PRIVATE_TOOL_URLS: ${ALLOW_PRIVATE_TOOL_URLS:-false}
ALLOW_PRIVATE_WEBHOOK_URLS: ${ALLOW_PRIVATE_WEBHOOK_URLS:-false}
UPSTREAM_FALLBACK_ENABLED: ${UPSTREAM_FALLBACK_ENABLED:-false}
UPSTREAM_BASE_URL: ${UPSTREAM_BASE_URL:-https://api.openai.com}
UPSTREAM_API_KEY: ${UPSTREAM_API_KEY:-}
SHADOW_BASE_URL: ${SHADOW_BASE_URL:-}
SHADOW_API_KEY: ${SHADOW_API_KEY:-}
SHADOW_SAMPLE_RATE: ${SHADOW_SAMPLE_RATE:-0}
x-backend-service: &backend-service
image: ai-gateway-go:${GATEWAY_VERSION:-0.10.0}
build: *backend-build
environment: *gateway-environment
read_only: true
tmpfs:
- /tmp:size=64m,mode=1777
cap_drop:
- ALL
security_opt:
- no-new-privileges:true
services:
postgres:
image: postgres:17-alpine
environment:
POSTGRES_DB: ${POSTGRES_DB:-gateway}
POSTGRES_USER: ${POSTGRES_USER:-gateway}
POSTGRES_PASSWORD: ${POSTGRES_PASSWORD:?POSTGRES_PASSWORD is required}
volumes:
- postgres-data:/var/lib/postgresql/data
healthcheck:
test: ["CMD-SHELL", "pg_isready -U $${POSTGRES_USER} -d $${POSTGRES_DB}"]
interval: 5s
timeout: 3s
retries: 30
restart: unless-stopped
redis-critical:
image: redis:8.2-alpine
environment:
REDIS_PASSWORD: ${REDIS_PASSWORD:?REDIS_PASSWORD is required}
command:
- sh
- -c
- exec redis-server --requirepass "$${REDIS_PASSWORD}" --appendonly yes --appendfsync everysec --maxmemory-policy noeviction
volumes:
- redis-critical-data:/data
healthcheck:
test: ["CMD-SHELL", "redis-cli --no-auth-warning -a \"$${REDIS_PASSWORD}\" ping | grep -q PONG"]
interval: 5s
timeout: 3s
retries: 30
restart: unless-stopped
redis-cache:
image: redis:8.2-alpine
environment:
REDIS_PASSWORD: ${REDIS_PASSWORD:?REDIS_PASSWORD is required}
REDIS_CACHE_MAXMEMORY: ${REDIS_CACHE_MAXMEMORY:-256mb}
command:
- sh
- -c
- exec redis-server --requirepass "$${REDIS_PASSWORD}" --maxmemory "$${REDIS_CACHE_MAXMEMORY}" --maxmemory-policy allkeys-lfu --save ""
healthcheck:
test: ["CMD-SHELL", "redis-cli --no-auth-warning -a \"$${REDIS_PASSWORD}\" ping | grep -q PONG"]
interval: 5s
timeout: 3s
retries: 30
restart: unless-stopped
migrator:
<<: *backend-service
entrypoint: ["gateway-migrator"]
depends_on:
postgres:
condition: service_healthy
restart: "no"
gateway-api:
<<: *backend-service
ports:
- "${GATEWAY_BIND_IP:-127.0.0.1}:${GATEWAY_PORT:-8080}:8080"
depends_on:
migrator:
condition: service_completed_successfully
redis-critical:
condition: service_healthy
redis-cache:
condition: service_healthy
healthcheck:
test: ["CMD-SHELL", "wget -q -O /dev/null http://127.0.0.1:8080/readyz"]
interval: 10s
timeout: 3s
retries: 30
start_period: 10s
restart: unless-stopped
admin-web:
image: ai-gateway-admin:${GATEWAY_VERSION:-0.10.0}
build:
context: ..
dockerfile: deploy/Dockerfile.web
args:
APP: admin
ports:
- "${WEB_BIND_IP:-127.0.0.1}:${ADMIN_PORT:-8081}:80"
depends_on:
gateway-api:
condition: service_healthy
restart: unless-stopped
portal-web:
image: ai-gateway-portal:${GATEWAY_VERSION:-0.10.0}
build:
context: ..
dockerfile: deploy/Dockerfile.web
args:
APP: portal
ports:
- "${WEB_BIND_IP:-127.0.0.1}:${PORTAL_PORT:-8082}:80"
depends_on:
gateway-api:
condition: service_healthy
restart: unless-stopped
outbox-worker:
<<: *backend-service
entrypoint: ["gateway-outbox-worker"]
depends_on:
migrator:
condition: service_completed_successfully
redis-critical:
condition: service_healthy
restart: unless-stopped
maintenance-worker:
<<: *backend-service
entrypoint: ["gateway-maintenance"]
depends_on:
migrator:
condition: service_completed_successfully
restart: unless-stopped
notification-worker:
<<: *backend-service
entrypoint: ["gateway-notification-worker"]
depends_on:
migrator:
condition: service_completed_successfully
redis-critical:
condition: service_healthy
restart: unless-stopped
bootstrap-admin:
<<: *backend-service
profiles: ["tools"]
entrypoint: ["gateway-bootstrap"]
environment:
<<: *gateway-environment
BOOTSTRAP_ADMIN_USERNAME: ${BOOTSTRAP_ADMIN_USERNAME:-admin}
BOOTSTRAP_ADMIN_PASSWORD: ${BOOTSTRAP_ADMIN_PASSWORD:-}
depends_on:
migrator:
condition: service_completed_successfully
restart: "no"
volumes:
postgres-data:
redis-critical-data:
+154
View File
@@ -0,0 +1,154 @@
services:
postgres:
image: postgres:17-alpine
environment:
POSTGRES_DB: gateway
POSTGRES_USER: gateway
POSTGRES_PASSWORD: gateway
volumes:
- postgres-data:/var/lib/postgresql/data
healthcheck:
test: ["CMD-SHELL", "pg_isready -U gateway -d gateway"]
interval: 5s
timeout: 3s
retries: 20
redis-critical:
image: redis:8.2-alpine
command: ["redis-server", "--appendonly", "yes", "--appendfsync", "everysec", "--maxmemory-policy", "noeviction"]
volumes:
- redis-critical-data:/data
healthcheck:
test: ["CMD", "redis-cli", "ping"]
interval: 5s
timeout: 3s
retries: 20
redis-cache:
image: redis:8.2-alpine
command: ["redis-server", "--maxmemory", "256mb", "--maxmemory-policy", "allkeys-lfu", "--save", ""]
healthcheck:
test: ["CMD", "redis-cli", "ping"]
interval: 5s
timeout: 3s
retries: 20
migrator:
build:
context: ..
dockerfile: deploy/Dockerfile
args:
VERSION: ${GATEWAY_VERSION:-0.10.0}
entrypoint: ["gateway-migrator"]
environment: &gateway-environment
APP_ENV: local
HTTP_ADDR: :8080
DATABASE_URL: postgres://gateway:gateway@postgres:5432/gateway?sslmode=disable
REDIS_CRITICAL_URL: redis://redis-critical:6379/0
REDIS_CACHE_URL: redis://redis-cache:6379/0
GATEWAY_BOOTSTRAP_API_KEY: ${GATEWAY_BOOTSTRAP_API_KEY:-}
GATEWAY_BOOTSTRAP_API_KEY_ENABLED: ${GATEWAY_BOOTSTRAP_API_KEY_ENABLED:-false}
CREDENTIAL_MASTER_KEY: ${CREDENTIAL_MASTER_KEY:-AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=}
CREDENTIAL_KEK_VERSION: 1
CREDENTIAL_KEK_KEYRING: ${CREDENTIAL_KEK_KEYRING:-}
ALLOW_PRIVATE_TOOL_URLS: ${ALLOW_PRIVATE_TOOL_URLS:-false}
ALLOW_PRIVATE_WEBHOOK_URLS: ${ALLOW_PRIVATE_WEBHOOK_URLS:-false}
UPSTREAM_BASE_URL: ${UPSTREAM_BASE_URL:-https://api.openai.com}
UPSTREAM_API_KEY: ${UPSTREAM_API_KEY:-}
SHADOW_BASE_URL: ${SHADOW_BASE_URL:-}
SHADOW_API_KEY: ${SHADOW_API_KEY:-}
SHADOW_SAMPLE_RATE: ${SHADOW_SAMPLE_RATE:-0}
depends_on:
postgres:
condition: service_healthy
gateway-api:
build:
context: ..
dockerfile: deploy/Dockerfile
args:
VERSION: ${GATEWAY_VERSION:-0.10.0}
environment: *gateway-environment
ports:
- "${GATEWAY_PORT:-8080}:8080"
depends_on:
migrator:
condition: service_completed_successfully
redis-critical:
condition: service_healthy
redis-cache:
condition: service_healthy
restart: unless-stopped
admin-web:
build:
context: ..
dockerfile: deploy/Dockerfile.web
args:
APP: admin
ports:
- "${ADMIN_PORT:-8081}:80"
depends_on:
gateway-api:
condition: service_started
restart: unless-stopped
portal-web:
build:
context: ..
dockerfile: deploy/Dockerfile.web
args:
APP: portal
ports:
- "${PORTAL_PORT:-8082}:80"
depends_on:
gateway-api:
condition: service_started
restart: unless-stopped
outbox-worker:
build:
context: ..
dockerfile: deploy/Dockerfile
args:
VERSION: ${GATEWAY_VERSION:-0.10.0}
entrypoint: ["gateway-outbox-worker"]
environment: *gateway-environment
depends_on:
migrator:
condition: service_completed_successfully
redis-critical:
condition: service_healthy
restart: unless-stopped
maintenance-worker:
build:
context: ..
dockerfile: deploy/Dockerfile
args:
VERSION: ${GATEWAY_VERSION:-0.10.0}
entrypoint: ["gateway-maintenance"]
environment: *gateway-environment
depends_on:
migrator:
condition: service_completed_successfully
restart: unless-stopped
notification-worker:
build:
context: ..
dockerfile: deploy/Dockerfile
args:
VERSION: ${GATEWAY_VERSION:-0.10.0}
entrypoint: ["gateway-notification-worker"]
environment: *gateway-environment
depends_on:
migrator:
condition: service_completed_successfully
redis-critical:
condition: service_healthy
restart: unless-stopped
volumes:
postgres-data:
redis-critical-data:
+71
View File
@@ -0,0 +1,71 @@
server {
listen 80;
server_name _;
# Keep Location headers relative (Location: /admin/) instead of letting
# nginx absolute_redirect rebuild them from $host + the listening port.
# The gateway is commonly published behind a non-standard port (e.g. 18081),
# and an absolute redirect would drop that port and send browsers to :80.
absolute_redirect off;
# Match the gateway's HTTP_MAX_BODY_BYTES (default 32 MiB). nginx's default
# of 1 MiB otherwise rejects large prompts and knowledge-base imports with
# 413 before they ever reach the gateway.
client_max_body_size 32m;
root /usr/share/nginx/html;
index index.html;
location = / {
return 302 /__APP__/;
}
location = /__APP__ {
return 301 /__APP__/;
}
location = /healthz {
proxy_pass http://gateway-api:8080/healthz;
proxy_set_header Host $host;
proxy_set_header X-Forwarded-Proto $scheme;
proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
}
location = /readyz {
proxy_pass http://gateway-api:8080/readyz;
proxy_set_header Host $host;
proxy_set_header X-Forwarded-Proto $scheme;
proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
}
location /api/ {
proxy_pass http://gateway-api:8080;
proxy_http_version 1.1;
proxy_set_header Host $host;
proxy_set_header X-Request-ID $request_id;
proxy_set_header X-Forwarded-Proto $scheme;
# nginx 是控制台(/api/)的唯一入口:用 $remote_addr 覆盖 X-Forwarded-For
# 避免客户端自带 X-Forwarded-For 头伪造来源 IP、绕过登录限流。
proxy_set_header X-Forwarded-For $remote_addr;
}
location /v1/ {
proxy_pass http://gateway-api:8080;
proxy_http_version 1.1;
proxy_buffering off;
proxy_read_timeout 600s;
proxy_set_header Host $host;
proxy_set_header X-Request-ID $request_id;
proxy_set_header X-Forwarded-Proto $scheme;
proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
}
# Vite emits assets and client-side routes below /admin/ or /portal/.
location /__APP__/ {
try_files $uri $uri/ /__APP__/index.html;
}
location / {
return 404;
}
}
+42
View File
@@ -0,0 +1,42 @@
# Copy this file to deploy/production.env and replace every CHANGE_ME value.
# Keep production.env outside version control and restrict it to the deploy user.
GATEWAY_VERSION=0.10.0
POSTGRES_DB=gateway
POSTGRES_USER=gateway
POSTGRES_PASSWORD=CHANGE_ME_LONG_RANDOM_POSTGRES_PASSWORD
REDIS_PASSWORD=CHANGE_ME_LONG_RANDOM_REDIS_PASSWORD
# URL-encode special characters in passwords used in URLs.
DATABASE_URL=postgres://gateway:CHANGE_ME_URL_ENCODED_POSTGRES_PASSWORD@postgres:5432/gateway?sslmode=disable
REDIS_CRITICAL_URL=redis://:CHANGE_ME_URL_ENCODED_REDIS_PASSWORD@redis-critical:6379/0
REDIS_CACHE_URL=redis://:CHANGE_ME_URL_ENCODED_REDIS_PASSWORD@redis-cache:6379/0
# Generate once with: openssl rand -base64 32
# Never rotate this value without following a KEK rotation procedure.
CREDENTIAL_MASTER_KEY=CHANGE_ME_BASE64_32_BYTE_KEY
CREDENTIAL_KEK_VERSION=1
CREDENTIAL_KEK_KEYRING=
# Disabled by default. Create database-backed API keys after bootstrap.
GATEWAY_BOOTSTRAP_API_KEY_ENABLED=false
GATEWAY_BOOTSTRAP_API_KEY=
UPSTREAM_FALLBACK_ENABLED=false
UPSTREAM_BASE_URL=https://api.openai.com
UPSTREAM_API_KEY=
ALLOW_PRIVATE_PROVIDER_URLS=false
ALLOW_PRIVATE_TOOL_URLS=false
ALLOW_PRIVATE_WEBHOOK_URLS=false
GATEWAY_BIND_IP=127.0.0.1
GATEWAY_PORT=8080
WEB_BIND_IP=127.0.0.1
ADMIN_PORT=8081
PORTAL_PORT=8082
REDIS_CACHE_MAXMEMORY=256mb
# Used only for the one-time bootstrap-admin command; remove after use.
BOOTSTRAP_ADMIN_USERNAME=admin
BOOTSTRAP_ADMIN_PASSWORD=CHANGE_ME_AT_LEAST_12_CHARACTERS
+18
View File
@@ -0,0 +1,18 @@
# ADR-0001:基线存储只使用 PostgreSQL 与 Redis
状态:已接受
日期:2026-08-10
## 决策
第一版部署不包含 MinIO/S3 与 ClickHouse。
- PostgreSQL 保存权威业务数据、outbox、有限长度的审计请求/响应预览和聚合统计。
- critical Redis 保存限流、配额、幂等、配置协调与短期可靠流,使用 AOF 和 `noeviction`
- cache Redis 保存可丢弃缓存,使用 `allkeys-lfu`,故障时网关降级运行。
- 不把完整且无上限的请求正文、响应正文和知识库原始文件写入本地磁盘。
- 知识库只保存经过提取、大小不超过 2 MiB 的纯文本及其分块;`Retriever` 隔离检索实现,基线使用 PostgreSQL FTS 和中文词片评分。
## 影响
基线更容易部署和运维,但不承诺无限期保存完整审计正文,也不针对超大规模多年期分析查询优化。未来只有在容量和合规需求明确时,才通过 `AuditBodyStore``AnalyticsSink` 接口接入对象存储或分析库。
+35
View File
@@ -0,0 +1,35 @@
# 双系统切换与回退手册
## 进入切换窗口前
1. 对照 `route-contract-matrix.md`,确认所有仍有调用量的旧契约已经完成客户端迁移或兼容层验证。
2. 连续观察影子指标;`gateway_shadow_errors_total` 不增长,契约 mismatch 已逐项解释,且 shadow 专用 Key 没有生产权限扩张。
3. 使用生产等价数据量执行 `gateway-loadtest`,确认错误率、p95、数据库连接池、critical Redis 延迟和审计队列丢弃数均在阈值内。
4. 冻结旧系统配置写入,执行最后一次幂等数据导入和数量/外键校验。
5. 确认 Go 与旧网关的 `/readyz` 均为成功,并备份 Nginx 活动 upstream include。
## 原子切换
`scripts/cutover.sh` 只修改一个专用 Nginx include。它先探活目标,写临时文件后原子替换,执行 `nginx -t`,再 reload;验证或 reload 失败会恢复备份。
```bash
export NGINX_ACTIVE_INCLUDE=/etc/nginx/conf.d/ai-gateway-active.conf
export GO_UPSTREAM=10.0.0.21:8080
export LEGACY_UPSTREAM=10.0.0.20:8000
export GO_HEALTH_URL=http://10.0.0.21:8080/readyz
export LEGACY_HEALTH_URL=http://10.0.0.20:8000/readyz
export PUBLIC_HEALTH_URL=https://gateway.example.com/readyz
scripts/cutover.sh preflight
scripts/cutover.sh go
scripts/cutover.sh status
```
## 回退触发条件
- 5 分钟窗口错误率超过既定 SLO,或出现鉴权/配额错误的系统性抬升;
- p95/p99 持续超过容量验收阈值;
- 审计队列出现丢弃、critical Redis 不稳定、数据库连接池耗尽;
- 与旧系统相比出现未经批准的响应结构差异或关键业务功能缺失。
执行 `scripts/cutover.sh rollback`。回退只切流量,不回写/删除 Go 数据;outbox 与审计 Worker 保持运行,便于事后核对。确认旧系统恢复后保存本次 include 备份、指标和日志,再分析问题。
+625
View File
@@ -0,0 +1,625 @@
# 旧 Python 网关数据字典与 ID 迁移映射
> 由 `scripts/export_legacy_dictionary.py` 从旧工程 SQLAlchemy metadata 生成。
> `audit_log_views` 是兼容查询投影,不作为独立持久化实体,因此这里统计 38 个实体。
## 映射约定
- 旧整数主键统一通过 `gateway.legacy_id_mappings` 映射为 UUID,不在新业务表保留双主键。
- UUID 使用固定 namespace 的 UUIDv5 生成,输入为 `source_system/entity_type/legacy_id`,重复导入保持幂等。
- 外键迁移必须先查映射表;缺失映射时整批失败,不允许写入悬空引用。
- 旧密文不能直接复用:Provider、TOTP、内部服务 Key 等敏感字段在导入时解密后用新用途标签重新加密。
- `audit_logs` 的完整请求/响应正文不进入新库,只迁移受长度约束的预览与结构化元数据。
## 实体去向
| 旧表 | 新域/目标 | 状态 |
|---|---|---|
| `admin_accounts` | gateway.admin_accounts(已落地) | 已建表,待批量导入器 |
| `application_revisions` | gateway.application_versions | 已建表,待批量导入器转换 config 引用 |
| `application_runs` | gateway.application_runs | 已建表,待批量导入器 |
| `applications` | gateway.applications | 已建表,待批量导入器 |
| `audit_logs` | gateway.audit_events(已落地,不迁移无限长正文) | 已建表,待批量导入器 |
| `audit_search_records` | 对应 M2–M4 领域表 | 已登记,随对应里程碑建表 |
| `audit_settings` | 对应 M2–M4 领域表 | 已登记,随对应里程碑建表 |
| `content_audit_scopes` | 对应 M2–M4 领域表 | 已登记,随对应里程碑建表 |
| `departments` | gateway.departments(已落地) | 已建表,待批量导入器 |
| `dept_budgets` | 对应 M2–M4 领域表 | 已登记,随对应里程碑建表 |
| `feishu_auth_settings` | 对应 M2–M4 领域表 | 已登记,随对应里程碑建表 |
| `filter_rules` | 对应 M2–M4 领域表 | 已登记,随对应里程碑建表 |
| `hallucination_events` | 对应 M2–M4 领域表 | 已登记,随对应里程碑建表 |
| `hallucination_policies` | 对应 M2–M4 领域表 | 已登记,随对应里程碑建表 |
| `hallucination_settings` | 对应 M2–M4 领域表 | 已登记,随对应里程碑建表 |
| `internal_service_keys` | 对应 M2–M4 领域表 | 已登记,随对应里程碑建表 |
| `knowledge_bases` | gateway.knowledge_bases | 已建表,旧 embedding_model 转检索器迁移策略待导入器处理 |
| `knowledge_chunks` | gateway.knowledge_chunks | 已建表,基线重新分块/建 PostgreSQL 搜索向量 |
| `knowledge_documents` | gateway.knowledge_documents | 已建表,只导入不超过 2 MiB 的提取文本 |
| `managed_conversations` | 对应 M2–M4 领域表 | 已登记,随对应里程碑建表 |
| `managed_messages` | 对应 M2–M4 领域表 | 已登记,随对应里程碑建表 |
| `model_aliases` | 对应 M2–M4 领域表 | 已登记,随对应里程碑建表 |
| `model_pricing` | 对应 M2–M4 领域表 | 已登记,随对应里程碑建表 |
| `model_qps_configs` | 对应 M2–M4 领域表 | 已登记,随对应里程碑建表 |
| `model_requests` | 对应 M2–M4 领域表 | 已登记,随对应里程碑建表 |
| `notify_config` | gateway.notification_channels | 已建表,待解密后重新加密签名密钥 |
| `pricing_config` | 对应 M2–M4 领域表 | 已登记,随对应里程碑建表 |
| `prompt_categories` | gateway.prompt_categories | 已建表,待批量导入器 |
| `prompt_favorites` | 对应 M2–M4 领域表 | 已登记,随对应里程碑建表 |
| `prompt_templates` | gateway.prompt_templates | 已建表,待批量导入器 |
| `prompt_versions` | gateway.prompt_versions | 已建表,待批量导入器 |
| `providers` | gateway.providers + gateway.provider_models(已落地,models 拆表) | 已建表,待批量导入器 |
| `rate_limit_rules` | 对应 M2–M4 领域表 | 已登记,随对应里程碑建表 |
| `routing_groups` | 对应 M2–M4 领域表 | 已登记,随对应里程碑建表 |
| `routing_targets` | 对应 M2–M4 领域表 | 已登记,随对应里程碑建表 |
| `tool_definitions` | gateway.tool_definitions | 已建表,待解密后用新用途标签加密请求头 |
| `users` | gateway.portal_users(已落地) | 已建表,待批量导入器 |
| `virtual_keys` | gateway.api_keys(已落地,字段需转换) | 已建表,待批量导入器 |
## 字段字典
### `admin_accounts`
| 字段 | 旧类型 | 可空 | 默认值 | 键/引用 |
|---|---|---:|---|---|
| `id` | `INTEGER` | 否 | `—` | PK |
| `username` | `VARCHAR(64)` | 否 | `—` | UNIQUE |
| `password_hash` | `VARCHAR(128)` | 否 | `—` | — |
| `display_name` | `VARCHAR(64)` | 是 | `` | — |
| `role` | `VARCHAR(16)` | 是 | `superadmin` | — |
| `active` | `BOOLEAN` | 是 | `True` | — |
| `last_login` | `DATETIME` | 是 | `` | — |
| `failed_logins` | `INTEGER` | 是 | `0` | — |
| `locked_until` | `DATETIME` | 是 | `` | — |
| `created_at` | `DATETIME` | 是 | `now` | — |
| `totp_secret` | `VARCHAR(255)` | 是 | `` | — |
| `totp_enabled` | `BOOLEAN` | 是 | `False` | — |
| `totp_last_step` | `INTEGER` | 是 | `` | — |
| `totp_backup_codes` | `TEXT` | 是 | `[]` | — |
| `totp_confirmed_at` | `DATETIME` | 是 | `` | — |
### `application_revisions`
| 字段 | 旧类型 | 可空 | 默认值 | 键/引用 |
|---|---|---:|---|---|
| `id` | `INTEGER` | 否 | `` | PK |
| `application_id` | `INTEGER` | 否 | `` | FK → `applications.id` |
| `version` | `INTEGER` | 否 | `` | — |
| `config` | `TEXT` | 否 | `` | — |
| `change_note` | `VARCHAR(256)` | 是 | `` | — |
| `published_by` | `VARCHAR(128)` | 是 | `` | — |
| `created_at` | `DATETIME` | 是 | `now` | — |
### `application_runs`
| 字段 | 旧类型 | 可空 | 默认值 | 键/引用 |
|---|---|---:|---|---|
| `id` | `INTEGER` | 否 | `` | PK |
| `application_id` | `INTEGER` | 否 | `` | FK → `applications.id` |
| `version` | `INTEGER` | 是 | `0` | — |
| `user_id` | `INTEGER` | 是 | `` | FK → `users.id` |
| `status` | `VARCHAR(16)` | 是 | `success` | — |
| `latency_ms` | `INTEGER` | 是 | `0` | — |
| `retrieval_count` | `INTEGER` | 是 | `0` | — |
| `tool_calls` | `INTEGER` | 是 | `0` | — |
| `error` | `VARCHAR(512)` | 是 | `` | — |
| `created_at` | `DATETIME` | 是 | `now` | — |
### `applications`
| 字段 | 旧类型 | 可空 | 默认值 | 键/引用 |
|---|---|---:|---|---|
| `id` | `INTEGER` | 否 | `` | PK |
| `app_code` | `VARCHAR(64)` | 否 | `` | UNIQUE |
| `name` | `VARCHAR(128)` | 否 | `` | — |
| `description` | `TEXT` | 是 | `` | — |
| `dept_id` | `INTEGER` | 是 | `` | FK → `departments.id` |
| `owner_user_id` | `INTEGER` | 是 | `` | FK → `users.id` |
| `status` | `VARCHAR(16)` | 是 | `active` | — |
| `allowed_models` | `TEXT` | 是 | `[]` | — |
| `quota_tokens` | `INTEGER` | 是 | `` | — |
| `rpm_limit` | `INTEGER` | 是 | `` | — |
| `draft_config` | `TEXT` | 是 | `{}` | — |
| `published_config` | `TEXT` | 是 | `{}` | — |
| `release_version` | `INTEGER` | 是 | `0` | — |
| `published_at` | `DATETIME` | 是 | `` | — |
| `runtime_key_id` | `INTEGER` | 是 | `` | — |
| `runtime_key_secret` | `TEXT` | 是 | `` | — |
| `created_by` | `VARCHAR(128)` | 是 | `` | — |
| `created_at` | `DATETIME` | 是 | `now` | — |
| `updated_at` | `DATETIME` | 是 | `now` | — |
### `audit_logs`
| 字段 | 旧类型 | 可空 | 默认值 | 键/引用 |
|---|---|---:|---|---|
| `id` | `INTEGER` | 否 | `` | PK |
| `event_id` | `VARCHAR(32)` | 是 | `` | UNIQUE |
| `ts` | `DATETIME` | 是 | `now` | — |
| `key_id` | `INTEGER` | 是 | `` | — |
| `key_name` | `VARCHAR(128)` | 是 | `` | — |
| `user_id` | `INTEGER` | 是 | `` | — |
| `owner` | `VARCHAR(128)` | 是 | `` | — |
| `model` | `VARCHAR(64)` | 是 | `` | — |
| `client_ip` | `VARCHAR(64)` | 是 | `` | — |
| `request_body` | `TEXT` | 是 | `` | — |
| `response_body` | `TEXT` | 是 | `` | — |
| `request_body_ref` | `VARCHAR(512)` | 是 | `` | — |
| `response_body_ref` | `VARCHAR(512)` | 是 | `` | — |
| `request_body_size` | `INTEGER` | 是 | `0` | — |
| `response_body_size` | `INTEGER` | 是 | `0` | — |
| `request_body_sha256` | `VARCHAR(64)` | 是 | `` | — |
| `response_body_sha256` | `VARCHAR(64)` | 是 | `` | — |
| `prompt_tokens` | `INTEGER` | 是 | `0` | — |
| `completion_tokens` | `INTEGER` | 是 | `0` | — |
| `filter_hits` | `TEXT` | 是 | `[]` | — |
| `outcome` | `VARCHAR(16)` | 是 | `pass` | — |
| `blocked` | `BOOLEAN` | 是 | `False` | — |
| `status_code` | `INTEGER` | 是 | `200` | — |
| `latency_ms` | `INTEGER` | 是 | `0` | — |
| `search_text` | `TEXT` | 是 | `` | — |
### `audit_search_records`
| 字段 | 旧类型 | 可空 | 默认值 | 键/引用 |
|---|---|---:|---|---|
| `id` | `INTEGER` | 否 | `` | PK |
| `viewer` | `VARCHAR(128)` | 否 | `` | — |
| `query_hash` | `VARCHAR(64)` | 否 | `` | — |
| `start_at` | `DATETIME` | 否 | `` | — |
| `end_at` | `DATETIME` | 否 | `` | — |
| `owner_filter` | `VARCHAR(128)` | 是 | `` | — |
| `model_filter` | `VARCHAR(128)` | 是 | `` | — |
| `result_count` | `INTEGER` | 是 | `0` | — |
| `duration_ms` | `INTEGER` | 是 | `0` | — |
| `ts` | `DATETIME` | 是 | `now` | — |
### `audit_settings`
| 字段 | 旧类型 | 可空 | 默认值 | 键/引用 |
|---|---|---:|---|---|
| `id` | `INTEGER` | 否 | `` | PK |
| `image_audit_enabled` | `BOOLEAN` | 是 | `True` | — |
| `media_audit_async` | `BOOLEAN` | 是 | `False` | — |
| `text_audit_enabled` | `BOOLEAN` | 是 | `False` | — |
| `deep_scan_enabled` | `BOOLEAN` | 是 | `False` | — |
| `text_audit_api_base` | `VARCHAR(256)` | 是 | `` | — |
| `text_audit_api_key` | `VARCHAR(512)` | 是 | `` | — |
| `text_audit_model` | `VARCHAR(128)` | 是 | `` | — |
| `image_audit_api_base` | `VARCHAR(256)` | 是 | `` | — |
| `image_audit_api_key` | `VARCHAR(512)` | 是 | `` | — |
| `image_audit_model` | `VARCHAR(128)` | 是 | `` | — |
| `updated_at` | `DATETIME` | 是 | `now` | — |
### `content_audit_scopes`
| 字段 | 旧类型 | 可空 | 默认值 | 键/引用 |
|---|---|---:|---|---|
| `id` | `INTEGER` | 否 | `` | PK |
| `kind` | `VARCHAR(16)` | 否 | `` | — |
| `scope` | `VARCHAR(32)` | 否 | `` | — |
| `enabled` | `BOOLEAN` | 是 | `True` | — |
| `mode` | `VARCHAR(16)` | 是 | `async` | — |
| `created_at` | `DATETIME` | 是 | `now` | — |
### `departments`
| 字段 | 旧类型 | 可空 | 默认值 | 键/引用 |
|---|---|---:|---|---|
| `id` | `INTEGER` | 否 | `` | PK |
| `name` | `VARCHAR(64)` | 否 | `` | UNIQUE |
| `created_at` | `DATETIME` | 是 | `now` | — |
### `dept_budgets`
| 字段 | 旧类型 | 可空 | 默认值 | 键/引用 |
|---|---|---:|---|---|
| `id` | `INTEGER` | 否 | `` | PK |
| `dept_id` | `INTEGER` | 否 | `` | UNIQUEFK → `departments.id` |
| `monthly_point_budget` | `FLOAT` | 是 | `0.0` | — |
| `alert_threshold_pct` | `INTEGER` | 是 | `80` | — |
| `notify_enabled` | `BOOLEAN` | 是 | `True` | — |
| `last_alert_period` | `VARCHAR(7)` | 是 | `` | — |
| `last_alert_threshold` | `INTEGER` | 是 | `0` | — |
| `updated_at` | `DATETIME` | 是 | `now` | — |
### `feishu_auth_settings`
| 字段 | 旧类型 | 可空 | 默认值 | 键/引用 |
|---|---|---:|---|---|
| `id` | `INTEGER` | 否 | `` | PK |
| `enabled` | `BOOLEAN` | 是 | `False` | — |
| `app_id` | `VARCHAR(128)` | 是 | `` | — |
| `app_secret` | `VARCHAR(512)` | 是 | `` | — |
| `default_dept` | `VARCHAR(128)` | 是 | `飞书用户` | — |
| `allowed_tenants` | `TEXT` | 是 | `[]` | — |
| `provisioning_mode` | `VARCHAR(16)` | 是 | `pending` | — |
| `updated_at` | `DATETIME` | 是 | `now` | — |
### `filter_rules`
| 字段 | 旧类型 | 可空 | 默认值 | 键/引用 |
|---|---|---:|---|---|
| `id` | `INTEGER` | 否 | `` | PK |
| `name` | `VARCHAR(64)` | 否 | `` | — |
| `rule_type` | `VARCHAR(16)` | 是 | `keyword` | — |
| `pattern` | `TEXT` | 否 | `` | — |
| `action` | `VARCHAR(16)` | 是 | `block` | — |
| `severity` | `VARCHAR(8)` | 是 | `` | — |
| `enabled` | `BOOLEAN` | 是 | `True` | — |
| `builtin` | `BOOLEAN` | 是 | `False` | — |
| `scope` | `VARCHAR(32)` | 是 | `global` | — |
| `dimension` | `VARCHAR(32)` | 是 | `sensitive_content` | — |
| `created_at` | `DATETIME` | 是 | `now` | — |
### `hallucination_events`
| 字段 | 旧类型 | 可空 | 默认值 | 键/引用 |
|---|---|---:|---|---|
| `id` | `INTEGER` | 否 | `` | PK |
| `request_id` | `VARCHAR(64)` | 是 | `` | — |
| `key_id` | `INTEGER` | 是 | `` | FK → `virtual_keys.id` |
| `application_id` | `INTEGER` | 是 | `` | FK → `applications.id` |
| `dept_id` | `INTEGER` | 是 | `` | FK → `departments.id` |
| `model` | `VARCHAR(128)` | 是 | `` | — |
| `mode` | `VARCHAR(16)` | 是 | `async` | — |
| `action` | `VARCHAR(16)` | 是 | `observe` | — |
| `verdict` | `VARCHAR(24)` | 是 | `insufficient` | — |
| `support_score` | `INTEGER` | 是 | `0` | — |
| `question` | `TEXT` | 是 | `` | — |
| `answer` | `TEXT` | 是 | `` | — |
| `claims` | `TEXT` | 是 | `[]` | — |
| `evidence` | `TEXT` | 是 | `[]` | — |
| `latency_ms` | `INTEGER` | 是 | `0` | — |
| `error` | `VARCHAR(512)` | 是 | `` | — |
| `created_at` | `DATETIME` | 是 | `now` | — |
### `hallucination_policies`
| 字段 | 旧类型 | 可空 | 默认值 | 键/引用 |
|---|---|---:|---|---|
| `id` | `INTEGER` | 否 | `` | PK |
| `scope` | `VARCHAR(32)` | 否 | `global` | — |
| `enabled` | `BOOLEAN` | 是 | `False` | — |
| `mode` | `VARCHAR(16)` | 是 | `async` | — |
| `action` | `VARCHAR(16)` | 是 | `observe` | — |
| `kb_ids` | `TEXT` | 是 | `[]` | — |
| `support_threshold` | `INTEGER` | 是 | `70` | — |
| `evidence_threshold` | `FLOAT` | 是 | `0.35` | — |
| `top_k` | `INTEGER` | 是 | `4` | — |
| `max_claims` | `INTEGER` | 是 | `8` | — |
| `created_at` | `DATETIME` | 是 | `now` | — |
| `updated_at` | `DATETIME` | 是 | `now` | — |
### `hallucination_settings`
| 字段 | 旧类型 | 可空 | 默认值 | 键/引用 |
|---|---|---:|---|---|
| `id` | `INTEGER` | 否 | `1` | PK |
| `api_base` | `VARCHAR(256)` | 是 | `` | — |
| `api_key` | `VARCHAR(512)` | 是 | `` | — |
| `model` | `VARCHAR(128)` | 是 | `` | — |
| `timeout_seconds` | `INTEGER` | 是 | `15` | — |
| `updated_at` | `DATETIME` | 是 | `now` | — |
### `internal_service_keys`
| 字段 | 旧类型 | 可空 | 默认值 | 键/引用 |
|---|---|---:|---|---|
| `id` | `INTEGER` | 否 | `` | PK |
| `virtual_key_id` | `INTEGER` | 是 | `` | FK → `virtual_keys.id` |
| `raw_key_encrypted` | `VARCHAR(512)` | 是 | `` | — |
| `created_at` | `DATETIME` | 是 | `now` | — |
### `knowledge_bases`
| 字段 | 旧类型 | 可空 | 默认值 | 键/引用 |
|---|---|---:|---|---|
| `id` | `INTEGER` | 否 | `` | PK |
| `name` | `VARCHAR(128)` | 否 | `` | UNIQUE |
| `description` | `TEXT` | 是 | `` | — |
| `embedding_model` | `VARCHAR(128)` | 否 | `` | — |
| `chunk_size` | `INTEGER` | 是 | `800` | — |
| `chunk_overlap` | `INTEGER` | 是 | `100` | — |
| `dept_scope` | `TEXT` | 是 | `[]` | — |
| `enabled` | `BOOLEAN` | 是 | `True` | — |
| `created_by` | `VARCHAR(128)` | 是 | `` | — |
| `created_at` | `DATETIME` | 是 | `now` | — |
| `updated_at` | `DATETIME` | 是 | `now` | — |
### `knowledge_chunks`
| 字段 | 旧类型 | 可空 | 默认值 | 键/引用 |
|---|---|---:|---|---|
| `id` | `INTEGER` | 否 | `` | PK |
| `kb_id` | `INTEGER` | 否 | `` | FK → `knowledge_bases.id` |
| `doc_id` | `INTEGER` | 否 | `` | FK → `knowledge_documents.id` |
| `chunk_index` | `INTEGER` | 是 | `0` | — |
| `content` | `TEXT` | 是 | `` | — |
| `embedding` | `TEXT` | 是 | `[]` | — |
| `token_count` | `INTEGER` | 是 | `0` | — |
| `created_at` | `DATETIME` | 是 | `now` | — |
### `knowledge_documents`
| 字段 | 旧类型 | 可空 | 默认值 | 键/引用 |
|---|---|---:|---|---|
| `id` | `INTEGER` | 否 | `` | PK |
| `kb_id` | `INTEGER` | 否 | `` | FK → `knowledge_bases.id` |
| `filename` | `VARCHAR(256)` | 是 | `` | — |
| `source_type` | `VARCHAR(16)` | 是 | `upload` | — |
| `status` | `VARCHAR(16)` | 是 | `pending` | — |
| `status_message` | `TEXT` | 是 | `` | — |
| `content` | `TEXT` | 是 | `` | — |
| `char_count` | `INTEGER` | 是 | `0` | — |
| `chunk_count` | `INTEGER` | 是 | `0` | — |
| `uploaded_by` | `VARCHAR(128)` | 是 | `` | — |
| `created_at` | `DATETIME` | 是 | `now` | — |
| `updated_at` | `DATETIME` | 是 | `now` | — |
### `managed_conversations`
| 字段 | 旧类型 | 可空 | 默认值 | 键/引用 |
|---|---|---:|---|---|
| `id` | `VARCHAR(48)` | 否 | `` | PK |
| `user_id` | `INTEGER` | 否 | `` | FK → `users.id` |
| `application_id` | `INTEGER` | 否 | `` | FK → `applications.id` |
| `title` | `VARCHAR(160)` | 是 | `新会话` | — |
| `status` | `VARCHAR(16)` | 否 | `active` | — |
| `next_seq` | `INTEGER` | 否 | `1` | — |
| `head_hash` | `VARCHAR(64)` | 否 | `` | — |
| `busy` | `BOOLEAN` | 否 | `False` | — |
| `busy_token` | `VARCHAR(48)` | 否 | `` | — |
| `busy_since` | `DATETIME` | 是 | `` | — |
| `created_at` | `DATETIME` | 是 | `now` | — |
| `updated_at` | `DATETIME` | 是 | `now` | — |
### `managed_messages`
| 字段 | 旧类型 | 可空 | 默认值 | 键/引用 |
|---|---|---:|---|---|
| `id` | `INTEGER` | 否 | `` | PK |
| `conversation_id` | `VARCHAR(48)` | 否 | `` | FK → `managed_conversations.id` |
| `seq` | `INTEGER` | 否 | `` | — |
| `role` | `VARCHAR(16)` | 否 | `` | — |
| `content` | `TEXT` | 否 | `` | — |
| `content_sha256` | `VARCHAR(64)` | 否 | `` | — |
| `prev_hash` | `VARCHAR(64)` | 否 | `` | — |
| `chain_hash` | `VARCHAR(64)` | 否 | `` | — |
| `created_at` | `DATETIME` | 是 | `now` | — |
### `model_aliases`
| 字段 | 旧类型 | 可空 | 默认值 | 键/引用 |
|---|---|---:|---|---|
| `id` | `INTEGER` | 否 | `` | PK |
| `alias` | `VARCHAR(128)` | 否 | `` | UNIQUE |
| `target_model` | `VARCHAR(128)` | 否 | `` | — |
| `enabled` | `BOOLEAN` | 是 | `True` | — |
| `note` | `VARCHAR(256)` | 是 | `` | — |
| `created_at` | `DATETIME` | 是 | `now` | — |
| `updated_at` | `DATETIME` | 是 | `now` | — |
### `model_pricing`
| 字段 | 旧类型 | 可空 | 默认值 | 键/引用 |
|---|---|---:|---|---|
| `id` | `INTEGER` | 否 | `` | PK |
| `model` | `VARCHAR(128)` | 否 | `` | — |
| `input_price_per_1k` | `FLOAT` | 是 | `0.0` | — |
| `output_price_per_1k` | `FLOAT` | 是 | `0.0` | — |
| `effective_from` | `DATETIME` | 是 | `now` | — |
| `note` | `VARCHAR(256)` | 是 | `` | — |
| `created_at` | `DATETIME` | 是 | `now` | — |
### `model_qps_configs`
| 字段 | 旧类型 | 可空 | 默认值 | 键/引用 |
|---|---|---:|---|---|
| `id` | `INTEGER` | 否 | `` | PK |
| `model` | `VARCHAR(128)` | 否 | `` | UNIQUE |
| `qps` | `INTEGER` | 否 | `` | — |
| `enabled` | `BOOLEAN` | 是 | `True` | — |
| `updated_at` | `DATETIME` | 是 | `now` | — |
### `model_requests`
| 字段 | 旧类型 | 可空 | 默认值 | 键/引用 |
|---|---|---:|---|---|
| `id` | `INTEGER` | 否 | `` | PK |
| `ts` | `DATETIME` | 是 | `now` | — |
| `user_id` | `INTEGER` | 否 | `` | FK → `users.id` |
| `dept_id` | `INTEGER` | 是 | `` | — |
| `user_snapshot` | `VARCHAR(128)` | 是 | `` | — |
| `model_name` | `VARCHAR(128)` | 否 | `` | — |
| `provider_name` | `VARCHAR(64)` | 是 | `` | — |
| `purpose` | `TEXT` | 是 | `` | — |
| `quota_required` | `INTEGER` | 是 | `` | — |
| `status` | `VARCHAR(16)` | 是 | `pending` | — |
| `target_key_id` | `INTEGER` | 是 | `` | — |
| `reviewer` | `VARCHAR(128)` | 是 | `` | — |
| `review_comment` | `TEXT` | 是 | `` | — |
| `reviewed_at` | `DATETIME` | 是 | `` | — |
| `created_at` | `DATETIME` | 是 | `now` | — |
### `notify_config`
| 字段 | 旧类型 | 可空 | 默认值 | 键/引用 |
|---|---|---:|---|---|
| `id` | `INTEGER` | 否 | `` | PK |
| `enabled` | `BOOLEAN` | 是 | `False` | — |
| `webhook_url` | `VARCHAR(512)` | 是 | `` | — |
| `signing_secret` | `VARCHAR(512)` | 是 | `` | — |
| `notify_actions` | `VARCHAR(32)` | 是 | `block,redact` | — |
| `dept_scope` | `TEXT` | 是 | `[]` | — |
| `updated_at` | `DATETIME` | 是 | `now` | — |
### `pricing_config`
| 字段 | 旧类型 | 可空 | 默认值 | 键/引用 |
|---|---|---:|---|---|
| `id` | `INTEGER` | 否 | `` | PK |
| `points_per_yuan` | `FLOAT` | 是 | `100.0` | — |
| `overhead_factor` | `FLOAT` | 是 | `1.0` | — |
| `updated_at` | `DATETIME` | 是 | `now` | — |
### `prompt_categories`
| 字段 | 旧类型 | 可空 | 默认值 | 键/引用 |
|---|---|---:|---|---|
| `id` | `INTEGER` | 否 | `` | PK |
| `name` | `VARCHAR(64)` | 否 | `` | UNIQUE |
| `description` | `VARCHAR(256)` | 是 | `` | — |
| `created_at` | `DATETIME` | 是 | `now` | — |
### `prompt_favorites`
| 字段 | 旧类型 | 可空 | 默认值 | 键/引用 |
|---|---|---:|---|---|
| `id` | `INTEGER` | 否 | `` | PK |
| `user_id` | `INTEGER` | 否 | `` | FK → `users.id` |
| `template_id` | `INTEGER` | 否 | `` | FK → `prompt_templates.id` |
| `created_at` | `DATETIME` | 是 | `now` | — |
### `prompt_templates`
| 字段 | 旧类型 | 可空 | 默认值 | 键/引用 |
|---|---|---:|---|---|
| `id` | `INTEGER` | 否 | `` | PK |
| `name` | `VARCHAR(128)` | 否 | `` | UNIQUE |
| `description` | `TEXT` | 是 | `` | — |
| `category_id` | `INTEGER` | 是 | `` | FK → `prompt_categories.id` |
| `tags` | `TEXT` | 是 | `[]` | — |
| `dept_scope` | `TEXT` | 是 | `[]` | — |
| `current_version_id` | `INTEGER` | 是 | `` | — |
| `use_count` | `INTEGER` | 是 | `0` | — |
| `enabled` | `BOOLEAN` | 是 | `True` | — |
| `created_by` | `VARCHAR(128)` | 是 | `` | — |
| `created_at` | `DATETIME` | 是 | `now` | — |
| `updated_at` | `DATETIME` | 是 | `now` | — |
### `prompt_versions`
| 字段 | 旧类型 | 可空 | 默认值 | 键/引用 |
|---|---|---:|---|---|
| `id` | `INTEGER` | 否 | `` | PK |
| `template_id` | `INTEGER` | 否 | `` | FK → `prompt_templates.id` |
| `version_number` | `INTEGER` | 否 | `` | — |
| `content` | `TEXT` | 否 | `` | — |
| `variables` | `TEXT` | 是 | `[]` | — |
| `change_note` | `VARCHAR(256)` | 是 | `` | — |
| `created_by` | `VARCHAR(128)` | 是 | `` | — |
| `created_at` | `DATETIME` | 是 | `now` | — |
### `providers`
| 字段 | 旧类型 | 可空 | 默认值 | 键/引用 |
|---|---|---:|---|---|
| `id` | `INTEGER` | 否 | `` | PK |
| `name` | `VARCHAR(64)` | 否 | `` | UNIQUE |
| `base_url` | `VARCHAR(256)` | 否 | `` | — |
| `api_key` | `VARCHAR(512)` | 是 | `` | — |
| `models` | `TEXT` | 是 | `[]` | — |
| `backends` | `TEXT` | 是 | `` | — |
| `enabled` | `BOOLEAN` | 是 | `True` | — |
| `created_at` | `DATETIME` | 是 | `now` | — |
### `rate_limit_rules`
| 字段 | 旧类型 | 可空 | 默认值 | 键/引用 |
|---|---|---:|---|---|
| `id` | `INTEGER` | 否 | `` | PK |
| `name` | `VARCHAR(64)` | 否 | `` | — |
| `enabled` | `BOOLEAN` | 是 | `True` | — |
| `dimension` | `VARCHAR(16)` | 是 | `consumer` | — |
| `field_name` | `VARCHAR(64)` | 是 | `` | — |
| `match_type` | `VARCHAR(16)` | 是 | `any` | — |
| `pattern` | `VARCHAR(256)` | 是 | `` | — |
| `limit_count` | `INTEGER` | 否 | `60` | — |
| `window_seconds` | `INTEGER` | 是 | `60` | — |
| `created_at` | `DATETIME` | 是 | `now` | — |
| `updated_at` | `DATETIME` | 是 | `now` | — |
### `routing_groups`
| 字段 | 旧类型 | 可空 | 默认值 | 键/引用 |
|---|---|---:|---|---|
| `id` | `INTEGER` | 否 | `` | PK |
| `virtual_model` | `VARCHAR(128)` | 否 | `` | UNIQUE |
| `strategy` | `VARCHAR(16)` | 是 | `priority` | — |
| `description` | `VARCHAR(256)` | 是 | `` | — |
| `enabled` | `BOOLEAN` | 是 | `True` | — |
| `created_by` | `VARCHAR(128)` | 是 | `` | — |
| `created_at` | `DATETIME` | 是 | `now` | — |
| `updated_at` | `DATETIME` | 是 | `now` | — |
### `routing_targets`
| 字段 | 旧类型 | 可空 | 默认值 | 键/引用 |
|---|---|---:|---|---|
| `id` | `INTEGER` | 否 | `` | PK |
| `group_id` | `INTEGER` | 否 | `` | FK → `routing_groups.id` |
| `target_model` | `VARCHAR(128)` | 否 | `` | — |
| `priority` | `INTEGER` | 是 | `100` | — |
| `capability_tags` | `TEXT` | 是 | `[]` | — |
| `enabled` | `BOOLEAN` | 是 | `True` | — |
| `created_at` | `DATETIME` | 是 | `now` | — |
### `tool_definitions`
| 字段 | 旧类型 | 可空 | 默认值 | 键/引用 |
|---|---|---:|---|---|
| `id` | `INTEGER` | 否 | `` | PK |
| `name` | `VARCHAR(128)` | 否 | `` | UNIQUE |
| `description` | `TEXT` | 是 | `` | — |
| `endpoint_url` | `VARCHAR(512)` | 否 | `` | — |
| `http_method` | `VARCHAR(8)` | 是 | `POST` | — |
| `headers_template` | `TEXT` | 是 | `` | — |
| `input_schema` | `TEXT` | 是 | `{}` | — |
| `timeout_seconds` | `INTEGER` | 是 | `15` | — |
| `dept_scope` | `TEXT` | 是 | `[]` | — |
| `enabled` | `BOOLEAN` | 是 | `True` | — |
| `created_by` | `VARCHAR(128)` | 是 | `` | — |
| `created_at` | `DATETIME` | 是 | `now` | — |
| `updated_at` | `DATETIME` | 是 | `now` | — |
### `users`
| 字段 | 旧类型 | 可空 | 默认值 | 键/引用 |
|---|---|---:|---|---|
| `id` | `INTEGER` | 否 | `` | PK |
| `account` | `VARCHAR(128)` | 否 | `` | UNIQUE |
| `name` | `VARCHAR(64)` | 是 | `` | — |
| `dept_id` | `INTEGER` | 是 | `` | FK → `departments.id` |
| `password_hash` | `VARCHAR(128)` | 是 | `` | — |
| `feishu_open_id` | `VARCHAR(128)` | 是 | `` | UNIQUE |
| `feishu_union_id` | `VARCHAR(128)` | 是 | `` | — |
| `auth_source` | `VARCHAR(16)` | 是 | `local` | — |
| `active` | `BOOLEAN` | 是 | `True` | — |
| `provisioning_status` | `VARCHAR(16)` | 是 | `active` | — |
| `failed_logins` | `INTEGER` | 是 | `0` | — |
| `locked_until` | `DATETIME` | 是 | `` | — |
| `deleted_name_snapshot` | `VARCHAR(128)` | 是 | `` | — |
| `created_at` | `DATETIME` | 是 | `now` | — |
| `totp_secret` | `VARCHAR(255)` | 是 | `` | — |
| `totp_enabled` | `BOOLEAN` | 是 | `False` | — |
| `totp_last_step` | `INTEGER` | 是 | `` | — |
| `totp_backup_codes` | `TEXT` | 是 | `[]` | — |
| `totp_confirmed_at` | `DATETIME` | 是 | `` | — |
### `virtual_keys`
| 字段 | 旧类型 | 可空 | 默认值 | 键/引用 |
|---|---|---:|---|---|
| `id` | `INTEGER` | 否 | `` | PK |
| `key_hash` | `VARCHAR(64)` | 否 | `` | UNIQUE |
| `key_prefix` | `VARCHAR(16)` | 是 | `` | — |
| `name` | `VARCHAR(128)` | 是 | `` | — |
| `user_id` | `INTEGER` | 是 | `` | FK → `users.id` |
| `application_id` | `INTEGER` | 是 | `` | FK → `applications.id` |
| `owner` | `VARCHAR(128)` | 是 | `` | — |
| `allowed_models` | `TEXT` | 是 | `[]` | — |
| `quota_tokens` | `INTEGER` | 是 | `` | — |
| `used_tokens` | `INTEGER` | 是 | `0` | — |
| `rpm_limit` | `INTEGER` | 是 | `` | — |
| `active` | `BOOLEAN` | 是 | `True` | — |
| `expires_at` | `DATETIME` | 是 | `` | — |
| `created_at` | `DATETIME` | 是 | `now` | — |
+42
View File
@@ -0,0 +1,42 @@
# 旧数据导入手册
当前工程提供两阶段、可重复执行的迁移入口,避免直接把旧 SQLite 密文或不完整外键写进新业务表。
## 1. 只读导出
在旧系统停止配置写入后,对数据库快照执行:
```bash
python3 scripts/export_legacy_data.py /secure-snapshot/gateway.db > /secure-transfer/gateway.jsonl
```
默认不导出密码、API Key、TOTP、请求头等敏感列。需要迁移 Provider/工具等凭据时,只能在隔离迁移环境使用 `--include-encrypted`;旧密文不能作为新密文直接写入,必须由专用转换步骤用旧密钥解密,再用新的用途标签和活动 KEK 加密。
## 2. 校验和暂存
先执行不连接数据库的校验:
```bash
gateway-legacy-import --input /secure-transfer/gateway.jsonl --dry-run
```
再写入 PostgreSQL 不可变暂存区:
```bash
export DATABASE_URL='postgres://...'
gateway-legacy-import --input /secure-transfer/gateway.jsonl
```
每条记录会验证 SHA-256,使用固定 namespace 生成 UUIDv5,并同时写入 `legacy_import_records``legacy_id_mappings`。相同快照重复执行会复用 batch;同一旧 ID 的内容发生变化时失败关闭,不会静默覆盖先前暂存数据。
## 3. 领域转换与验收
暂存成功不等于已写入生产业务表。必须按外键拓扑执行领域转换:部门 → 身份 → Provider/API Key → Prompt/知识/工具 → 应用版本 → 审计历史。每批转换后核对:
- 源记录数、目标记录数、明确跳过数和失败数之和相等;
- 所有旧外键都能从 `legacy_id_mappings` 找到映射;
- Provider、工具、通知、TOTP 等密文可使用新 keyring 解密,且数据库中不存在旧明文;
- 审计只迁移有界预览与结构化元数据,不迁移无限长原始正文;
- 应用发布版本引用的 Prompt、知识库和工具均存在且启用。
当前仓库没有用户的实际 SQLite 数据文件,因此已完成的是可验证暂存管道;生产领域转换必须在拿到脱敏快照及旧密钥托管授权后执行,不能凭空假设旧密文格式。
+115
View File
@@ -0,0 +1,115 @@
# 重构进度
更新时间:2026-08-11
## 已完成:M0 工程基线
- Go 1.26 模块、配置校验、结构化日志和优雅退出。
- PostgreSQL `pgx` 连接池、critical/cache 双 Redis 客户端。
- 独立且带校验和、事务和 advisory lock 的数据库迁移器。
- OpenAI 兼容入口、bootstrap key 鉴权、请求体限制、连接池和 SSE 透传。
- `healthz`、依赖感知 `readyz`、基础 Prometheus 指标。
- PostgreSQL outbox、API Key、Provider 与分区审计表的首版 schema。
- 不包含 MinIO/S3、ClickHouse 的 Docker Compose 基线。
- Art Design Pro v3.0.2 管理端与门户端,已执行官方精简流程并通过生产构建。
- OpenAPI 运行时契约和共享 TypeScript 系统客户端起点。
## 已完成:M1 身份、Provider 与配置面
已完成:
- 管理员、门户用户 schema 和独立 bootstrap 命令。
- 兼容 Python 版本 60 万轮/旧 12 万轮 PBKDF2-SHA256,并在成功登录时升级旧哈希。
- critical Redis 可撤销不透明会话、数据库 active/role 复核、失败次数原子更新和账号锁定。
- `/api/v1/admin/*``/api/v1/portal/*` 登录、身份、菜单、退出契约。
- 管理端与门户端 TOTP 配置、二次登录、停用、备用码重置;含临时挑战令牌、时间步原子防重放和备用码一次性消费。
- 两套 Art 登录页已接入动态验证码/备用码,并通过生产构建。
- 数据库 API Key 创建、列表、scope、过期、Redis 共享缓存和即时撤销;明文只显示一次,创建/撤销与 outbox 同事务。
- Art 管理端 API Key 页面、动态菜单入口及生产构建。
- Provider AES-256-GCM 凭据、SSRF 地址校验、revision 和事务 outbox。
- Provider 数据库运行时原子快照、默认/显式路由、能力约束、共享连接池、周期刷新和上一有效版本回退。
- Art 管理端供应商页面及生产后端权限模式。
- Provider 连接测试、模型目录同步、历史模型停用策略和 `provider.models_synced` 事务 outbox。
- Art 管理端连接测试、同步结果和模型目录查看交互。
- 版本化 KEK keyring、历史密钥解密、Provider 凭据事务轮换和管理端轮换入口。
- PostgreSQL 权威 Provider 配置、critical Redis 多实例变更通知、进程内不可变快照与周期兜底刷新。
- 管理员/门户账号管理、内置角色、直接权限字符串、权限感知动态菜单和即时数据库复核。
- bootstrap API Key 显式关闭开关及兼容入口使用量指标。
- 层级部门 CRUD、循环检测、在用停用保护、门户账号部门绑定和部门 outbox。
- OIDC 配置面、加密 Client Secret、Authorization Code + PKCE、RS256/JWKS 校验、自动开户和一次性交换码。
- SAML 配置面、SP metadata、SP-initiated Redirect/POST 流程、SHA-256+ XML 签名断言校验、RelayState/InResponseTo 绑定、断言防重放、自动开户和部门绑定。
- 旧 Python 网关 38 个持久化实体字段字典、实体去向清单、确定性 UUIDv5 生成器和幂等旧 ID 映射表。
- SAML 与 OIDC 均已通过 PostgreSQL + Redis + 本地模拟 IdP 的端到端登录、自动开户、一次性交换和重放拦截测试。
## 已完成:M2 流量治理与路由
已完成:
- API Key 级每分钟请求上限和 UTC 自然月请求配额,`0` 明确表示不限制。
- critical Redis Lua 原子准入;限流返回 `429`/`Retry-After`/`X-RateLimit-*`,配置限制时 Redis 故障失败关闭。
- API Key 限额的 PostgreSQL 权威配置、鉴权缓存传递、Art 管理表单与 OpenAPI 契约。
- PostgreSQL + Redis + 模拟上游端到端测试:RPM=2 精确放行 2 次,月配额=3 精确放行 3 次。
- API Key UTC 自然月 Token 配额:请求前按输入体积与最大输出预算原子预留,响应结束按上游 usage 原子校准。
- OpenAI/Responses/Anthropic 的 JSON 与 SSE usage 增量提取;不缓冲完整流式响应,上游未返回 usage 时按预留预算保守计量。
- Art 管理端展示当月 Token 用量并可即时调整 RPM、月请求和月 Token 配额;更新同时清除共享鉴权缓存并写入 outbox。
- Provider 级响应头超时、有限指数退避和熔断;仅 GET/HEAD 或携带 `Idempotency-Key` 的可重放请求允许重试,所有阈值均可通过环境变量配置。
- PostgreSQL 模型别名/路由规则、事务 outbox、管理 CRUD、Redis 多实例变更通知和进程内不可变路由快照。
- 按 endpoint、API Key、tenant 条件匹配;最高优先级分组内使用请求 ID 做确定性加权选择,显式 Provider 作为候选约束。
- 已登记的模型别名在条件不匹配、规则停用或显式 Provider 不属于候选集时失败关闭,避免把别名透传到默认上游绕过策略。
- 代理在命中规则时改写上游模型并返回 `X-Gateway-Model`;未配置规则时不读取请求体,保持原代理热路径性能。
- Art 管理端模型路由页面,可维护别名、上游模型、Provider、权重、优先级、条件和启停状态。
## 已完成:M3 审计、usage、内容策略与成本
已完成:
- 网关调用的有界异步审计队列、批量 PostgreSQL COPY、失败保留/重试、队列满降级和优雅停机排空。
- 请求 ID、API Key、tenant、Provider、原始/目标模型、协议、HTTP 状态、延迟及输入/输出 Token 采集。
- PostgreSQL 按日 API Key/Provider/模型 usage 聚合,不在请求热路径同步写数据库。
- `audit:read`/`usage:read` 权限、筛选查询 API、Art 管理端审计与用量页面。
- 审计 accepted/dropped/written/flush failures Prometheus 指标。
- 独立 `gateway-outbox-worker`,使用 PostgreSQL `FOR UPDATE SKIP LOCKED` 租约支持多实例并行抢占。
- Redis Lua 原子事件 ID 去重与 Stream 写入,数据库确认丢失后重投不会产生重复 Stream 消息。
- 指数退避、最大尝试次数、死信状态、`outbox:read`/`outbox:manage` 权限、管理 API 与 Art 人工重试页面。
- `event_consumptions` 与业务 handler 共用 PostgreSQL 事务的消费者幂等入口,handler 回滚时消费标记同步回滚。
- 独立 `gateway-maintenance` 周期任务,按 UTC 月预创建审计分区,事务迁移 default 分区数据,并按配置直接丢弃完整过期分区、精确清理边界月和长期 usage。
- 审计维护使用 PostgreSQL advisory lock 与表级事务锁,多副本启动时仍只有一个实例执行;重复运行保持幂等。
- PostgreSQL 内容策略 CRUD、事务 outbox、`content_policy:read/manage` 权限与多实例周期刷新;运行时使用不可变 RE2 编译快照。
- 内容策略支持 `audit``block``redact`,按 endpoint、模型和 API Key 限定范围;只遍历提示词文本字段,不改写工具 schema、图片/Base64 等非文本载荷。
- 默认敏感信息脱敏策略、阻断响应、`X-Gateway-Content-Redacted` 响应头,以及不包含原始命中值的审计标签。
- Art 管理端内容策略页面,可维护优先级、范围、规则、替换文本和启停状态。
- 带生效/失效时间的模型价格版本,支持 Provider + 精确模型或末尾 `*` 前缀,精确规则优先并使用请求时刻选择版本。
- 输入/输出价格均以“每百万 Token 的微货币单位”精确存储;usage 完成时核算请求成本,审计与按日聚合均保存 `cost_microunits`
- `pricing:read/manage` 权限、模型价格管理 API/Art 页面,审计与用量页面展示估算成本。
## 已完成:M4 Prompt、知识、工具、应用编排和通知
- Prompt 分类/模板和只增不改的版本,显式变量声明、必填/默认值校验、历史版本激活、管理端和 API Key 渲染接口。
- PostgreSQL 知识库保存不超过 2 MiB 的提取文本,不依赖 MinIO/S3;段落感知重叠分块,FTS 与中文二元词片混合召回,并通过 `Retriever` 接口隔离未来向量实现。
- 声明式 HTTP 工具、基础 JSON Schema 入参验证、KEK 加密请求头、注册及拨号时 SSRF 校验、禁止重定向、1 MiB 响应上限和工具调用记录。
- AI 应用草稿、引用完整性校验、不可变发布版本和运行记录;应用入口组合 Prompt/RAG/OpenAI 工具调用循环,内部每个模型请求仍走原网关全治理链路。
- 通知通道与可靠投递记录,独立 Worker 消费 outbox,精确/前缀事件订阅、HMAC-SHA256 签名、失败留痕与人工重试。
- 内容策略命中由异步审计批处理在同一事务写入 `content_policy.matched` outbox,不在模型请求热路径同步投递 Webhook,也不包含原始敏感命中值。
- Art Design Pro 管理端新增 Prompt、知识库、工具、应用和通知五个页面,权限与动态菜单均由服务端控制。
- 独立前端多阶段镜像将 Art 管理端生产构建交付给 Nginx,并同源代理管理 API、模型 API、SSE 与健康检查;Compose 可直接启动完整管理面。
## 已完成:M5 契约、迁移与切换工程
- 通过 Python AST 自动盘点旧源码,确认实际为 201 个路由装饰器,而不是早期估算的 141 个;与 OpenAPI 生成可重复的覆盖/替代/退役/缺口矩阵。
- 可选影子中间件只复制确定性采样的非流式 JSON POST,使用独立影子 Key,不转发生产凭据;异步比较状态码和 JSON 结构签名并输出 Prometheus 指标。
- 独立 `gateway-loadtest` 支持并发、时长、超时、错误率和 p95 阈值,输出机器可读 JSON 并以退出码阻断不达标发布。
- 原子 Nginx upstream 切换脚本在目标探活、`nginx -t` 和 reload 任一步失败时恢复备份;提供明确的回退触发条件与演练手册。
- 旧 SQLite 只读 JSONL 导出、逐记录 SHA-256、确定性 UUIDv5 和 PostgreSQL 不可变暂存管道已完成;默认排除敏感列,相同快照幂等,旧记录变化失败关闭。
- 201 条旧路由已全部形成明确决策:85 条同契约覆盖、112 条由新契约替代、4 条基于安全或架构 ADR 退役、未决缺口 0;OpenAPI 0.10.0 覆盖全部 146 条 Go 字面量路由。
## 已完成:M6 门户工作台与扩展治理
- 门户资产目录、Prompt 查看/搜索/收藏、个人审计/Token/成本统计、接入说明和模型访问申请;所有列表按登录用户与部门范围服务端过滤。
- 管理端模型申请审批、事实核验 Provider 设置、作用域策略与核验事件页面;事实核验复用 Provider 加密凭据,不重复保存 API Key/Base URL。
- 已发布应用的门户单轮入口和服务端托管会话;运行凭证使用独立 AEAD purpose 加密,明文不返回浏览器。
- 会话串行租约避免并发回答交错,用户/助手消息以不可变顺序和 SHA-256 哈希链保存,重放前验证完整性,最多 200 条消息。
- 应用目录/资产依赖分析、历史版本回滚、知识文档重新分块、Prompt 详情、运行时快照刷新、系统信息和 24 小时监控汇总。
- Art Design Pro 管理端新增模型治理页面;门户端新增资产目录、Prompt 广场、个人用量与模型权限四个生产页面。
- Compose 同时交付 API、管理端 `8081` 和门户端 `8082`;同一前端 Dockerfile 通过受限 build arg 构建两套独立 SPA。
工程实现与本地生产等价验证已关闭。正式上线仍是外部发布门禁:拿到实际旧数据库脱敏快照和旧密钥迁移授权后执行领域转换,并在真实新旧双系统与生产等价流量下完成影子观察、容量验收及切换/回退演练;这些动作不会在缺少生产数据和授权时伪造为已执行。
+213
View File
@@ -0,0 +1,213 @@
# 旧 Python → Go 路由契约矩阵
> 由 `scripts/compare_route_contracts.py` 通过 AST 与 OpenAPI 自动生成;`replaced` 表示能力有新契约承接,仍需兼容层或客户端迁移,不能视为原路由原样可用。
- 旧源码路由装饰器:201
- 同方法同规范路径:85
- 新契约替代:112
- 明确退役:4
- 尚无契约:0
| 状态 | 方法 | 旧路径 | Go 契约/预期目标 | 来源 |
|---|---|---|---|---|
| 新契约替代 | `GET` | `/admin/api/admins` | `/api/v1/admin/identities/admins` | `admin.py:290` |
| 新契约替代 | `POST` | `/admin/api/admins` | `/api/v1/admin/identities/admins` | `admin.py:306` |
| 新契约替代 | `POST` | `/admin/api/admins/change_password` | `/api/v1/admin/password` | `admin.py:361` |
| 新契约替代 | `DELETE` | `/admin/api/admins/{admin_id}` | `/api/v1/admin/identities/*/{identity_id} active=false` | `admin.py:381` |
| 新契约替代 | `PATCH` | `/admin/api/admins/{admin_id}` | `/api/v1/admin/identities/admins/{identity_id} (PUT with full revision payload)` | `admin.py:329` |
| 已覆盖 | `GET` | `/admin/api/applications` | `/api/v1/admin/applications` | `applications.py:158` |
| 已覆盖 | `POST` | `/admin/api/applications` | `/api/v1/admin/applications` | `applications.py:287` |
| 已覆盖 | `GET` | `/admin/api/applications/asset-dependencies` | `/api/v1/admin/applications/asset-dependencies` | `applications.py:172` |
| 已覆盖 | `GET` | `/admin/api/applications/catalog` | `/api/v1/admin/applications/catalog` | `applications.py:165` |
| 已覆盖 | `DELETE` | `/admin/api/applications/{app_id}` | `/api/v1/admin/applications/{application_id}` | `applications.py:370` |
| 新契约替代 | `PATCH` | `/admin/api/applications/{app_id}` | `/api/v1/admin/applications/{application_id} (PUT with full revision payload)` | `applications.py:312` |
| 新契约替代 | `GET` | `/admin/api/applications/{app_id}/composition` | `/api/v1/admin/applications/{application_id}` | `applications.py:190` |
| 新契约替代 | `PUT` | `/admin/api/applications/{app_id}/composition` | `/api/v1/admin/applications/{application_id}` | `applications.py:202` |
| 新契约替代 | `POST` | `/admin/api/applications/{app_id}/keys` | `/api/v1/admin/api-keys (scoped credentials)` | `applications.py:384` |
| 已覆盖 | `POST` | `/admin/api/applications/{app_id}/publish` | `/api/v1/admin/applications/{application_id}/publish` | `applications.py:214` |
| 已覆盖 | `POST` | `/admin/api/applications/{app_id}/rollback/{version}` | `/api/v1/admin/applications/{application_id}/rollback/{version}` | `applications.py:233` |
| 已覆盖 | `GET` | `/admin/api/applications/{app_id}/runs` | `/api/v1/admin/applications/{application_id}/runs` | `applications.py:254` |
| 新契约替代 | `POST` | `/admin/api/applications/{app_id}/test` | `/v1/applications/{application_code}/chat/completions` | `applications.py:272` |
| 新契约替代 | `GET` | `/admin/api/audit_settings` | `/api/v1/admin/content-policies and /api/v1/admin/notification-channels` | `admin.py:2297` |
| 新契约替代 | `PUT` | `/admin/api/audit_settings` | `/api/v1/admin/content-policies and /api/v1/admin/notification-channels` | `admin.py:2340` |
| 新契约替代 | `POST` | `/admin/api/audit_settings/test_image_model` | `/api/v1/admin/content-policies and /api/v1/admin/notification-channels` | `admin.py:2379` |
| 新契约替代 | `GET` | `/admin/api/content_audit_scopes` | `/api/v1/admin/content-policies` | `admin.py:2413` |
| 新契约替代 | `POST` | `/admin/api/content_audit_scopes` | `/api/v1/admin/content-policies` | `admin.py:2436` |
| 新契约替代 | `DELETE` | `/admin/api/content_audit_scopes/{scope_id}` | `/api/v1/admin/content-policies` | `admin.py:2470` |
| 新契约替代 | `GET` | `/admin/api/cost_report` | `/api/v1/admin/usage/daily and /api/v1/admin/model-prices` | `admin.py:3936` |
| 新契约替代 | `GET` | `/admin/api/cost_report/export` | `/api/v1/admin/usage/daily and /api/v1/admin/model-prices` | `admin.py:3967` |
| 已覆盖 | `GET` | `/admin/api/departments` | `/api/v1/admin/departments` | `admin.py:550` |
| 已覆盖 | `POST` | `/admin/api/departments` | `/api/v1/admin/departments` | `admin.py:565` |
| 新契约替代 | `DELETE` | `/admin/api/departments/{dept_id}` | `/api/v1/admin/departments/{department_id} active=false` | `admin.py:615` |
| 新契约替代 | `PATCH` | `/admin/api/departments/{dept_id}` | `/api/v1/admin/departments/{department_id} (PUT with full revision payload)` | `admin.py:584` |
| 新契约替代 | `GET` | `/admin/api/dept_budgets` | `/api/v1/admin/usage/daily and /api/v1/admin/model-prices` | `admin.py:3781` |
| 新契约替代 | `POST` | `/admin/api/dept_budgets` | `/api/v1/admin/usage/daily and /api/v1/admin/model-prices` | `admin.py:3807` |
| 新契约替代 | `DELETE` | `/admin/api/dept_budgets/{budget_id}` | `/api/v1/admin/usage/daily and /api/v1/admin/model-prices` | `admin.py:3839` |
| 新契约替代 | `GET` | `/admin/api/dept_stats` | `/api/v1/admin/usage/daily and /api/v1/admin/model-prices` | `admin.py:2934` |
| 已覆盖 | `GET` | `/admin/api/fact-check/events` | `/api/v1/admin/fact-check/events` | `hallucination.py:390` |
| 已覆盖 | `GET` | `/admin/api/fact-check/events/{event_id}` | `/api/v1/admin/fact-check/events/{event_id}` | `hallucination.py:405` |
| 已覆盖 | `GET` | `/admin/api/fact-check/policies` | `/api/v1/admin/fact-check/policies` | `hallucination.py:348` |
| 已覆盖 | `POST` | `/admin/api/fact-check/policies` | `/api/v1/admin/fact-check/policies` | `hallucination.py:355` |
| 已覆盖 | `DELETE` | `/admin/api/fact-check/policies/{policy_id}` | `/api/v1/admin/fact-check/policies/{policy_id}` | `hallucination.py:378` |
| 已覆盖 | `GET` | `/admin/api/fact-check/settings` | `/api/v1/admin/fact-check/settings` | `hallucination.py:312` |
| 已覆盖 | `PUT` | `/admin/api/fact-check/settings` | `/api/v1/admin/fact-check/settings` | `hallucination.py:323` |
| 新契约替代 | `GET` | `/admin/api/feishu_auth` | `/api/v1/admin/identity-providers (generic OIDC/SAML)` | `admin.py:3314` |
| 新契约替代 | `POST` | `/admin/api/feishu_auth` | `/api/v1/admin/identity-providers (generic OIDC/SAML)` | `admin.py:3335` |
| 新契约替代 | `POST` | `/admin/api/feishu_auth/test` | `/api/v1/admin/identity-providers (generic OIDC/SAML)` | `admin.py:3364` |
| 新契约替代 | `GET` | `/admin/api/kb` | `/api/v1/admin/knowledge-bases` | `knowledge.py:395` |
| 新契约替代 | `POST` | `/admin/api/kb` | `/api/v1/admin/knowledge-bases` | `knowledge.py:410` |
| 新契约替代 | `DELETE` | `/admin/api/kb/{kb_id}` | `/api/v1/admin/knowledge-bases/{knowledge_base_id}` | `knowledge.py:475` |
| 新契约替代 | `PATCH` | `/admin/api/kb/{kb_id}` | `/api/v1/admin/knowledge-bases/{knowledge_base_id} (PUT with full revision payload)` | `knowledge.py:439` |
| 新契约替代 | `GET` | `/admin/api/kb/{kb_id}/documents` | `/api/v1/admin/knowledge-bases/{knowledge_base_id}/documents` | `knowledge.py:494` |
| 新契约替代 | `POST` | `/admin/api/kb/{kb_id}/documents/paste` | `/api/v1/admin/knowledge-bases/{knowledge_base_id}/documents` | `knowledge.py:550` |
| 新契约替代 | `POST` | `/admin/api/kb/{kb_id}/documents/upload` | `/api/v1/admin/knowledge-bases/{knowledge_base_id}/documents` | `knowledge.py:508` |
| 新契约替代 | `DELETE` | `/admin/api/kb/{kb_id}/documents/{doc_id}` | `/api/v1/admin/knowledge-bases/{knowledge_base_id}/documents/{document_id}` | `knowledge.py:590` |
| 新契约替代 | `POST` | `/admin/api/kb/{kb_id}/documents/{doc_id}/reprocess` | `/api/v1/admin/knowledge-bases/{knowledge_base_id}/documents/{document_id}/reprocess` | `knowledge.py:574` |
| 新契约替代 | `POST` | `/admin/api/kb/{kb_id}/test_search` | `/api/v1/admin/knowledge-bases/{knowledge_base_id}/search` | `knowledge.py:613` |
| 新契约替代 | `GET` | `/admin/api/keys` | `/api/v1/admin/api-keys` | `admin.py:912` |
| 新契约替代 | `POST` | `/admin/api/keys` | `/api/v1/admin/api-keys` | `admin.py:943` |
| 新契约替代 | `DELETE` | `/admin/api/keys/{key_id}` | `/api/v1/admin/api-keys/{api_key_id}` | `admin.py:1017` |
| 新契约替代 | `PATCH` | `/admin/api/keys/{key_id}` | `/api/v1/admin/api-keys/{api_key_id}/limits` | `admin.py:977` |
| 新契约替代 | `GET` | `/admin/api/log_search` | `/api/v1/admin/audit-events` | `admin.py:2817` |
| 新契约替代 | `GET` | `/admin/api/log_views` | `/api/v1/admin/audit-events` | `admin.py:2879` |
| 已覆盖 | `POST` | `/admin/api/login` | `/api/v1/admin/login` | `admin.py:151` |
| 已覆盖 | `POST` | `/admin/api/login/totp` | `/api/v1/admin/login/totp` | `admin.py:193` |
| 新契约替代 | `DELETE` | `/admin/api/logs` | `/api/v1/admin/audit-events (bounded immutable audit query)` | `admin.py:2633` |
| 新契约替代 | `GET` | `/admin/api/logs` | `/api/v1/admin/audit-events` | `admin.py:2540` |
| 新契约替代 | `GET` | `/admin/api/logs/search` | `/api/v1/admin/audit-events (bounded immutable audit query)` | `admin.py:2689` |
| 新契约替代 | `GET` | `/admin/api/logs/{log_id}/view_history` | `/api/v1/admin/audit-events (bounded immutable audit query)` | `admin.py:2862` |
| 明确退役 | `POST` | `/admin/api/logs/{log_id}/view_raw` | `retired: unbounded raw-body access is excluded by storage/compliance ADR` | `admin.py:2653` |
| 新契约替代 | `GET` | `/admin/api/malicious_domains` | `/api/v1/admin/content-policies` | `admin.py:2501` |
| 新契约替代 | `POST` | `/admin/api/malicious_domains` | `/api/v1/admin/content-policies` | `admin.py:2512` |
| 新契约替代 | `DELETE` | `/admin/api/malicious_domains/{domain}` | `/api/v1/admin/content-policies` | `admin.py:2528` |
| 已覆盖 | `GET` | `/admin/api/model-requests` | `/api/v1/admin/model-requests` | `admin.py:3506` |
| 已覆盖 | `POST` | `/admin/api/model-requests/{req_id}/approve` | `/api/v1/admin/model-requests/{request_id}/approve` | `admin.py:3522` |
| 已覆盖 | `POST` | `/admin/api/model-requests/{req_id}/reject` | `/api/v1/admin/model-requests/{request_id}/reject` | `admin.py:3569` |
| 新契约替代 | `GET` | `/admin/api/model_aliases` | `/api/v1/admin/model-routes` | `admin.py:2204` |
| 新契约替代 | `POST` | `/admin/api/model_aliases` | `/api/v1/admin/model-routes` | `admin.py:2224` |
| 新契约替代 | `DELETE` | `/admin/api/model_aliases/{alias_id}` | `/api/v1/admin/model-routes` | `admin.py:2262` |
| 新契约替代 | `PUT` | `/admin/api/model_aliases/{alias_id}` | `/api/v1/admin/model-routes` | `admin.py:2240` |
| 新契约替代 | `GET` | `/admin/api/model_pricing` | `/api/v1/admin/model-prices` | `admin.py:3615` |
| 新契约替代 | `POST` | `/admin/api/model_pricing` | `/api/v1/admin/model-prices` | `admin.py:3646` |
| 新契约替代 | `DELETE` | `/admin/api/model_pricing/{pricing_id}` | `/api/v1/admin/model-prices/{price_id}` | `admin.py:3677` |
| 新契约替代 | `GET` | `/admin/api/model_qps` | `/api/v1/admin/api-keys and /api/v1/admin/model-routes` | `admin.py:2114` |
| 新契约替代 | `POST` | `/admin/api/model_qps` | `/api/v1/admin/api-keys and /api/v1/admin/model-routes` | `admin.py:2125` |
| 新契约替代 | `DELETE` | `/admin/api/model_qps/{config_id}` | `/api/v1/admin/api-keys and /api/v1/admin/model-routes` | `admin.py:2147` |
| 已覆盖 | `GET` | `/admin/api/models` | `/api/v1/admin/models` | `admin.py:1093` |
| 已覆盖 | `GET` | `/admin/api/monitoring/overview` | `/api/v1/admin/monitoring/overview` | `admin.py:3153` |
| 新契约替代 | `GET` | `/admin/api/notify_config` | `/api/v1/admin/notification-channels` | `admin.py:3392` |
| 新契约替代 | `POST` | `/admin/api/notify_config` | `/api/v1/admin/notification-channels` | `admin.py:3415` |
| 新契约替代 | `POST` | `/admin/api/notify_config/test` | `/api/v1/admin/notification-channels` | `admin.py:3447` |
| 新契约替代 | `GET` | `/admin/api/pricing_config` | `/api/v1/admin/usage/daily and /api/v1/admin/model-prices` | `admin.py:3702` |
| 新契约替代 | `PUT` | `/admin/api/pricing_config` | `/api/v1/admin/usage/daily and /api/v1/admin/model-prices` | `admin.py:3717` |
| 已覆盖 | `GET` | `/admin/api/prompts` | `/api/v1/admin/prompts` | `prompts.py:227` |
| 已覆盖 | `POST` | `/admin/api/prompts` | `/api/v1/admin/prompts` | `prompts.py:246` |
| 新契约替代 | `GET` | `/admin/api/prompts/categories` | `/api/v1/admin/prompt-categories` | `prompts.py:135` |
| 新契约替代 | `POST` | `/admin/api/prompts/categories` | `/api/v1/admin/prompt-categories` | `prompts.py:150` |
| 新契约替代 | `DELETE` | `/admin/api/prompts/categories/{cat_id}` | `/api/v1/admin/prompt-categories` | `prompts.py:168` |
| 已覆盖 | `DELETE` | `/admin/api/prompts/{tpl_id}` | `/api/v1/admin/prompts/{prompt_id}` | `prompts.py:323` |
| 已覆盖 | `GET` | `/admin/api/prompts/{tpl_id}` | `/api/v1/admin/prompts/{prompt_id}` | `prompts.py:279` |
| 新契约替代 | `PATCH` | `/admin/api/prompts/{tpl_id}` | `/api/v1/admin/prompts/{prompt_id} (PUT with full revision payload)` | `prompts.py:292` |
| 新契约替代 | `POST` | `/admin/api/prompts/{tpl_id}/preview` | `/api/v1/admin/prompts/{prompt_id}/render` | `prompts.py:392` |
| 已覆盖 | `POST` | `/admin/api/prompts/{tpl_id}/versions` | `/api/v1/admin/prompts/{prompt_id}/versions` | `prompts.py:344` |
| 已覆盖 | `POST` | `/admin/api/prompts/{tpl_id}/versions/{version_id}/activate` | `/api/v1/admin/prompts/{prompt_id}/versions/{version}/activate` | `prompts.py:371` |
| 已覆盖 | `GET` | `/admin/api/providers` | `/api/v1/admin/providers` | `admin.py:1083` |
| 已覆盖 | `POST` | `/admin/api/providers` | `/api/v1/admin/providers` | `admin.py:1106` |
| 新契约替代 | `DELETE` | `/admin/api/providers/{provider_id}` | `/api/v1/admin/providers/{provider_id} enabled=false` | `admin.py:1152` |
| 新契约替代 | `PATCH` | `/admin/api/providers/{provider_id}` | `/api/v1/admin/providers/{provider_id} (PUT with full revision payload)` | `admin.py:1169` |
| 已覆盖 | `PUT` | `/admin/api/providers/{provider_id}` | `/api/v1/admin/providers/{provider_id}` | `admin.py:1128` |
| 新契约替代 | `GET` | `/admin/api/providers/{provider_id}/backends` | `/api/v1/admin/providers/{provider_id}/models and /api/v1/admin/model-routes` | `admin.py:1333` |
| 新契约替代 | `PUT` | `/admin/api/providers/{provider_id}/backends` | `/api/v1/admin/providers/{provider_id}/models and /api/v1/admin/model-routes` | `admin.py:1358` |
| 新契约替代 | `DELETE` | `/admin/api/providers/{provider_id}/backends/{backend_name}` | `/api/v1/admin/providers/{provider_id}/models and /api/v1/admin/model-routes` | `admin.py:1455` |
| 新契约替代 | `PATCH` | `/admin/api/providers/{provider_id}/backends/{backend_name}` | `/api/v1/admin/providers/{provider_id}/models and /api/v1/admin/model-routes` | `admin.py:1414` |
| 新契约替代 | `POST` | `/admin/api/providers/{provider_id}/fetch_models` | `/api/v1/admin/providers/{provider_id}/models/sync` | `admin.py:1255` |
| 已覆盖 | `POST` | `/admin/api/providers/{provider_id}/test` | `/api/v1/admin/providers/{provider_id}/test` | `admin.py:1231` |
| 新契约替代 | `GET` | `/admin/api/rate_limit_rules` | `/api/v1/admin/api-keys and /api/v1/admin/model-routes` | `admin.py:2033` |
| 新契约替代 | `POST` | `/admin/api/rate_limit_rules` | `/api/v1/admin/api-keys and /api/v1/admin/model-routes` | `admin.py:2044` |
| 新契约替代 | `DELETE` | `/admin/api/rate_limit_rules/{rule_id}` | `/api/v1/admin/api-keys and /api/v1/admin/model-routes` | `admin.py:2090` |
| 新契约替代 | `PATCH` | `/admin/api/rate_limit_rules/{rule_id}` | `/api/v1/admin/api-keys and /api/v1/admin/model-routes` | `admin.py:2064` |
| 已覆盖 | `POST` | `/admin/api/reload` | `/api/v1/admin/reload` | `admin.py:3292` |
| 新契约替代 | `GET` | `/admin/api/routing_groups` | `/api/v1/admin/model-routes` | `smart_router.py:309` |
| 新契约替代 | `POST` | `/admin/api/routing_groups` | `/api/v1/admin/model-routes` | `smart_router.py:326` |
| 新契约替代 | `DELETE` | `/admin/api/routing_groups/{group_id}` | `/api/v1/admin/model-routes` | `smart_router.py:372` |
| 新契约替代 | `PATCH` | `/admin/api/routing_groups/{group_id}` | `/api/v1/admin/model-routes` | `smart_router.py:346` |
| 新契约替代 | `POST` | `/admin/api/routing_groups/{group_id}/preview` | `/api/v1/admin/model-routes` | `smart_router.py:468` |
| 新契约替代 | `POST` | `/admin/api/routing_groups/{group_id}/targets` | `/api/v1/admin/model-routes` | `smart_router.py:390` |
| 新契约替代 | `DELETE` | `/admin/api/routing_groups/{group_id}/targets/{target_id}` | `/api/v1/admin/model-routes` | `smart_router.py:448` |
| 新契约替代 | `PATCH` | `/admin/api/routing_groups/{group_id}/targets/{target_id}` | `/api/v1/admin/model-routes` | `smart_router.py:422` |
| 新契约替代 | `GET` | `/admin/api/rule_dimensions` | `/api/v1/admin/content-policies` | `admin.py:1854` |
| 新契约替代 | `GET` | `/admin/api/rules` | `/api/v1/admin/content-policies` | `admin.py:1717` |
| 新契约替代 | `POST` | `/admin/api/rules` | `/api/v1/admin/content-policies` | `admin.py:1753` |
| 新契约替代 | `POST` | `/admin/api/rules/batch` | `/api/v1/admin/content-policies` | `admin.py:1775` |
| 新契约替代 | `DELETE` | `/admin/api/rules/{rule_id}` | `/api/v1/admin/content-policies` | `admin.py:1835` |
| 新契约替代 | `PATCH` | `/admin/api/rules/{rule_id}` | `/api/v1/admin/content-policies` | `admin.py:1808` |
| 新契约替代 | `GET` | `/admin/api/stats` | `/api/v1/admin/usage/daily and /api/v1/admin/model-prices` | `admin.py:3022` |
| 新契约替代 | `GET` | `/admin/api/system_info` | `/api/v1/admin/system-info` | `admin.py:2488` |
| 已覆盖 | `GET` | `/admin/api/tools` | `/api/v1/admin/tools` | `tools.py:218` |
| 已覆盖 | `POST` | `/admin/api/tools` | `/api/v1/admin/tools` | `tools.py:229` |
| 已覆盖 | `DELETE` | `/admin/api/tools/{tool_id}` | `/api/v1/admin/tools/{tool_id}` | `tools.py:305` |
| 新契约替代 | `PATCH` | `/admin/api/tools/{tool_id}` | `/api/v1/admin/tools/{tool_id} (PUT with full revision payload)` | `tools.py:260` |
| 明确退役 | `GET` | `/admin/api/tools/{tool_id}/headers` | `retired: decrypted secret headers are intentionally write-only` | `tools.py:320` |
| 新契约替代 | `POST` | `/admin/api/tools/{tool_id}/test_invoke` | `/api/v1/admin/tools/{tool_id}/test` | `tools.py:338` |
| 新契约替代 | `POST` | `/admin/api/totp/backup_codes/regenerate` | `/api/v1/admin/totp/backup-codes/regenerate` | `admin.py:521` |
| 已覆盖 | `POST` | `/admin/api/totp/confirm` | `/api/v1/admin/totp/confirm` | `admin.py:464` |
| 已覆盖 | `POST` | `/admin/api/totp/disable` | `/api/v1/admin/totp/disable` | `admin.py:487` |
| 已覆盖 | `POST` | `/admin/api/totp/setup` | `/api/v1/admin/totp/setup` | `admin.py:439` |
| 已覆盖 | `GET` | `/admin/api/totp/status` | `/api/v1/admin/totp/status` | `admin.py:422` |
| 新契约替代 | `GET` | `/admin/api/users` | `/api/v1/admin/identities/portal-users` | `admin.py:683` |
| 新契约替代 | `POST` | `/admin/api/users` | `/api/v1/admin/identities/portal-users` | `admin.py:721` |
| 新契约替代 | `DELETE` | `/admin/api/users/{user_id}` | `/api/v1/admin/identities/*/{identity_id} active=false` | `admin.py:826` |
| 新契约替代 | `PATCH` | `/admin/api/users/{user_id}` | `/api/v1/admin/identities/portal-users/{identity_id} (PUT with full revision payload)` | `admin.py:745` |
| 新契约替代 | `POST` | `/admin/api/users/{user_id}/approve` | `/api/v1/admin/identities/portal-users/{identity_id} active=true` | `admin.py:788` |
| 新契约替代 | `DELETE` | `/admin/api/users/{user_id}/feishu` | `/api/v1/admin/identity-providers (generic identity link lifecycle)` | `admin.py:806` |
| 新契约替代 | `POST` | `/admin/api/users/{user_id}/reset_password` | `/api/v1/admin/identities/portal-users/{identity_id} password field` | `admin.py:771` |
| 已覆盖 | `GET` | `/admin/api/whoami` | `/api/v1/admin/whoami` | `admin.py:256` |
| 明确退役 | `GET` | `/admin/ui` | `Art Design Pro SPA deployment` | `main.py:3838` |
| 已覆盖 | `GET` | `/healthz` | `/healthz` | `main.py:3756` |
| 已覆盖 | `GET` | `/metrics` | `/metrics` | `main.py:3793` |
| 已覆盖 | `GET` | `/portal/api/applications` | `/api/v1/portal/applications` | `applications.py:408` |
| 已覆盖 | `POST` | `/portal/api/apps/{app_code}/chat` | `/api/v1/portal/apps/{application_code}/chat` | `orchestration.py:227` |
| 已覆盖 | `POST` | `/portal/api/apps/{app_code}/conversations` | `/api/v1/portal/apps/{application_code}/conversations` | `orchestration.py:245` |
| 已覆盖 | `GET` | `/portal/api/apps/{app_code}/conversations/{conversation_id}` | `/api/v1/portal/apps/{application_code}/conversations/{conversation_id}` | `orchestration.py:260` |
| 已覆盖 | `POST` | `/portal/api/apps/{app_code}/conversations/{conversation_id}/messages` | `/api/v1/portal/apps/{application_code}/conversations/{conversation_id}/messages` | `orchestration.py:275` |
| 已覆盖 | `GET` | `/portal/api/catalog` | `/api/v1/portal/catalog` | `portal.py:938` |
| 已覆盖 | `GET` | `/portal/api/cost` | `/api/v1/portal/cost` | `portal.py:550` |
| 已覆盖 | `GET` | `/portal/api/docs-info` | `/api/v1/portal/docs-info` | `portal.py:690` |
| 新契约替代 | `GET` | `/portal/api/feishu/callback` | `/api/v1/portal/sso (generic OIDC/SAML)` | `portal.py:214` |
| 新契约替代 | `GET` | `/portal/api/feishu/enabled` | `/api/v1/portal/sso (generic OIDC/SAML)` | `portal.py:189` |
| 新契约替代 | `GET` | `/portal/api/feishu/login` | `/api/v1/portal/sso (generic OIDC/SAML)` | `portal.py:195` |
| 已覆盖 | `GET` | `/portal/api/knowledge` | `/api/v1/portal/knowledge` | `portal.py:830` |
| 已覆盖 | `POST` | `/portal/api/login` | `/api/v1/portal/login` | `portal.py:76` |
| 已覆盖 | `POST` | `/portal/api/login/totp` | `/api/v1/portal/login/totp` | `portal.py:123` |
| 已覆盖 | `POST` | `/portal/api/logout` | `/api/v1/portal/logout` | `portal.py:181` |
| 已覆盖 | `GET` | `/portal/api/logs` | `/api/v1/portal/logs` | `portal.py:597` |
| 已覆盖 | `GET` | `/portal/api/logs/{log_id}` | `/api/v1/portal/logs/{audit_event_id}` | `portal.py:657` |
| 已覆盖 | `GET` | `/portal/api/me` | `/api/v1/portal/me` | `portal.py:432` |
| 已覆盖 | `GET` | `/portal/api/model-requests` | `/api/v1/portal/model-requests` | `portal.py:803` |
| 已覆盖 | `POST` | `/portal/api/model-requests` | `/api/v1/portal/model-requests` | `portal.py:754` |
| 已覆盖 | `GET` | `/portal/api/model-requests/available` | `/api/v1/portal/model-requests/available` | `portal.py:716` |
| 已覆盖 | `POST` | `/portal/api/password` | `/api/v1/portal/password` | `portal.py:273` |
| 已覆盖 | `GET` | `/portal/api/prompts` | `/api/v1/portal/prompts` | `portal.py:855` |
| 已覆盖 | `GET` | `/portal/api/prompts/{tpl_id}` | `/api/v1/portal/prompts/{prompt_id}` | `portal.py:887` |
| 已覆盖 | `DELETE` | `/portal/api/prompts/{tpl_id}/favorite` | `/api/v1/portal/prompts/{prompt_id}/favorite` | `portal.py:926` |
| 已覆盖 | `POST` | `/portal/api/prompts/{tpl_id}/favorite` | `/api/v1/portal/prompts/{prompt_id}/favorite` | `portal.py:910` |
| 已覆盖 | `GET` | `/portal/api/stats` | `/api/v1/portal/stats` | `portal.py:457` |
| 已覆盖 | `GET` | `/portal/api/tools` | `/api/v1/portal/tools` | `portal.py:838` |
| 新契约替代 | `POST` | `/portal/api/totp/backup_codes/regenerate` | `/api/v1/portal/totp/backup-codes/regenerate` | `portal.py:409` |
| 已覆盖 | `POST` | `/portal/api/totp/confirm` | `/api/v1/portal/totp/confirm` | `portal.py:354` |
| 已覆盖 | `POST` | `/portal/api/totp/disable` | `/api/v1/portal/totp/disable` | `portal.py:377` |
| 已覆盖 | `POST` | `/portal/api/totp/setup` | `/api/v1/portal/totp/setup` | `portal.py:328` |
| 已覆盖 | `GET` | `/portal/api/totp/status` | `/api/v1/portal/totp/status` | `portal.py:310` |
| 明确退役 | `GET` | `/portal/ui` | `Art Design Pro SPA deployment` | `main.py:3846` |
| 已覆盖 | `GET` | `/readyz` | `/readyz` | `main.py:3762` |
| 新契约替代 | `POST` | `/v1/apps/{app_code}/chat/completions` | `/v1/applications/{application_code}/chat/completions` | `orchestration.py:213` |
| 已覆盖 | `POST` | `/v1/chat/completions` | `/v1/chat/completions` | `main.py:2502` |
| 已覆盖 | `POST` | `/v1/embeddings` | `/v1/embeddings` | `main.py:2541` |
| 已覆盖 | `POST` | `/v1/knowledge/search` | `/v1/knowledge/search` | `main.py:3543` |
| 已覆盖 | `POST` | `/v1/messages` | `/v1/messages` | `main.py:2515` |
| 已覆盖 | `GET` | `/v1/models` | `/v1/models` | `main.py:3512` |
| 已覆盖 | `GET` | `/v1/prompts` | `/v1/prompts` | `main.py:3657` |
| 已覆盖 | `POST` | `/v1/prompts/{prompt_name}/render` | `/v1/prompts/{prompt_name}/render` | `main.py:3680` |
| 已覆盖 | `POST` | `/v1/responses` | `/v1/responses` | `main.py:2530` |
| 已覆盖 | `GET` | `/v1/tools` | `/v1/tools` | `main.py:3601` |
| 已覆盖 | `POST` | `/v1/tools/{tool_name}/invoke` | `/v1/tools/{tool_code}/invoke` | `main.py:3610` |
@@ -0,0 +1,170 @@
# AI Gateway Go — 旗舰版(Ultra)需求规划与实现情况报告
- **报告日期**: 2026-08-12
- **工程**: AI Gateway 全量 Go 重构(替代原 Python/FastAPI 网关)
- **活跃工作树**: `/home/ben/ai-gateway-src/ai-gateway-go-deploy-0.10.0`
- **当前版本**: 0.10.0(Go 1.26,PostgreSQL 17 + 双 Redis,22 个迁移)
---
## 一、项目现状概览
| 项 | 状态 |
|---|---|
| 部署形态 | Docker Compose(项目名 `deploy`),gateway-api `:8080` / admin-web `:8081` / portal-web `:8082` |
| 数据层 | PostgreSQL 17(权威配置 + 审计分区)+ critical/cache 双 Redis;22 个迁移已应用 |
| 迁移 | `000001``000022`(含旗舰版资源市场 `000022_resource_marketplace.sql`) |
| 验证 | `go build ./...``go vet`、全套单测、资源市场集成测试连真实库 **全部通过** |
| 前端 | Art Design Pro 管理端 + 门户端,已构建进镜像并运行 |
| 版本控制 | 无 git(目录无 `.git`),靠目录快照;建议尽快接入版本控制 |
---
## 二、已完成里程碑(M0M7)
### M0 工程基线
Go 1.26 模块、配置校验、结构化日志、优雅退出;pgx 连接池、双 Redis 客户端;带校验和/事务/advisory lock 的独立迁移器;OpenAI 兼容入口、bootstrap key、请求体限制、SSE 透传;healthz/readyz/Prometheus;outbox/API Key/Provider/审计分区首版 schema。
### M1 身份、Provider 与配置面
管理员/门户用户 + PBKDF2 旧密码兼容升级;可撤销会话/账号锁定;admin/portal 登录、身份、菜单、退出;TOTP 两步验证 + 备用码;API Key 全生命周期 + Redis 即时撤销;Provider AES-256-GCM 凭据 + SSRF 校验 + 事务 outbox + 原子快照/路由/能力约束;KEK 版本化轮换;层级部门 CRUD;OIDC(PKCE/RS256/JWKS)与 SAML 2.0(SP-initiated/断言防重放/自动开户)端到端通过。
### M2 流量治理与路由
API Key 级 RPM + UTC 月请求配额 + Token 配额(请求前预留、usage 回写校准),Redis Lua 原子准入,429/Retry-After/X-RateLimit-*;Provider 超时/退避/熔断(幂等安全重试);模型别名条件路由(确定性加权 + 失败关闭 + X-Gateway-Model);Art 管理端限流/路由页面。
### M3 审计、usage、内容策略与成本
有界异步审计批量 COPY + 失败保留 + 优雅停机;按日 usage 聚合;audit/usage 权限隔离查询;outbox-worker(SKIP LOCKED 多实例 + Redis Stream 原子去重 + 死信重试);审计月分区维护(advisory lock 防多实例重复);内容策略(RE2 不可变快照,audit/block/redact,默认脱敏);时间版本模型定价 + 按日成本核算;Art 审计/用量/内容策略/定价页面。
### M4 Prompt、知识、工具、应用编排和通知
Prompt 分类/模板/不可变版本/必填校验;知识库(2 MiB 有界正文、段落感知重叠分块、FTS+中文二元词片检索、Retriever 接口);声明式 HTTP 工具(JSON Schema 校验、KEK 加密头、双层 SSRF、禁止重定向、1 MiB 响应限制);AI 应用草稿 + 不可发布版本 → `/v1/applications/{code}/chat/completions`(仍过鉴权/配额/内容策略/路由/成本/审计);通知 worker(HMAC-SHA256 Webhook、精确/通配事件、失败重试)。
### M5 契约、迁移与切换工程
旧 Python 网关 38 个持久化实体数据字典;确定性 UUIDv5 旧 ID 映射;201 条旧路由全覆盖/替代/退役决策,未决契约缺口为 0;OpenAPI 0.10.0 覆盖全部 Go 路由;legacy 导入工具 + cutover runbook + 路由契约矩阵 + 校验脚本。
### M6 门户工作台与扩展治理
门户资产目录、Prompt 查看/搜索/收藏、个人审计/Token/成本统计、接入说明、模型访问申请与管理员审批;事实核验(复用 Provider 凭据 + 作用域策略 + 事件契约);已发布应用门户单轮入口 + 服务端托管会话(独立 AEAD purpose 加密凭证、串行租约、SHA-256 哈希链、最多 200 条);Compose 同时交付 API/管理端 8081/门户端 8082。
### M7 旗舰版资源市场(本次新增)
三类可发布资产 **MCP 服务器 / Skills / 数字员工** + 共享分类与标签 + 市场安装(工作区绑定 + 权限):
- 数据层:6 张表(`marketplace_categories` / `mcp_servers` / `skills` / `digital_employees` / `digital_employee_runs` / `marketplace_installations`)
- 后端:管理端 CRUD + 在线编辑/发布(`admin_marketplace.go`);统一市场目录合并查询(`marketplace.go`);运行时渲染 skill、列出/调用 MCP 工具、运行数字员工(含知识检索 + 工具执行 + 访问控制,`runtime_marketplace.go`);MCP 客户端(SSE + JSON-RPC 工具发现/调用/缓存,`mcp_client.go`)
- 前端:管理端 4 页(市场总览 / MCP 服务器 / Skills / 数字员工)+ 门户资源市场页
- 测试:`TestMarketplaceLifecycle` / `TestWorkbenchPostgreSQLLifecycle` 连真实 PostgreSQL 通过
---
## 三、旗舰版(Ultra)功能矩阵对照
状态图例:**✅ 已完成** | **⚠️ 部分覆盖** | **❌ 未实现**
### 管理平台
| 功能模块 | 功能 | 社区 Free | 专业 Pro | 旗舰 Ultra | 当前状态 |
|---|---|---|---|---|---|
| 管理平台 | AI 助手:通过 AI 助手查看平台信息和权限控制 | — | ✓ | ✓ | ❌ |
| 管理平台 | 收藏:常用功能收藏 | ✓ | ✓ | ✓ | ❌ |
| 管理平台 | 概览:平台总览与基础看板 | — | ✓ | ✓ | ⚠️ 仅有 24h 监控汇总(`monitoring/overview`),非完整看板 |
| 管理平台 | 数据报表:token/工具/渠道/审批授权/安全事件汇总统计 | — | ✓ | ✓ | ❌(仅有 usage/stats 单维) |
| 资源市场 | 自定义 MCP/Skills/数字员工资源 | ✓ | ✓ | ✓ | ✅ M7 |
| 资源市场 | 在线编辑和发布 MCP/Skills/数字员工 | 不支持 | ✓ | ✓ | ✅ 管理端 CRUD + publish 路由 |
| 资源市场 | 资源分类管理和标签管理 | ✓ | ✓ | ✓ | ✅ `marketplace_categories` |
| 配置管理 | 平台配置敏感参数,skill/mcp 运行时动态注入环境变量 | — | ✓ | ✓ | ❌ |
| 智能体管理 | 智能体节点:运行监控/动态创建/裸机安装/节点池分配/公有私有池路由 | — | ✓ | ✓ | ❌ |
| 智能体管理 | LLMTrace:会话中大模型调用、工具调用执行性能跟踪 | — | ✓ | ✓ | ❌ |
| 智能体管理 | 智能体会话:会话列表,区分普通会话与数字员工会话 | — | ✓ | ✓ | ❌(有应用托管会话,非会话列表体系) |
| 模型管理 | 对接国内外主流大模型供应商 | ✓ | ✓ | ✓ | ✅ providers + 模型目录同步 |
| 模型管理 | 供应商中添加配置大模型 | ✓ | ✓ | ✓ | ✅ |
| 模型管理 | 对接本地模型(Ollama/vLLM) | 不支持 | ✓ | ✓ | ✅ 走 OpenAI 兼容通用 provider |
| 模型管理 | 不同大模型 Token 配额、使用量统计 | 不支持 | ✓ | ✓ | ✅ quota + usage 聚合 |
| 模型管理 | 大模型使用权限分级管控(用户/角色) | ✓ | ✓ | ✓ | ✅ RBAC scope |
| 知识库 | 外部文档导入、自动切分、向量化、语义匹配召回 | — | ✓ | ✓ | ⚠️ 导入/切分/检索 ✅;向量化+语义匹配 ❌(现为 FTS+中文分词) |
| 记忆管理 | 记忆集合:个人/部门/全局多层记忆,提炼与语义匹配召回 | — | ✓ | ✓ | ❌ |
| 记忆管理 | 根据调用自动裁剪衰减片段 | — | ✓ | ✓ | ❌ |
| 记忆管理 | 记忆授权:提炼内容/沉淀经验授权给其他用户 | — | ✓ | ✓ | ❌ |
| 权限管理 | 完整 RBAC 角色权限管理 | 不支持 | ✓ | ✓ | ✅ |
| 权限管理 | 用户/部门/角色管理 | ✓ | ✓ | ✓ | ✅ |
| 权限管理 | 系统历史操作审计日志 | — | ✓ | ✓ | ✅ `admin/audit-events` |
| 权限管理 | 资源(mcp/skills/数字员工)权限管控、资源授权 | ✓ | ✓ | ✓ | ✅ marketplace 访问控制 |
| 权限管理 | 资源/大模型/渠道使用申请审批 | — | ✓ | ✓ | ⚠️ 仅模型申请审批(model-requests);资源安装/渠道无审批流 |
| API 集成 | API Key 调用大模型/mcp/skills 组合或数字员工 | — | ✓ | ✓ | ✅ API Key + runtime marketplace |
| 渠道管理 | Web 聊天界面 | ✓ | ✓ | ✓ | ✅ portal |
| 渠道管理 | 企业微信/个人微信/钉钉/飞书渠道 | ✓ | ✓ | ✓ | ❌ |
| 渠道管理 | 渠道权限管控、使用权限授权(用户/角色) | ✓ | ✓ | ✓ | ❌ |
| 渠道管理 | 多渠道治理、审计、使用量统计 | ✓ | ✓ | ✓ | ❌ |
| 安装部署 | x86_64 安装包 | ✓ | ✓ | ✓ | ⚠️ Docker Compose 交付,非传统安装包 |
| 安装部署 | ARM64 安装包 | — | ✓ | ✓ | ❌ |
| 部署方式 | 单机 / 冷备 / 集群 | 单机 | 单机/冷备 | 单机/冷备/集群 | ⚠️ outbox 支持多实例 SKIP LOCKED(集群一部分);冷备/完整集群方案未做 |
| 租户管理 | 单租户使用 | ✓ | ✓ | ✓ | ✅ |
| 租户管理 | 平台管理员多租户管理 | — | — | ✓ | ❌ |
| 安全策略 | 运行时安全:网络/工具命令执行安全校验审批/工具调用频率限制 | — | ✓ | ✓ | ⚠️ 工具 SSRF/拨号防护 ✅;命令执行审批、工具限流 ❌ |
| 安全策略 | 供应链安全:skill/mcp 资源安全扫描 | — | ✓ | ✓ | ❌ |
| 安全策略 | 数据安全:工具数据输入输出脱敏 + 大模型回答隐私敏感信息拦截替换 | — | ✓ | ✓ | ⚠️ 提示词输入脱敏 ✅;工具输出/回答拦截替换 ❌ |
| 站内消息 | 平台推送站内消息与动态 | — | ✓ | ✓ | ❌(通知 worker 仅 webhook) |
| 审批授权 | 资源/模型/渠道使用申请流程审批管理 | — | ✓ | ✓ | ⚠️ 仅模型申请 |
| 文件管理 | 平台文件资源与对象存储文件浏览管理 | — | ✓ | ✓ | ❌(MinIO/S3 不在基线) |
| 审计日志 | 系统全量历史操作审计日志查询 | — | ✓ | ✓ | ✅ |
| 企业报表 | 企业级运营数据报表统计分析 | — | ✓ | ✓ | ❌ |
| License | 平台 License 授权管理与有效期管控 | ✓ | ✓ | ✓ | ❌ |
### 工作台
| 功能 | 功能说明 | 当前状态 |
|---|---|---|
| 聊天 | 新建会话、授权大模型对话 | ✅ |
| 聊天 | 删除会话、会话重命名 | ✅ |
| 定时任务 | 创建定时任务(配置提示词/渠道/mcp/skills/数字员工/会话ID) | ❌ |
| 定时任务 | 启动/修改/立即执行/删除定时任务配置 | ❌ |
| 定时任务 | 查看定时任务执行历史 | ❌ |
| 个人渠道 | 配置个人微信/企业微信/钉钉/飞书 | ❌ |
| 个人渠道 | 个人微信/企业微信快速扫码对接 | ❌ |
| 个人渠道 | 绑定特定大模型执行对话 | ❌ |
| 我的资源 | 查看被授权资源(mcp/skills/数字员工)详细信息 | ✅ portal/marketplace |
| 我的资源 | 从被授权数字员工进入会话 | ⚠️ runtime 可跑,无会话列表入口 |
| 我的资源 | 通过权限申请从插件市场安装 MCP/Skills/数字员工 | ✅ marketplace install |
| 我的资源 | 查看被授权大模型使用量 | ✅ portal/stats |
| 我的资源 | 查看插件资源权限等级(可查看/仅使用/管理) | ⚠️ 有访问控制,三档等级未成体系 |
| 我的资源 | 申请大模型/token量/skill/mcp/数字员工/渠道权限 | ⚠️ 仅模型申请 |
| 配置管理 | 个人配置环境变量,供 skill/mcp 使用 | ❌ |
| 个人文件仓库 | 对话产生的报告/文件存入个人仓库 | ❌ |
| 安全策略 | 个人智能体安全策略(网络/工具命令校验审批/限流/脱敏/隐私拦截) | ❌ |
| 个人中心 | 账号信息、密码修改、登录记录查看 | ✅ |
| 消息通知 | 系统消息、审批待办、任务执行结果提醒 | ⚠️ webhook 投递 ✅;站内消息/待办/结果提醒 ❌ |
---
## 四、差距汇总
- **完全未实现(❌,约 20 项)**:AI 助手、收藏、数据报表/企业报表、配置管理(env 注入)、智能体管理三项(节点/LLMTrace/会话)、记忆管理三项、渠道管理全项、多租户、供应链安全扫描、站内消息、完整审批流、文件管理(对象存储)、License、定时任务全项、个人渠道、个人文件仓库、个人安全策略、ARM64。
- **部分覆盖需补齐(⚠️,约 10 项)**:平台概览看板、知识库向量化语义召回、资源/渠道审批、工具输出脱敏与大模型回答拦截替换、工具命令审批与工具限流、集群部署方案、站内消息/审批待办/任务结果、数字员工会话入口、资源权限等级三档、全类型权限申请。
---
## 五、后续里程碑规划(M8–M13)
| 里程碑 | 内容 | 依赖 |
|---|---|---|
| **M8 基础设施层** | 对象存储(MinIO)、向量化(pgvector)、定时任务调度器、站内消息 | — |
| **M9 智能体与可观测** | LLMTrace、智能体会话、智能体节点(监控/节点池/路由)、AI 助手 | M8 |
| **M10 记忆管理** | 多层记忆集合、语义召回、裁剪衰减、记忆授权 | M8(pgvector) |
| **M11 渠道与审批** | 渠道管理(企业微信/个人微信/钉钉/飞书)、个人渠道、完整审批流、资源权限等级 | M8 |
| **M12 数据安全与供应链** | 知识库向量化、输出脱敏/回答拦截、工具命令审批/限流、供应链扫描、个人安全策略 | M8(pgvector) |
| **M13 平台运营** | 数据报表/企业报表、完整看板、收藏、多租户、License、ARM64/集群部署 | M8 |
已建任务跟踪:`#4``#9`
---
## 六、验证情况(当前基线)
- `go build ./...` ✅、`go vet ./internal/workbench/`
- 全套单测(全部包)✅;资源市场单测(含 MCP 客户端)✅
- 集成测试 `TestMarketplaceLifecycle``TestWorkbenchPostgreSQLLifecycle` 连真实 PostgreSQL ✅
- 部署冒烟:healthz/readyz ✅、admin :8081 / portal :8082 302 ✅、22 迁移应用 ✅
## 七、部署与已知坑
- 活跃工作树:`/home/ben/ai-gateway-src/ai-gateway-go-deploy-0.10.0`;compose 项目名 `deploy`,workdir 在 `deploy/` 下。
- **重启恢复**:postgres/redis restart 策略为 `no`,重启后需进 `deploy/` 执行 `docker compose up -d`;若容器 `networks` 为空需 `--force-recreate`
- **旧 Python 网关**:`llm-gateway.service`(systemd)监听 8080 已 `systemctl disable`,不再开机抢占。
- **本机无 go 工具链**:编译/测试用 `docker run --rm -v $PWD:/src -w /src -e GOCACHE=/tmp/gocache golang:1.26.5-alpine sh -c 'go build ./...'`;集成测试加 `--network deploy_default` + `WORKBENCH_TEST_DATABASE_URL=postgres://gateway:gateway@postgres:5432/gateway?sslmode=disable`
- **上线门禁未过**:旧库脱敏快照迁移、影子观察、容量验收、切换/回退演练需真实生产数据与授权后方可执行。
+27
View File
@@ -0,0 +1,27 @@
module aigateway.local/core
go 1.26.0
toolchain go1.26.5
require (
github.com/beevik/etree v1.5.0
github.com/crewjam/saml v0.5.1
github.com/jackc/pgx/v5 v5.10.0
github.com/redis/go-redis/v9 v9.21.0
github.com/russellhaering/goxmldsig v1.4.0
)
require (
github.com/cespare/xxhash/v2 v2.3.0 // indirect
github.com/golang-jwt/jwt/v4 v4.5.2 // indirect
github.com/jackc/pgpassfile v1.0.0 // indirect
github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect
github.com/jackc/puddle/v2 v2.2.2 // indirect
github.com/jonboulle/clockwork v0.2.2 // indirect
github.com/mattermost/xml-roundtrip-validator v0.1.0 // indirect
go.uber.org/atomic v1.11.0 // indirect
golang.org/x/crypto v0.33.0 // indirect
golang.org/x/sync v0.17.0 // indirect
golang.org/x/text v0.29.0 // indirect
)
+78
View File
@@ -0,0 +1,78 @@
github.com/beevik/etree v1.1.0/go.mod h1:r8Aw8JqVegEf0w2fDnATrX9VpkMcyFeM0FhwO62wh+A=
github.com/beevik/etree v1.5.0 h1:iaQZFSDS+3kYZiGoc9uKeOkUY3nYMXOKLl6KIJxiJWs=
github.com/beevik/etree v1.5.0/go.mod h1:gPNJNaBGVZ9AwsidazFZyygnd+0pAU38N4D+WemwKNs=
github.com/bsm/ginkgo/v2 v2.12.0 h1:Ny8MWAHyOepLGlLKYmXG4IEkioBysk6GpaRTLC8zwWs=
github.com/bsm/ginkgo/v2 v2.12.0/go.mod h1:SwYbGRRDovPVboqFv0tPTcG1sN61LM1Z4ARdbAV9g4c=
github.com/bsm/gomega v1.27.10 h1:yeMWxP2pV2fG3FgAODIY8EiRE3dy0aeFYt4l7wh6yKA=
github.com/bsm/gomega v1.27.10/go.mod h1:JyEr/xRbxbtgWNi8tIEVPUYZ5Dzef52k01W3YH0H+O0=
github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs=
github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
github.com/crewjam/saml v0.5.1 h1:g+mfp0CrLuLRZCK793PgJcZeg5dS/0CDwoeAX2zcwNI=
github.com/crewjam/saml v0.5.1/go.mod h1:r0fDkmFe5URDgPrmtH0IYokva6fac3AUdstiPhyEolQ=
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/golang-jwt/jwt/v4 v4.5.2 h1:YtQM7lnr8iZ+j5q71MGKkNw9Mn7AjHM68uc9g5fXeUI=
github.com/golang-jwt/jwt/v4 v4.5.2/go.mod h1:m21LjoU+eqJr34lmDMbreY2eSTRJ1cv77w39/MY0Ch0=
github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=
github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM=
github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg=
github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo=
github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM=
github.com/jackc/pgx/v5 v5.10.0 h1:VhSvgU2jSli8o3AqIEOTJr7rZwAEUVo4E4XhR94Zfr0=
github.com/jackc/pgx/v5 v5.10.0/go.mod h1:mal1tBGAFfLHvZzaYh77YS/eC6IX9OWbRV1QIIM0Jn4=
github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo=
github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4=
github.com/jonboulle/clockwork v0.2.2 h1:UOGuzwb1PwsrDAObMuhUnj0p5ULPj8V/xJ7Kx9qUBdQ=
github.com/jonboulle/clockwork v0.2.2/go.mod h1:Pkfl5aHPm1nk2H9h0bjmnJD/BcgbGXUBGnn1kMkgxc8=
github.com/klauspost/cpuid/v2 v2.2.10 h1:tBs3QSyvjDyFTq3uoc/9xFpCuOsJQFNPiAhYdw2skhE=
github.com/klauspost/cpuid/v2 v2.2.10/go.mod h1:hqwkgyIinND0mEev00jJYCxPNVRVXFQeu1XKlok6oO0=
github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo=
github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI=
github.com/kr/pretty v0.3.0/go.mod h1:640gp4NfQd8pI5XOwp5fnNeVWj67G7CFk/SaSQn7NBk=
github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
github.com/mattermost/xml-roundtrip-validator v0.1.0 h1:RXbVD2UAl7A7nOTR4u7E3ILa4IbtvKBHw64LDsmu9hU=
github.com/mattermost/xml-roundtrip-validator v0.1.0/go.mod h1:qccnGMcpgwcNaBnxqpJpWWUiPNr5H3O8eDgGV9gT5To=
github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA=
github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
github.com/redis/go-redis/v9 v9.21.0 h1:FPBE4hhbAke+TLmcY3WkpbDffJEomdqPn3HYiqAtL9E=
github.com/redis/go-redis/v9 v9.21.0/go.mod h1:v/M13XI1PVCDcm01VtPFOADfZtHf8YW3baQf57KlIkA=
github.com/rogpeppe/go-internal v1.6.1/go.mod h1:xXDCJY+GAPziupqXw64V24skbSoqbTEfhy4qGm1nDQc=
github.com/rogpeppe/go-internal v1.8.0/go.mod h1:WmiCO8CzOY8rg0OYDC4/i/2WRWAB6poM+XZ2dLUbcbE=
github.com/russellhaering/goxmldsig v1.4.0 h1:8UcDh/xGyQiyrW+Fq5t8f+l2DLB1+zlhYzkPUJ7Qhys=
github.com/russellhaering/goxmldsig v1.4.0/go.mod h1:gM4MDENBQf7M+V824SGfyIUVFWydB7n0KkEubVJl+Tw=
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U=
github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U=
github.com/zeebo/xxh3 v1.1.0 h1:s7DLGDK45Dyfg7++yxI0khrfwq9661w9EN78eP/UZVs=
github.com/zeebo/xxh3 v1.1.0/go.mod h1:IisAie1LELR4xhVinxWS5+zf1lA4p0MW4T+w+W07F5s=
go.uber.org/atomic v1.11.0 h1:ZvwS0R+56ePWxUNi+Atn9dWONBPp/AUETXlHW0DxSjE=
go.uber.org/atomic v1.11.0/go.mod h1:LUxbIzbOniOlMKjJjyPfpl4v+PKK2cNJn91OQbhoJI0=
golang.org/x/crypto v0.33.0 h1:IOBPskki6Lysi0lo9qQvbxiQ+FvsCC/YWOecCHAixus=
golang.org/x/crypto v0.33.0/go.mod h1:bVdXmD7IV/4GdElGPozy6U7lWdRXA4qyRVGJV57uQ5M=
golang.org/x/sync v0.17.0 h1:l60nONMj9l5drqw6jlhIELNv9I0A4OFgRsG9k2oT9Ug=
golang.org/x/sync v0.17.0/go.mod h1:9KTHXmSnoGruLpwFjVSX0lNNA75CykiMECbovNTZqGI=
golang.org/x/sys v0.30.0 h1:QjkSwP/36a20jFYWkSue1YwXzLmsV5Gfq7Eiy72C1uc=
golang.org/x/sys v0.30.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
golang.org/x/text v0.29.0 h1:1neNs90w9YzJ9BocxfsQNHKuAT4pkghyXc4nhZ6sJvk=
golang.org/x/text v0.29.0/go.mod h1:7MhJOA9CD2qZyOKYazxdYMF85OwPdEr9jTtBpO7ydH4=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=
gopkg.in/errgo.v2 v2.1.0/go.mod h1:hNsd1EY+bozCKY1Ytp96fpM3vjJbqLJn88ws8XvfDNI=
gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
gotest.tools v2.2.0+incompatible h1:VsBPFP1AI068pPrMxtb/S8Zkgf9xEmTLJjfM+P5UIEo=
gotest.tools v2.2.0+incompatible/go.mod h1:DsYFclhRJ6vuDpmuTbkuFWG+y2sxOXAzmJt81HFBacw=
+203
View File
@@ -0,0 +1,203 @@
package apikey
import (
"encoding/json"
"errors"
"net/http"
"strings"
"time"
"aigateway.local/core/internal/identity"
"aigateway.local/core/internal/platform/apiresponse"
)
type AdminHTTPHandler struct {
repository *Repository
authenticator *Authenticator
usage *UsageStore
identity *identity.Service
mux *http.ServeMux
}
type createRequest struct {
Name string `json:"name"`
Scopes []string `json:"scopes"`
RequestsPerMinute int `json:"requests_per_minute"`
MonthlyRequestQuota int64 `json:"monthly_request_quota"`
MonthlyTokenQuota int64 `json:"monthly_token_quota"`
ExpiresAt *time.Time `json:"expires_at"`
}
type limitsRequest struct {
RequestsPerMinute int `json:"requests_per_minute"`
MonthlyRequestQuota int64 `json:"monthly_request_quota"`
MonthlyTokenQuota int64 `json:"monthly_token_quota"`
}
func NewAdminHTTPHandler(repository *Repository, authenticator *Authenticator, identityService *identity.Service) *AdminHTTPHandler {
h := &AdminHTTPHandler{repository: repository, authenticator: authenticator, identity: identityService, mux: http.NewServeMux()}
h.mux.HandleFunc("GET /api/v1/admin/api-keys", h.list)
h.mux.HandleFunc("POST /api/v1/admin/api-keys", h.create)
h.mux.HandleFunc("PUT /api/v1/admin/api-keys/{api_key_id}/limits", h.updateLimits)
h.mux.HandleFunc("DELETE /api/v1/admin/api-keys/{api_key_id}", h.revoke)
return h
}
func (h *AdminHTTPHandler) SetUsageStore(store *UsageStore) { h.usage = store }
func (h *AdminHTTPHandler) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
h.mux.ServeHTTP(writer, request)
}
func (h *AdminHTTPHandler) list(writer http.ResponseWriter, request *http.Request) {
if _, ok := h.requirePermission(writer, request, identity.PermissionAPIKeyRead); !ok {
return
}
records, err := h.repository.List(request.Context())
if err != nil {
h.writeError(writer, err)
return
}
ids := make([]string, 0, len(records))
for _, record := range records {
ids = append(ids, record.ID)
}
usage, err := h.usage.MonthlyTokens(request.Context(), ids, time.Now())
if err != nil {
h.writeError(writer, err)
return
}
items := make([]map[string]any, 0, len(records))
for _, record := range records {
item := publicRecord(record)
item["monthly_token_usage"] = usage[record.ID]
items = append(items, item)
}
apiresponse.OK(writer, items)
}
func (h *AdminHTTPHandler) updateLimits(writer http.ResponseWriter, request *http.Request) {
account, ok := h.requirePermission(writer, request, identity.PermissionAPIKeyManage)
if !ok {
return
}
var input limitsRequest
decoder := json.NewDecoder(http.MaxBytesReader(writer, request.Body, 1<<20))
decoder.DisallowUnknownFields()
if decoder.Decode(&input) != nil || !validLimits(input.RequestsPerMinute, input.MonthlyRequestQuota, input.MonthlyTokenQuota) {
apiresponse.Error(writer, http.StatusBadRequest, "API Key 限流或月配额无效")
return
}
record, hash, err := h.repository.UpdateLimits(request.Context(), request.PathValue("api_key_id"), input.RequestsPerMinute, input.MonthlyRequestQuota, input.MonthlyTokenQuota, account.ID)
if err != nil {
h.writeError(writer, err)
return
}
if err := h.authenticator.Invalidate(request.Context(), hash); err != nil {
h.writeError(writer, err)
return
}
response := publicRecord(record)
if h.usage != nil {
usage, usageErr := h.usage.MonthlyTokens(request.Context(), []string{record.ID}, time.Now())
if usageErr == nil {
response["monthly_token_usage"] = usage[record.ID]
}
}
apiresponse.OK(writer, response)
}
func validLimits(requestsPerMinute int, monthlyRequestQuota, monthlyTokenQuota int64) bool {
return requestsPerMinute >= 0 && requestsPerMinute <= 1_000_000 &&
monthlyRequestQuota >= 0 && monthlyRequestQuota <= 1_000_000_000_000 &&
monthlyTokenQuota >= 0 && monthlyTokenQuota <= 1_000_000_000_000_000
}
func (h *AdminHTTPHandler) create(writer http.ResponseWriter, request *http.Request) {
account, ok := h.requirePermission(writer, request, identity.PermissionAPIKeyManage)
if !ok {
return
}
var input createRequest
decoder := json.NewDecoder(http.MaxBytesReader(writer, request.Body, 1<<20))
decoder.DisallowUnknownFields()
if decoder.Decode(&input) != nil {
apiresponse.Error(writer, http.StatusBadRequest, "请求格式无效")
return
}
input.Name = strings.TrimSpace(input.Name)
if input.Name == "" || len(input.Name) > 128 || len(input.Scopes) == 0 || input.ExpiresAt != nil && !input.ExpiresAt.After(time.Now()) {
apiresponse.Error(writer, http.StatusBadRequest, "名称、权限范围或过期时间无效")
return
}
if !validLimits(input.RequestsPerMinute, input.MonthlyRequestQuota, input.MonthlyTokenQuota) {
apiresponse.Error(writer, http.StatusBadRequest, "API Key 限流或月配额无效")
return
}
for _, scope := range input.Scopes {
if scope != "gateway:invoke" && scope != "*" {
apiresponse.Error(writer, http.StatusBadRequest, "包含不支持的权限范围")
return
}
}
record, secret, err := h.repository.Create(request.Context(), input.Name, input.Scopes, input.RequestsPerMinute, input.MonthlyRequestQuota, input.MonthlyTokenQuota, input.ExpiresAt, account.ID)
if err != nil {
h.writeError(writer, err)
return
}
response := publicRecord(record)
response["key"] = secret
response["warning"] = "密钥只显示一次,请立即安全保存"
apiresponse.OK(writer, response)
}
func (h *AdminHTTPHandler) revoke(writer http.ResponseWriter, request *http.Request) {
account, ok := h.requirePermission(writer, request, identity.PermissionAPIKeyManage)
if !ok {
return
}
hash, err := h.repository.Revoke(request.Context(), request.PathValue("api_key_id"), account.ID)
if err != nil {
h.writeError(writer, err)
return
}
if err := h.authenticator.Invalidate(request.Context(), hash); err != nil {
h.writeError(writer, err)
return
}
apiresponse.OK(writer, map[string]bool{"revoked": true})
}
func (h *AdminHTTPHandler) requirePermission(writer http.ResponseWriter, request *http.Request, permission string) (identity.Account, bool) {
account, err := h.identity.Authenticate(request.Context(), identity.KindAdmin, request.Header.Get("Authorization"))
if err != nil {
apiresponse.Error(writer, http.StatusUnauthorized, "登录状态无效或已过期")
return identity.Account{}, false
}
if !identity.HasPermission(account, permission) {
apiresponse.Error(writer, http.StatusForbidden, "缺少 API Key 操作权限")
return identity.Account{}, false
}
return account, true
}
func (h *AdminHTTPHandler) writeError(writer http.ResponseWriter, err error) {
switch {
case errors.Is(err, ErrInvalid):
apiresponse.Error(writer, http.StatusNotFound, "API Key 不存在或已撤销")
case errors.Is(err, ErrStore):
apiresponse.Error(writer, http.StatusServiceUnavailable, "API Key 服务暂不可用")
default:
apiresponse.Error(writer, http.StatusInternalServerError, "API Key 处理失败")
}
}
func publicRecord(record Record) map[string]any {
return map[string]any{
"id": record.ID, "name": record.Name, "key_prefix": record.KeyPrefix,
"scopes": record.Scopes, "enabled": record.Enabled, "expires_at": record.ExpiresAt,
"requests_per_minute": record.RequestsPerMinute, "monthly_request_quota": record.MonthlyRequestQuota,
"monthly_token_quota": record.MonthlyTokenQuota,
"last_used_at": record.LastUsedAt, "created_at": record.CreatedAt,
}
}
+75
View File
@@ -0,0 +1,75 @@
package apikey
import (
"context"
"crypto/rand"
"crypto/sha256"
"encoding/base64"
"encoding/hex"
"errors"
"strings"
"time"
)
var (
ErrInvalid = errors.New("invalid API key")
ErrStore = errors.New("API key store unavailable")
)
type Record struct {
ID string
TenantID *string
Name string
KeyPrefix string
KeyHash []byte
Scopes []string
Enabled bool
RequestsPerMinute int
MonthlyRequestQuota int64
MonthlyTokenQuota int64
ExpiresAt *time.Time
LastUsedAt *time.Time
CreatedAt time.Time
}
type Principal struct {
APIKeyID string `json:"api_key_id"`
TenantID *string `json:"tenant_id,omitempty"`
Scopes []string `json:"scopes"`
RequestsPerMinute int `json:"requests_per_minute"`
MonthlyRequestQuota int64 `json:"monthly_request_quota"`
MonthlyTokenQuota int64 `json:"monthly_token_quota"`
}
func Generate() (secret, prefix string, hash []byte, err error) {
random := make([]byte, 32)
if _, err = rand.Read(random); err != nil {
return "", "", nil, err
}
secret = "gw_" + base64.RawURLEncoding.EncodeToString(random)
prefix = secret[:16]
digest := sha256.Sum256([]byte(secret))
return secret, prefix, digest[:], nil
}
func Digest(secret string) ([]byte, string) {
digest := sha256.Sum256([]byte(strings.TrimSpace(secret)))
return digest[:], hex.EncodeToString(digest[:])
}
func HasScope(scopes []string, required string) bool {
for _, scope := range scopes {
if scope == "*" || scope == required {
return true
}
}
return false
}
type KeyAuthenticator interface {
Authenticate(context.Context, string) error
}
type PrincipalAuthenticator interface {
AuthenticatePrincipal(context.Context, string) (Principal, error)
}
+73
View File
@@ -0,0 +1,73 @@
package apikey
import (
"bytes"
"context"
"testing"
)
func TestBootstrapCompatibilityUsageCounter(t *testing.T) {
authenticator := NewAuthenticator(nil, nil, "temporary-bootstrap-key")
if err := authenticator.Authenticate(context.Background(), "temporary-bootstrap-key"); err != nil {
t.Fatal(err)
}
if authenticator.BootstrapUses() != 1 {
t.Fatalf("unexpected bootstrap usage count %d", authenticator.BootstrapUses())
}
}
func TestGenerateAndDigest(t *testing.T) {
first, prefix, hash, err := Generate()
if err != nil {
t.Fatal(err)
}
if len(first) < 40 || len(prefix) != 16 || prefix != first[:16] {
t.Fatal("invalid API key format")
}
computed, _ := Digest(first)
if !bytes.Equal(hash, computed) {
t.Fatal("stored digest differs")
}
second, _, _, err := Generate()
if err != nil || second == first {
t.Fatal("API keys must be independently random")
}
}
func TestScopes(t *testing.T) {
if !HasScope([]string{"gateway:invoke"}, "gateway:invoke") || !HasScope([]string{"*"}, "gateway:invoke") {
t.Fatal("expected scope is missing")
}
if HasScope([]string{"gateway:read"}, "gateway:invoke") {
t.Fatal("unexpected scope accepted")
}
}
// TestInvokeScope pins the regression where portal application runtime
// credentials ("application:run") could never reach the gateway: the scope
// check only admitted "gateway:invoke", so every hosted application chat
// returned 401 even though the credential is legitimate.
func TestInvokeScope(t *testing.T) {
accept := map[string][]string{
"gateway key": {"gateway:invoke"},
"admin wildcard": {"*"},
"application runtime key": {"application:run"},
"runtime + read": {"application:run", "gateway:read"},
}
for name, scopes := range accept {
if !invokeScope(scopes) {
t.Fatalf("invokeScope(%v) = false, want true for %s", scopes, name)
}
}
reject := map[string][]string{
"read-only": {"gateway:read"},
"unrelated scope": {"workbench:run"},
"empty": {},
"runtime scope missing": {"gateway:read", "workbench:run"},
}
for name, scopes := range reject {
if invokeScope(scopes) {
t.Fatalf("invokeScope(%v) = true, want false for %s", scopes, name)
}
}
}
+124
View File
@@ -0,0 +1,124 @@
package apikey
import (
"context"
"crypto/subtle"
"encoding/json"
"errors"
"fmt"
"log/slog"
"strings"
"sync/atomic"
"time"
"github.com/redis/go-redis/v9"
)
type Authenticator struct {
repository *Repository
redis *redis.Client
bootstrap string
cacheTTL time.Duration
bootstrapUses atomic.Uint64
logger *slog.Logger
}
func NewAuthenticator(repository *Repository, client *redis.Client, bootstrap string) *Authenticator {
return &Authenticator{repository: repository, redis: client, bootstrap: bootstrap, cacheTTL: 30 * time.Second}
}
// SetLogger wires an optional logger used for best-effort cache diagnostics.
func (a *Authenticator) SetLogger(logger *slog.Logger) { a.logger = logger }
func (a *Authenticator) Authenticate(ctx context.Context, secret string) error {
_, err := a.AuthenticatePrincipal(ctx, secret)
return err
}
// invokeScope reports whether the key carries a scope that may reach the
// gateway. Besides regular gateway keys it admits the server-side
// application runtime credential ("application:run"), which is created only
// by the portal runtime (encrypted in PostgreSQL, tenant-bound, never shown
// to browsers) so that hosted application conversations can call the gateway
// on behalf of the app owner. Both call paths in the workbench runtime and
// the main proxy use this single method, so the runtime credential works
// end-to-end without widening any other surface.
func invokeScope(scopes []string) bool {
return HasScope(scopes, "gateway:invoke") || HasScope(scopes, "application:run")
}
func (a *Authenticator) AuthenticatePrincipal(ctx context.Context, secret string) (Principal, error) {
secret = strings.TrimSpace(secret)
if secret == "" {
return Principal{}, ErrInvalid
}
if a.bootstrap != "" && len(secret) == len(a.bootstrap) && subtle.ConstantTimeCompare([]byte(secret), []byte(a.bootstrap)) == 1 {
a.bootstrapUses.Add(1)
// The bootstrap credential is a real gateway key, not an anonymous
// pass. Return a stable, well-known identity so the admission and
// token-quota paths run (they allow it outright: RPM/quota are 0 by
// design for a migration key) and audit records attribute usage to
// "bootstrap" instead of an empty principal that silently skips every
// policy stage. An empty APIKeyID previously bypassed rate limiting,
// quota and audit attribution entirely.
return Principal{APIKeyID: "bootstrap", Scopes: []string{"gateway:invoke"}}, nil
}
hash, hexHash := Digest(secret)
if a.redis != nil {
payload, err := a.redis.Get(ctx, cacheKey(hexHash)).Bytes()
if err == nil {
var principal Principal
if json.Unmarshal(payload, &principal) == nil && invokeScope(principal.Scopes) {
return principal, nil
}
return Principal{}, ErrInvalid
}
if err != nil && !errors.Is(err, redis.Nil) {
return Principal{}, fmt.Errorf("%w: %v", ErrStore, err)
}
}
principal, err := a.repository.Validate(ctx, hash)
if err != nil {
return Principal{}, err
}
if !invokeScope(principal.Scopes) {
return Principal{}, ErrInvalid
}
// Best-effort cache: the database is the source of truth for validation.
// A cache write failure must not turn a key the database just accepted
// into a 503, or a transient Redis blip would take the whole gateway down.
if a.redis != nil {
payload, _ := json.Marshal(principal)
if err := a.redis.Set(ctx, cacheKey(hexHash), payload, a.cacheTTL).Err(); err != nil && a.logger != nil {
a.logger.Warn("api key cache write failed; continuing with database result", "error", err)
}
}
return principal, nil
}
func (a *Authenticator) BootstrapUses() uint64 {
if a == nil {
return 0
}
return a.bootstrapUses.Load()
}
func (a *Authenticator) Invalidate(ctx context.Context, hash []byte) error {
if a.redis == nil {
return nil
}
_, hexHash := DigestFromHash(hash)
return a.redis.Del(ctx, cacheKey(hexHash)).Err()
}
func DigestFromHash(hash []byte) ([]byte, string) {
const hex = "0123456789abcdef"
encoded := make([]byte, len(hash)*2)
for i, value := range hash {
encoded[i*2] = hex[value>>4]
encoded[i*2+1] = hex[value&15]
}
return hash, string(encoded)
}
func cacheKey(hexHash string) string { return "gateway:api-key:v1:" + hexHash }
+179
View File
@@ -0,0 +1,179 @@
package apikey
import (
"context"
"encoding/json"
"errors"
"fmt"
"time"
platformid "aigateway.local/core/internal/platform/id"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgxpool"
)
type Repository struct{ pool *pgxpool.Pool }
func NewRepository(pool *pgxpool.Pool) *Repository { return &Repository{pool: pool} }
func (r *Repository) List(ctx context.Context) ([]Record, error) {
if r.pool == nil {
return nil, ErrStore
}
rows, err := r.pool.Query(ctx, `
SELECT id::text, tenant_id::text, name, key_prefix, scopes, enabled,
requests_per_minute, monthly_request_quota, monthly_token_quota, expires_at, last_used_at, created_at
FROM gateway.api_keys ORDER BY created_at DESC`)
if err != nil {
return nil, fmt.Errorf("%w: %v", ErrStore, err)
}
defer rows.Close()
var records []Record
for rows.Next() {
var record Record
if err := rows.Scan(&record.ID, &record.TenantID, &record.Name, &record.KeyPrefix, &record.Scopes, &record.Enabled, &record.RequestsPerMinute, &record.MonthlyRequestQuota, &record.MonthlyTokenQuota, &record.ExpiresAt, &record.LastUsedAt, &record.CreatedAt); err != nil {
return nil, fmt.Errorf("%w: %v", ErrStore, err)
}
records = append(records, record)
}
return records, rows.Err()
}
func (r *Repository) Create(ctx context.Context, name string, scopes []string, requestsPerMinute int, monthlyRequestQuota, monthlyTokenQuota int64, expiresAt *time.Time, actorID string) (Record, string, error) {
if r.pool == nil {
return Record{}, "", ErrStore
}
secret, prefix, hash, err := Generate()
if err != nil {
return Record{}, "", err
}
id, err := platformid.NewUUID()
if err != nil {
return Record{}, "", err
}
eventID, err := platformid.NewUUID()
if err != nil {
return Record{}, "", err
}
tx, err := r.pool.Begin(ctx)
if err != nil {
return Record{}, "", fmt.Errorf("%w: %v", ErrStore, err)
}
defer func() { _ = tx.Rollback(ctx) }()
record := Record{ID: id, Name: name, KeyPrefix: prefix, KeyHash: hash, Scopes: scopes, Enabled: true, RequestsPerMinute: requestsPerMinute, MonthlyRequestQuota: monthlyRequestQuota, MonthlyTokenQuota: monthlyTokenQuota, ExpiresAt: expiresAt}
err = tx.QueryRow(ctx, `
INSERT INTO gateway.api_keys (id, name, key_prefix, key_hash, scopes, requests_per_minute, monthly_request_quota, monthly_token_quota, expires_at, created_by)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, nullif($10,'')::uuid)
RETURNING created_at`, id, name, prefix, hash, scopes, requestsPerMinute, monthlyRequestQuota, monthlyTokenQuota, expiresAt, actorID).Scan(&record.CreatedAt)
if err != nil {
return Record{}, "", fmt.Errorf("%w: %v", ErrStore, err)
}
payload, _ := json.Marshal(map[string]any{"api_key_id": id, "name": name, "key_prefix": prefix, "requests_per_minute": requestsPerMinute, "monthly_request_quota": monthlyRequestQuota, "monthly_token_quota": monthlyTokenQuota})
if _, err := tx.Exec(ctx, `
INSERT INTO gateway.outbox_events (event_id, event_type, event_version, aggregate_type, aggregate_id, payload)
VALUES ($1, 'api_key.created', 1, 'api_key', $2, $3)`, eventID, id, payload); err != nil {
return Record{}, "", fmt.Errorf("%w: %v", ErrStore, err)
}
if err := tx.Commit(ctx); err != nil {
return Record{}, "", fmt.Errorf("%w: %v", ErrStore, err)
}
return record, secret, nil
}
func (r *Repository) Validate(ctx context.Context, hash []byte) (Principal, error) {
if r.pool == nil {
return Principal{}, ErrStore
}
var principal Principal
err := r.pool.QueryRow(ctx, `
UPDATE gateway.api_keys
SET last_used_at = CASE WHEN last_used_at IS NULL OR last_used_at < clock_timestamp() - interval '5 minutes' THEN clock_timestamp() ELSE last_used_at END
WHERE key_hash = $1 AND enabled AND (expires_at IS NULL OR expires_at > clock_timestamp())
RETURNING id::text, tenant_id::text, scopes, requests_per_minute, monthly_request_quota, monthly_token_quota`, hash).Scan(&principal.APIKeyID, &principal.TenantID, &principal.Scopes, &principal.RequestsPerMinute, &principal.MonthlyRequestQuota, &principal.MonthlyTokenQuota)
if errors.Is(err, pgx.ErrNoRows) {
return Principal{}, ErrInvalid
}
if err != nil {
return Principal{}, fmt.Errorf("%w: %v", ErrStore, err)
}
return principal, nil
}
func (r *Repository) Revoke(ctx context.Context, id, actorID string) ([]byte, error) {
if r.pool == nil {
return nil, ErrStore
}
eventID, err := platformid.NewUUID()
if err != nil {
return nil, err
}
tx, err := r.pool.Begin(ctx)
if err != nil {
return nil, fmt.Errorf("%w: %v", ErrStore, err)
}
defer func() { _ = tx.Rollback(ctx) }()
var hash []byte
err = tx.QueryRow(ctx, `
UPDATE gateway.api_keys SET enabled = false, updated_at = clock_timestamp()
WHERE id = $1 AND enabled RETURNING key_hash`, id).Scan(&hash)
if errors.Is(err, pgx.ErrNoRows) {
return nil, ErrInvalid
}
if err != nil {
return nil, fmt.Errorf("%w: %v", ErrStore, err)
}
payload, _ := json.Marshal(map[string]any{"api_key_id": id, "actor_id": actorID})
if _, err := tx.Exec(ctx, `
INSERT INTO gateway.outbox_events (event_id, event_type, event_version, aggregate_type, aggregate_id, payload)
VALUES ($1, 'api_key.revoked', 1, 'api_key', $2, $3)`, eventID, id, payload); err != nil {
return nil, fmt.Errorf("%w: %v", ErrStore, err)
}
if err := tx.Commit(ctx); err != nil {
return nil, fmt.Errorf("%w: %v", ErrStore, err)
}
return hash, nil
}
func (r *Repository) UpdateLimits(ctx context.Context, id string, requestsPerMinute int, monthlyRequestQuota, monthlyTokenQuota int64, actorID string) (Record, []byte, error) {
if r.pool == nil {
return Record{}, nil, ErrStore
}
eventID, err := platformid.NewUUID()
if err != nil {
return Record{}, nil, err
}
tx, err := r.pool.Begin(ctx)
if err != nil {
return Record{}, nil, fmt.Errorf("%w: %v", ErrStore, err)
}
defer func() { _ = tx.Rollback(ctx) }()
var record Record
err = tx.QueryRow(ctx, `
UPDATE gateway.api_keys
SET requests_per_minute = $2, monthly_request_quota = $3, monthly_token_quota = $4, updated_at = clock_timestamp()
WHERE id = $1 AND enabled
RETURNING id::text, tenant_id::text, name, key_prefix, key_hash, scopes, enabled,
requests_per_minute, monthly_request_quota, monthly_token_quota, expires_at, last_used_at, created_at`,
id, requestsPerMinute, monthlyRequestQuota, monthlyTokenQuota,
).Scan(&record.ID, &record.TenantID, &record.Name, &record.KeyPrefix, &record.KeyHash, &record.Scopes, &record.Enabled,
&record.RequestsPerMinute, &record.MonthlyRequestQuota, &record.MonthlyTokenQuota, &record.ExpiresAt, &record.LastUsedAt, &record.CreatedAt)
if errors.Is(err, pgx.ErrNoRows) {
return Record{}, nil, ErrInvalid
}
if err != nil {
return Record{}, nil, fmt.Errorf("%w: %v", ErrStore, err)
}
payload, _ := json.Marshal(map[string]any{
"api_key_id": id, "actor_id": actorID, "requests_per_minute": requestsPerMinute,
"monthly_request_quota": monthlyRequestQuota, "monthly_token_quota": monthlyTokenQuota,
})
if _, err := tx.Exec(ctx, `
INSERT INTO gateway.outbox_events (event_id, event_type, event_version, aggregate_type, aggregate_id, payload)
VALUES ($1, 'api_key.limits_updated', 1, 'api_key', $2, $3)`, eventID, id, payload); err != nil {
return Record{}, nil, fmt.Errorf("%w: %v", ErrStore, err)
}
if err := tx.Commit(ctx); err != nil {
return Record{}, nil, fmt.Errorf("%w: %v", ErrStore, err)
}
return record, record.KeyHash, nil
}
+42
View File
@@ -0,0 +1,42 @@
package apikey
import (
"context"
"fmt"
"time"
"github.com/redis/go-redis/v9"
)
type UsageStore struct{ redis *redis.Client }
func NewUsageStore(client *redis.Client) *UsageStore { return &UsageStore{redis: client} }
func MonthlyTokenUsageKey(apiKeyID string, now time.Time) string {
return fmt.Sprintf("gateway:usage:api-key:%s:tokens:%s", apiKeyID, now.UTC().Format("200601"))
}
func (s *UsageStore) MonthlyTokens(ctx context.Context, apiKeyIDs []string, now time.Time) (map[string]int64, error) {
usage := make(map[string]int64, len(apiKeyIDs))
if len(apiKeyIDs) == 0 || s == nil || s.redis == nil {
return usage, nil
}
pipe := s.redis.Pipeline()
commands := make(map[string]*redis.StringCmd, len(apiKeyIDs))
for _, id := range apiKeyIDs {
commands[id] = pipe.Get(ctx, MonthlyTokenUsageKey(id, now))
}
_, err := pipe.Exec(ctx)
if err != nil && err != redis.Nil {
return nil, fmt.Errorf("%w: %v", ErrStore, err)
}
for id, command := range commands {
value, commandErr := command.Int64()
if commandErr == nil {
usage[id] = max(value, 0)
} else if commandErr != redis.Nil {
return nil, fmt.Errorf("%w: %v", ErrStore, commandErr)
}
}
return usage, nil
}
+144
View File
@@ -0,0 +1,144 @@
package audit
import (
"errors"
"net/http"
"strconv"
"strings"
"time"
"aigateway.local/core/internal/identity"
"aigateway.local/core/internal/platform/apiresponse"
)
type AdminHTTPHandler struct {
query *QueryService
identity *identity.Service
mux *http.ServeMux
}
func NewAdminHTTPHandler(query *QueryService, identityService *identity.Service) *AdminHTTPHandler {
h := &AdminHTTPHandler{query: query, identity: identityService, mux: http.NewServeMux()}
h.mux.HandleFunc("GET /api/v1/admin/audit-events", h.listEvents)
h.mux.HandleFunc("GET /api/v1/admin/usage/daily", h.listDailyUsage)
return h
}
func (h *AdminHTTPHandler) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
h.mux.ServeHTTP(writer, request)
}
func (h *AdminHTTPHandler) listEvents(writer http.ResponseWriter, request *http.Request) {
if !h.requirePermission(writer, request, identity.PermissionAuditRead) {
return
}
now := time.Now().UTC()
from, err := queryTime(request, "from", now.Add(-24*time.Hour))
if err != nil {
apiresponse.Error(writer, http.StatusBadRequest, "from 时间无效")
return
}
to, err := queryTime(request, "to", now.Add(time.Second))
if err != nil || !to.After(from) || to.Sub(from) > 366*24*time.Hour {
apiresponse.Error(writer, http.StatusBadRequest, "审计查询时间范围无效或超过 366 天")
return
}
limit := 50
if value := request.URL.Query().Get("limit"); value != "" {
limit, err = strconv.Atoi(value)
if err != nil || limit < 1 || limit > 200 {
apiresponse.Error(writer, http.StatusBadRequest, "limit 必须在 1 到 200 之间")
return
}
}
var before *time.Time
if value := request.URL.Query().Get("before"); value != "" {
parsed, parseErr := time.Parse(time.RFC3339Nano, value)
if parseErr != nil {
apiresponse.Error(writer, http.StatusBadRequest, "before 游标无效")
return
}
before = &parsed
}
var status *int
if value := request.URL.Query().Get("status"); value != "" {
parsed, parseErr := strconv.Atoi(value)
if parseErr != nil || parsed < 100 || parsed > 599 {
apiresponse.Error(writer, http.StatusBadRequest, "status 无效")
return
}
status = &parsed
}
items, err := h.query.ListEvents(request.Context(), EventFilter{
From: from, To: to, Before: before, APIKeyID: strings.TrimSpace(request.URL.Query().Get("api_key_id")),
Provider: strings.TrimSpace(request.URL.Query().Get("provider")), Model: strings.TrimSpace(request.URL.Query().Get("model")), StatusCode: status, Limit: limit,
})
if err != nil {
apiresponse.Error(writer, http.StatusServiceUnavailable, "审计查询服务暂不可用")
return
}
next := ""
if len(items) == limit {
next = items[len(items)-1].RecordedAt.Format(time.RFC3339Nano)
}
apiresponse.OK(writer, map[string]any{"items": items, "next_before": next})
}
func (h *AdminHTTPHandler) listDailyUsage(writer http.ResponseWriter, request *http.Request) {
if !h.requirePermission(writer, request, identity.PermissionUsageRead) {
return
}
now := time.Now().UTC()
from, err := queryDate(request, "from", now.AddDate(0, 0, -29))
if err != nil {
apiresponse.Error(writer, http.StatusBadRequest, "from 日期无效")
return
}
to, err := queryDate(request, "to", now)
if err != nil || to.Before(from) || to.Sub(from) > 366*24*time.Hour {
apiresponse.Error(writer, http.StatusBadRequest, "usage 查询日期范围无效或超过 366 天")
return
}
items, err := h.query.ListDailyUsage(request.Context(), UsageFilter{
From: from, To: to, APIKeyID: strings.TrimSpace(request.URL.Query().Get("api_key_id")),
Provider: strings.TrimSpace(request.URL.Query().Get("provider")), Model: strings.TrimSpace(request.URL.Query().Get("model")),
})
if err != nil {
apiresponse.Error(writer, http.StatusServiceUnavailable, "usage 查询服务暂不可用")
return
}
apiresponse.OK(writer, items)
}
func (h *AdminHTTPHandler) requirePermission(writer http.ResponseWriter, request *http.Request, permission string) bool {
account, err := h.identity.Authenticate(request.Context(), identity.KindAdmin, request.Header.Get("Authorization"))
if err != nil {
status := http.StatusUnauthorized
if !errors.Is(err, identity.ErrInvalidSession) && !errors.Is(err, identity.ErrNotFound) {
status = http.StatusServiceUnavailable
}
apiresponse.Error(writer, status, "登录状态无效或身份服务暂不可用")
return false
}
if !identity.HasPermission(account, permission) {
apiresponse.Error(writer, http.StatusForbidden, "缺少审计或 usage 查看权限")
return false
}
return true
}
func queryTime(request *http.Request, name string, fallback time.Time) (time.Time, error) {
value := request.URL.Query().Get(name)
if value == "" {
return fallback, nil
}
return time.Parse(time.RFC3339, value)
}
func queryDate(request *http.Request, name string, fallback time.Time) (time.Time, error) {
value := request.URL.Query().Get(name)
if value == "" {
value = fallback.Format("2006-01-02")
}
return time.Parse("2006-01-02", value)
}
+147
View File
@@ -0,0 +1,147 @@
package audit
import (
"context"
"errors"
"fmt"
"regexp"
"sort"
"time"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgxpool"
)
const maintenanceLockID int64 = 6720240816
var auditPartitionPattern = regexp.MustCompile(`^audit_events_(\d{4})(\d{2})$`)
type MaintenanceResult struct {
CreatedPartitions []string `json:"created_partitions"`
DroppedPartitions []string `json:"dropped_partitions"`
DeletedAuditRows int64 `json:"deleted_audit_rows"`
DeletedUsageRows int64 `json:"deleted_usage_rows"`
}
type Maintenance struct {
pool *pgxpool.Pool
auditRetention time.Duration
usageRetention time.Duration
monthsAhead int
}
func NewMaintenance(pool *pgxpool.Pool, auditRetention, usageRetention time.Duration, monthsAhead int) *Maintenance {
return &Maintenance{pool: pool, auditRetention: auditRetention, usageRetention: usageRetention, monthsAhead: monthsAhead}
}
func (m *Maintenance) Run(ctx context.Context, now time.Time) (MaintenanceResult, error) {
var result MaintenanceResult
if m == nil || m.pool == nil {
return result, errors.New("audit maintenance store unavailable")
}
now = now.UTC()
auditCutoff := now.Add(-m.auditRetention)
usageCutoff := now.Add(-m.usageRetention)
tx, err := m.pool.Begin(ctx)
if err != nil {
return result, fmt.Errorf("begin audit maintenance: %w", err)
}
defer func() { _ = tx.Rollback(ctx) }()
if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock($1)`, maintenanceLockID); err != nil {
return result, fmt.Errorf("lock audit maintenance: %w", err)
}
if _, err := tx.Exec(ctx, `LOCK TABLE gateway.audit_events IN ACCESS EXCLUSIVE MODE`); err != nil {
return result, fmt.Errorf("lock audit table: %w", err)
}
partitions, err := listAuditPartitions(ctx, tx)
if err != nil {
return result, err
}
for name := range partitions {
start, ok := auditPartitionMonth(name)
if !ok || start.AddDate(0, 1, 0).After(auditCutoff) {
continue
}
if _, err := tx.Exec(ctx, `DROP TABLE `+pgx.Identifier{"gateway", name}.Sanitize()); err != nil {
return result, fmt.Errorf("drop audit partition %s: %w", name, err)
}
result.DroppedPartitions = append(result.DroppedPartitions, name)
}
if _, err := tx.Exec(ctx, `ALTER TABLE gateway.audit_events DETACH PARTITION gateway.audit_events_default`); err != nil {
return result, fmt.Errorf("detach audit default partition: %w", err)
}
deleted, err := tx.Exec(ctx, `DELETE FROM gateway.audit_events_default WHERE recorded_at < $1`, auditCutoff)
if err != nil {
return result, fmt.Errorf("clean audit default partition: %w", err)
}
result.DeletedAuditRows += deleted.RowsAffected()
from := monthStart(auditCutoff)
through := monthStart(now).AddDate(0, m.monthsAhead+1, 0)
for start := from; start.Before(through); start = start.AddDate(0, 1, 0) {
end := start.AddDate(0, 1, 0)
name := "audit_events_" + start.Format("200601")
if _, exists := partitions[name]; !exists {
statement := fmt.Sprintf(`CREATE TABLE %s PARTITION OF gateway.audit_events FOR VALUES FROM ('%s') TO ('%s')`, pgx.Identifier{"gateway", name}.Sanitize(), start.Format(time.RFC3339), end.Format(time.RFC3339))
if _, err := tx.Exec(ctx, statement); err != nil {
return result, fmt.Errorf("create audit partition %s: %w", name, err)
}
result.CreatedPartitions = append(result.CreatedPartitions, name)
}
statement := `WITH moved AS (DELETE FROM gateway.audit_events_default WHERE recorded_at >= $1 AND recorded_at < $2 RETURNING *) INSERT INTO gateway.audit_events SELECT * FROM moved`
if _, err := tx.Exec(ctx, statement, start, end); err != nil {
return result, fmt.Errorf("move default audit rows into %s: %w", name, err)
}
}
if _, err := tx.Exec(ctx, `ALTER TABLE gateway.audit_events ATTACH PARTITION gateway.audit_events_default DEFAULT`); err != nil {
return result, fmt.Errorf("reattach audit default partition: %w", err)
}
deleted, err = tx.Exec(ctx, `DELETE FROM gateway.audit_events WHERE recorded_at < $1`, auditCutoff)
if err != nil {
return result, fmt.Errorf("apply exact audit retention: %w", err)
}
result.DeletedAuditRows += deleted.RowsAffected()
deleted, err = tx.Exec(ctx, `DELETE FROM gateway.usage_daily WHERE usage_date < $1::date`, usageCutoff.Format("2006-01-02"))
if err != nil {
return result, fmt.Errorf("apply usage retention: %w", err)
}
result.DeletedUsageRows = deleted.RowsAffected()
if err := tx.Commit(ctx); err != nil {
return result, fmt.Errorf("commit audit maintenance: %w", err)
}
sort.Strings(result.CreatedPartitions)
sort.Strings(result.DroppedPartitions)
return result, nil
}
func listAuditPartitions(ctx context.Context, tx pgx.Tx) (map[string]struct{}, error) {
rows, err := tx.Query(ctx, `SELECT child.relname FROM pg_inherits i JOIN pg_class parent ON parent.oid=i.inhparent JOIN pg_namespace n ON n.oid=parent.relnamespace JOIN pg_class child ON child.oid=i.inhrelid WHERE n.nspname='gateway' AND parent.relname='audit_events'`)
if err != nil {
return nil, fmt.Errorf("list audit partitions: %w", err)
}
defer rows.Close()
result := make(map[string]struct{})
for rows.Next() {
var name string
if err := rows.Scan(&name); err != nil {
return nil, fmt.Errorf("scan audit partition: %w", err)
}
if name != "audit_events_default" {
result[name] = struct{}{}
}
}
return result, rows.Err()
}
func auditPartitionMonth(name string) (time.Time, bool) {
match := auditPartitionPattern.FindStringSubmatch(name)
if match == nil {
return time.Time{}, false
}
parsed, err := time.Parse("200601", match[1]+match[2])
return parsed.UTC(), err == nil
}
func monthStart(value time.Time) time.Time {
value = value.UTC()
return time.Date(value.Year(), value.Month(), 1, 0, 0, 0, 0, time.UTC)
}
@@ -0,0 +1,73 @@
package audit
import (
"context"
"os"
"testing"
"time"
"github.com/jackc/pgx/v5/pgxpool"
)
func TestMaintenancePartitionsRetentionAndIdempotency(t *testing.T) {
databaseURL := os.Getenv("AUDIT_MAINTENANCE_TEST_DATABASE_URL")
if databaseURL == "" {
t.Skip("AUDIT_MAINTENANCE_TEST_DATABASE_URL is not configured")
}
ctx := context.Background()
pool, err := pgxpool.New(ctx, databaseURL)
if err != nil {
t.Fatal(err)
}
defer pool.Close()
const oldID = "44444444-4444-4444-8444-444444444444"
const currentID = "55555555-5555-4555-8555-555555555555"
_, err = pool.Exec(ctx, `INSERT INTO gateway.audit_events(id,request_id,protocol,status_code,recorded_at) VALUES
($1,'maintenance-old','/v1/test',200,'2026-01-15T00:00:00Z'),
($2,'maintenance-current','/v1/test',200,'2026-08-10T00:00:00Z')`, oldID, currentID)
if err != nil {
t.Fatal(err)
}
now := time.Date(2026, 8, 11, 0, 0, 0, 0, time.UTC)
first, err := NewMaintenance(pool, 300*24*time.Hour, 730*24*time.Hour, 1).Run(ctx, now)
if err != nil {
t.Fatal(err)
}
if !contains(first.CreatedPartitions, "audit_events_202601") || !contains(first.CreatedPartitions, "audit_events_202608") {
t.Fatalf("expected old and current partitions, got %#v", first.CreatedPartitions)
}
second, err := NewMaintenance(pool, 90*24*time.Hour, 730*24*time.Hour, 1).Run(ctx, now)
if err != nil {
t.Fatal(err)
}
if !contains(second.DroppedPartitions, "audit_events_202601") {
t.Fatalf("expected stale partition to be dropped, got %#v", second.DroppedPartitions)
}
var currentTable string
if err := pool.QueryRow(ctx, `SELECT tableoid::regclass::text FROM gateway.audit_events WHERE id=$1`, currentID).Scan(&currentTable); err != nil {
t.Fatal(err)
}
if currentTable != "gateway.audit_events_202608" && currentTable != "audit_events_202608" {
t.Fatalf("current row was not routed to monthly partition: %s", currentTable)
}
var oldCount int
if err := pool.QueryRow(ctx, `SELECT count(*) FROM gateway.audit_events WHERE id=$1`, oldID).Scan(&oldCount); err != nil || oldCount != 0 {
t.Fatalf("expired row still exists: count=%d err=%v", oldCount, err)
}
third, err := NewMaintenance(pool, 90*24*time.Hour, 730*24*time.Hour, 1).Run(ctx, now)
if err != nil {
t.Fatal(err)
}
if len(third.CreatedPartitions) != 0 || len(third.DroppedPartitions) != 0 {
t.Fatalf("maintenance must be idempotent: %#v", third)
}
}
func contains(values []string, wanted string) bool {
for _, value := range values {
if value == wanted {
return true
}
}
return false
}
+23
View File
@@ -0,0 +1,23 @@
package audit
import (
"testing"
"time"
)
func TestAuditPartitionMonth(t *testing.T) {
month, ok := auditPartitionMonth("audit_events_202608")
if !ok || month.Format("2006-01-02") != "2026-08-01" {
t.Fatalf("unexpected month: %v %v", month, ok)
}
if _, ok := auditPartitionMonth("audit_events_default"); ok {
t.Fatal("default partition must not parse as a monthly partition")
}
}
func TestMonthStartUsesUTC(t *testing.T) {
value := time.Date(2026, 8, 31, 23, 0, 0, 0, time.FixedZone("UTC-2", -2*3600))
if got := monthStart(value); got.Format(time.RFC3339) != "2026-09-01T00:00:00Z" {
t.Fatalf("unexpected UTC month start: %s", got)
}
}
+151
View File
@@ -0,0 +1,151 @@
package audit
import (
"context"
"encoding/json"
"fmt"
"strconv"
"strings"
"time"
"github.com/jackc/pgx/v5/pgxpool"
)
type QueryService struct{ pool *pgxpool.Pool }
func NewQueryService(pool *pgxpool.Pool) *QueryService { return &QueryService{pool: pool} }
type EventView struct {
ID string `json:"id"`
TenantID *string `json:"tenant_id"`
RequestID string `json:"request_id"`
APIKeyID *string `json:"api_key_id"`
APIKeyName *string `json:"api_key_name"`
ProviderCode *string `json:"provider_code"`
Model *string `json:"model"`
Protocol string `json:"protocol"`
StatusCode *int `json:"status_code"`
PromptTokens *int64 `json:"prompt_tokens"`
CompletionTokens *int64 `json:"completion_tokens"`
CostMicrounits *int64 `json:"cost_microunits"`
LatencyMS *int `json:"latency_ms"`
Labels json.RawMessage `json:"labels"`
RecordedAt time.Time `json:"recorded_at"`
}
type EventFilter struct {
From, To time.Time
Before *time.Time
APIKeyID string
Provider string
Model string
StatusCode *int
Limit int
}
func (s *QueryService) ListEvents(ctx context.Context, filter EventFilter) ([]EventView, error) {
if s == nil || s.pool == nil {
return nil, fmt.Errorf("audit store unavailable")
}
where := []string{"a.recorded_at >= $1", "a.recorded_at < $2"}
args := []any{filter.From, filter.To}
add := func(condition string, value any) {
args = append(args, value)
where = append(where, fmt.Sprintf(condition, len(args)))
}
if filter.Before != nil {
add("a.recorded_at < $%d", *filter.Before)
}
if filter.APIKeyID != "" {
add("a.api_key_id = $%d", filter.APIKeyID)
}
if filter.Provider != "" {
add("a.provider_code = $%d", filter.Provider)
}
if filter.Model != "" {
add("a.model = $%d", filter.Model)
}
if filter.StatusCode != nil {
add("a.status_code = $%d", *filter.StatusCode)
}
args = append(args, filter.Limit)
query := `SELECT a.id::text,a.tenant_id::text,a.request_id,a.api_key_id::text,k.name,
a.provider_code,a.model,a.protocol,a.status_code,a.prompt_tokens,a.completion_tokens,a.cost_microunits,a.latency_ms,a.labels,a.recorded_at
FROM gateway.audit_events a LEFT JOIN gateway.api_keys k ON k.id=a.api_key_id
WHERE ` + strings.Join(where, " AND ") + ` ORDER BY a.recorded_at DESC,a.id DESC LIMIT $` + strconv.Itoa(len(args))
rows, err := s.pool.Query(ctx, query, args...)
if err != nil {
return nil, fmt.Errorf("query audit events: %w", err)
}
defer rows.Close()
items := make([]EventView, 0)
for rows.Next() {
var item EventView
if err := rows.Scan(&item.ID, &item.TenantID, &item.RequestID, &item.APIKeyID, &item.APIKeyName,
&item.ProviderCode, &item.Model, &item.Protocol, &item.StatusCode, &item.PromptTokens,
&item.CompletionTokens, &item.CostMicrounits, &item.LatencyMS, &item.Labels, &item.RecordedAt); err != nil {
return nil, fmt.Errorf("scan audit event: %w", err)
}
items = append(items, item)
}
return items, rows.Err()
}
type DailyUsageView struct {
Date time.Time `json:"date"`
APIKeyID string `json:"api_key_id"`
APIKeyName string `json:"api_key_name"`
ProviderCode string `json:"provider_code"`
Model string `json:"model"`
Requests int64 `json:"requests"`
FailedRequests int64 `json:"failed_requests"`
PromptTokens int64 `json:"prompt_tokens"`
CompletionTokens int64 `json:"completion_tokens"`
CostMicrounits int64 `json:"cost_microunits"`
}
type UsageFilter struct {
From, To time.Time
APIKeyID string
Provider string
Model string
}
func (s *QueryService) ListDailyUsage(ctx context.Context, filter UsageFilter) ([]DailyUsageView, error) {
if s == nil || s.pool == nil {
return nil, fmt.Errorf("usage store unavailable")
}
where := []string{"u.usage_date >= $1", "u.usage_date <= $2"}
args := []any{filter.From, filter.To}
add := func(column string, value any) {
args = append(args, value)
where = append(where, column+" = $"+strconv.Itoa(len(args)))
}
if filter.APIKeyID != "" {
add("u.api_key_id", filter.APIKeyID)
}
if filter.Provider != "" {
add("u.provider_code", filter.Provider)
}
if filter.Model != "" {
add("u.model", filter.Model)
}
rows, err := s.pool.Query(ctx, `SELECT u.usage_date,u.api_key_id::text,k.name,u.provider_code,u.model,
u.requests,u.failed_requests,u.prompt_tokens,u.completion_tokens,u.cost_microunits
FROM gateway.usage_daily u JOIN gateway.api_keys k ON k.id=u.api_key_id WHERE `+
strings.Join(where, " AND ")+` ORDER BY u.usage_date DESC,k.name,u.provider_code,u.model`, args...)
if err != nil {
return nil, fmt.Errorf("query daily usage: %w", err)
}
defer rows.Close()
items := make([]DailyUsageView, 0)
for rows.Next() {
var item DailyUsageView
if err := rows.Scan(&item.Date, &item.APIKeyID, &item.APIKeyName, &item.ProviderCode, &item.Model,
&item.Requests, &item.FailedRequests, &item.PromptTokens, &item.CompletionTokens, &item.CostMicrounits); err != nil {
return nil, fmt.Errorf("scan daily usage: %w", err)
}
items = append(items, item)
}
return items, rows.Err()
}
+282
View File
@@ -0,0 +1,282 @@
package audit
import (
"context"
"encoding/json"
"fmt"
"log/slog"
"strings"
"sync/atomic"
"time"
platformid "aigateway.local/core/internal/platform/id"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgtype"
"github.com/jackc/pgx/v5/pgxpool"
)
type Event struct {
TenantID *string
RequestID string
APIKeyID *string
ProviderCode string
Model string
Protocol string
StatusCode int
PromptTokens int64
CompletionTokens int64
CostMicrounits int64
PriceID string
Currency string
LatencyMS int
Labels map[string]any
RecordedAt time.Time
}
type Stats struct {
Accepted uint64
Dropped uint64
Written uint64
Failures uint64
}
type Recorder struct {
pool *pgxpool.Pool
logger *slog.Logger
queue chan Event
batchSize int
flushInterval time.Duration
accepted atomic.Uint64
dropped atomic.Uint64
written atomic.Uint64
failures atomic.Uint64
}
func NewRecorder(pool *pgxpool.Pool, logger *slog.Logger, queueSize, batchSize int, flushInterval time.Duration) *Recorder {
if queueSize < 1 {
queueSize = 4096
}
if batchSize < 1 {
batchSize = 200
}
if flushInterval <= 0 {
flushInterval = time.Second
}
return &Recorder{pool: pool, logger: logger, queue: make(chan Event, queueSize), batchSize: batchSize, flushInterval: flushInterval}
}
func (r *Recorder) Record(event Event) bool {
if r == nil || r.pool == nil {
return false
}
if event.RecordedAt.IsZero() {
event.RecordedAt = time.Now().UTC()
}
select {
case r.queue <- event:
r.accepted.Add(1)
return true
default:
r.dropped.Add(1)
return false
}
}
func (r *Recorder) Stats() Stats {
if r == nil {
return Stats{}
}
return Stats{Accepted: r.accepted.Load(), Dropped: r.dropped.Load(), Written: r.written.Load(), Failures: r.failures.Load()}
}
func (r *Recorder) Prometheus() string {
stats := r.Stats()
return fmt.Sprintf("gateway_audit_events_accepted_total %d\ngateway_audit_events_dropped_total %d\ngateway_audit_events_written_total %d\ngateway_audit_flush_failures_total %d\n", stats.Accepted, stats.Dropped, stats.Written, stats.Failures)
}
func (r *Recorder) Run(ctx context.Context) {
ticker := time.NewTicker(r.flushInterval)
defer ticker.Stop()
pending := make([]Event, 0, r.batchSize)
flushAllowed := true
for {
var events <-chan Event = r.queue
if len(pending) >= cap(r.queue)+r.batchSize {
events = nil
}
select {
case event := <-events:
pending = append(pending, event)
if len(pending) >= r.batchSize && flushAllowed {
if r.flush(ctx, pending) == nil {
pending = pending[:0]
} else {
flushAllowed = false
}
}
case <-ticker.C:
flushAllowed = true
if len(pending) > 0 {
if r.flush(ctx, pending) == nil {
pending = pending[:0]
} else {
flushAllowed = false
}
}
case <-ctx.Done():
for {
select {
case event := <-r.queue:
pending = append(pending, event)
default:
flushCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
if len(pending) > 0 {
_ = r.flush(flushCtx, pending)
}
cancel()
return
}
}
}
}
}
type dailyKey struct {
date, apiKeyID, provider, model string
}
type dailyValue struct {
requests, failed, prompt, completion, cost int64
}
type policyAlert struct {
eventID, requestID string
tenantID *string
payload []byte
}
func (r *Recorder) flush(ctx context.Context, events []Event) error {
if len(events) == 0 {
return nil
}
tx, err := r.pool.Begin(ctx)
if err != nil {
return r.flushError(err)
}
defer func() { _ = tx.Rollback(ctx) }()
rows := make([][]any, 0, len(events))
daily := make(map[dailyKey]dailyValue)
alerts := make([]policyAlert, 0)
for _, event := range events {
id, err := platformid.NewUUID()
if err != nil {
return r.flushError(err)
}
labels, _ := json.Marshal(event.Labels)
if policies, matched := event.Labels["content_policies"]; matched {
alertID, idErr := platformid.NewUUID()
if idErr != nil {
return r.flushError(idErr)
}
payload, _ := json.Marshal(map[string]any{
"request_id": event.RequestID, "protocol": event.Protocol, "status_code": event.StatusCode,
"provider_code": event.ProviderCode, "model": event.Model, "content_policies": policies,
"content_redacted": event.Labels["content_redacted"], "recorded_at": event.RecordedAt,
})
alerts = append(alerts, policyAlert{eventID: alertID, requestID: event.RequestID, tenantID: event.TenantID, payload: payload})
}
rows = append(rows, []any{
uuidValue(id), uuidPointer(event.TenantID), event.RequestID, nil, uuidPointer(event.APIKeyID),
nullString(event.ProviderCode), nullString(event.Model), event.Protocol, event.StatusCode,
nullToken(event.PromptTokens), nullToken(event.CompletionTokens), nullCost(event.CostMicrounits), event.LatencyMS,
nil, nil, labels, event.RecordedAt,
})
if event.APIKeyID != nil && *event.APIKeyID != "" {
key := dailyKey{date: event.RecordedAt.UTC().Format("2006-01-02"), apiKeyID: *event.APIKeyID, provider: event.ProviderCode, model: event.Model}
value := daily[key]
value.requests++
if event.StatusCode >= 400 {
value.failed++
}
value.prompt += max(event.PromptTokens, 0)
value.completion += max(event.CompletionTokens, 0)
value.cost += max(event.CostMicrounits, 0)
daily[key] = value
}
}
_, err = tx.CopyFrom(ctx, pgx.Identifier{"gateway", "audit_events"}, []string{
"id", "tenant_id", "request_id", "actor_id", "api_key_id", "provider_code", "model", "protocol",
"status_code", "prompt_tokens", "completion_tokens", "cost_microunits", "latency_ms",
"request_preview", "response_preview", "labels", "recorded_at",
}, pgx.CopyFromRows(rows))
if err != nil {
return r.flushError(err)
}
batch := &pgx.Batch{}
for key, value := range daily {
batch.Queue(`INSERT INTO gateway.usage_daily(usage_date,api_key_id,provider_code,model,requests,failed_requests,prompt_tokens,completion_tokens,cost_microunits)
VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9)
ON CONFLICT(usage_date,api_key_id,provider_code,model) DO UPDATE SET
requests=gateway.usage_daily.requests+EXCLUDED.requests,
failed_requests=gateway.usage_daily.failed_requests+EXCLUDED.failed_requests,
prompt_tokens=gateway.usage_daily.prompt_tokens+EXCLUDED.prompt_tokens,
completion_tokens=gateway.usage_daily.completion_tokens+EXCLUDED.completion_tokens,
cost_microunits=gateway.usage_daily.cost_microunits+EXCLUDED.cost_microunits,
updated_at=clock_timestamp()`, key.date, key.apiKeyID, key.provider, key.model, value.requests, value.failed, value.prompt, value.completion, value.cost)
}
for _, alert := range alerts {
batch.Queue(`INSERT INTO gateway.outbox_events(event_id,event_type,event_version,tenant_id,aggregate_type,aggregate_id,payload)
VALUES($1,'content_policy.matched',1,$2,'gateway_request',$3,$4)`, alert.eventID, uuidPointer(alert.tenantID), alert.requestID, alert.payload)
}
results := tx.SendBatch(ctx, batch)
if err := results.Close(); err != nil {
return r.flushError(err)
}
if err := tx.Commit(ctx); err != nil {
return r.flushError(err)
}
r.written.Add(uint64(len(events)))
return nil
}
func (r *Recorder) flushError(err error) error {
r.failures.Add(1)
if r.logger != nil {
r.logger.Error("audit batch flush failed", "error", err)
}
return err
}
func uuidValue(value string) pgtype.UUID {
var result pgtype.UUID
_ = result.Scan(value)
return result
}
func uuidPointer(value *string) any {
if value == nil || strings.TrimSpace(*value) == "" {
return nil
}
return uuidValue(*value)
}
func nullString(value string) any {
if value == "" {
return nil
}
return value
}
func nullToken(value int64) any {
if value <= 0 {
return nil
}
return value
}
func nullCost(value int64) any {
if value <= 0 {
return nil
}
return value
}
+119
View File
@@ -0,0 +1,119 @@
package contentpolicy
import (
"encoding/json"
"errors"
"net/http"
"aigateway.local/core/internal/identity"
"aigateway.local/core/internal/platform/apiresponse"
)
type AdminHTTPHandler struct {
store *Store
engine *Engine
identity *identity.Service
mux *http.ServeMux
}
func NewAdminHTTPHandler(store *Store, engine *Engine, identityService *identity.Service) *AdminHTTPHandler {
h := &AdminHTTPHandler{store: store, engine: engine, identity: identityService, mux: http.NewServeMux()}
h.mux.HandleFunc("GET /api/v1/admin/content-policies", h.list)
h.mux.HandleFunc("POST /api/v1/admin/content-policies", h.create)
h.mux.HandleFunc("PUT /api/v1/admin/content-policies/{policy_id}", h.update)
h.mux.HandleFunc("DELETE /api/v1/admin/content-policies/{policy_id}", h.delete)
return h
}
func (h *AdminHTTPHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { h.mux.ServeHTTP(w, r) }
func (h *AdminHTTPHandler) list(w http.ResponseWriter, r *http.Request) {
if _, ok := h.require(w, r, identity.PermissionContentPolicyRead); !ok {
return
}
items, err := h.store.List(r.Context())
if err != nil {
apiresponse.Error(w, 503, "内容策略服务暂不可用")
return
}
apiresponse.OK(w, items)
}
func (h *AdminHTTPHandler) create(w http.ResponseWriter, r *http.Request) {
actor, ok := h.require(w, r, identity.PermissionContentPolicyManage)
if !ok {
return
}
p, ok := decodePolicy(w, r)
if !ok {
return
}
saved, err := h.store.Save(r.Context(), p, actor.ID, true)
h.finish(w, r, saved, err)
}
func (h *AdminHTTPHandler) update(w http.ResponseWriter, r *http.Request) {
actor, ok := h.require(w, r, identity.PermissionContentPolicyManage)
if !ok {
return
}
p, ok := decodePolicy(w, r)
if !ok {
return
}
p.ID = r.PathValue("policy_id")
saved, err := h.store.Save(r.Context(), p, actor.ID, false)
h.finish(w, r, saved, err)
}
func (h *AdminHTTPHandler) delete(w http.ResponseWriter, r *http.Request) {
actor, ok := h.require(w, r, identity.PermissionContentPolicyManage)
if !ok {
return
}
err := h.store.Delete(r.Context(), r.PathValue("policy_id"), actor.ID)
if errors.Is(err, ErrNotFound) {
apiresponse.Error(w, 404, "内容策略不存在")
return
}
if err != nil {
apiresponse.Error(w, 503, "内容策略删除失败")
return
}
_ = h.engine.Reload(r.Context())
apiresponse.OK(w, map[string]bool{"deleted": true})
}
func (h *AdminHTTPHandler) finish(w http.ResponseWriter, r *http.Request, p Policy, err error) {
if errors.Is(err, ErrNotFound) {
apiresponse.Error(w, 404, "内容策略不存在")
return
}
if err != nil {
apiresponse.Error(w, 400, err.Error())
return
}
_ = h.engine.Reload(r.Context())
apiresponse.OK(w, p)
}
func decodePolicy(w http.ResponseWriter, r *http.Request) (Policy, bool) {
var p Policy
d := json.NewDecoder(http.MaxBytesReader(w, r.Body, 1<<20))
d.DisallowUnknownFields()
if d.Decode(&p) != nil {
apiresponse.Error(w, 400, "请求格式无效")
return p, false
}
Normalize(&p)
if err := ValidateInput(p); err != nil {
apiresponse.Error(w, 400, err.Error())
return p, false
}
return p, true
}
func (h *AdminHTTPHandler) require(w http.ResponseWriter, r *http.Request, permission string) (identity.Account, bool) {
a, err := h.identity.Authenticate(r.Context(), identity.KindAdmin, r.Header.Get("Authorization"))
if err != nil {
apiresponse.Error(w, 401, "登录状态无效")
return a, false
}
if !identity.HasPermission(a, permission) {
apiresponse.Error(w, 403, "缺少内容策略权限")
return a, false
}
return a, true
}
+272
View File
@@ -0,0 +1,272 @@
package contentpolicy
import (
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"io"
"log/slog"
"net/http"
"regexp"
"slices"
"strings"
"sync/atomic"
"time"
"github.com/jackc/pgx/v5/pgxpool"
)
var ErrBodyTooLarge = errors.New("request body is too large")
type Rule struct {
Name string `json:"name"`
Pattern string `json:"pattern"`
Replacement string `json:"replacement"`
}
type Policy struct {
ID string `json:"id"`
Name string `json:"name"`
Description string `json:"description"`
Action string `json:"action"`
Priority int `json:"priority"`
Paths []string `json:"paths"`
Models []string `json:"models"`
APIKeyIDs []string `json:"api_key_ids"`
Rules []Rule `json:"rules"`
Enabled bool `json:"enabled"`
Revision int64 `json:"revision"`
CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"`
}
type compiledRule struct {
name, replacement string
expression *regexp.Regexp
}
type compiledPolicy struct {
Policy
rules []compiledRule
}
type snapshot struct{ policies []compiledPolicy }
type Engine struct {
pool *pgxpool.Pool
logger *slog.Logger
refresh time.Duration
current atomic.Pointer[snapshot]
}
type Match struct {
PolicyID string `json:"policy_id"`
PolicyName string `json:"policy_name"`
Action string `json:"action"`
Rules []string `json:"rules"`
}
type Result struct {
Matches []Match
Blocked, Redacted bool
}
func NewEngine(pool *pgxpool.Pool, refresh time.Duration, logger *slog.Logger) *Engine {
if refresh <= 0 {
refresh = 30 * time.Second
}
e := &Engine{pool: pool, refresh: refresh, logger: logger}
e.current.Store(&snapshot{})
return e
}
func (e *Engine) Run(ctx context.Context) {
_ = e.Reload(ctx)
ticker := time.NewTicker(e.refresh)
defer ticker.Stop()
for {
select {
case <-ticker.C:
if err := e.Reload(ctx); err != nil && e.logger != nil {
e.logger.Warn("content policy refresh failed", "error", err)
}
case <-ctx.Done():
return
}
}
}
func (e *Engine) Reload(ctx context.Context) error {
if e == nil || e.pool == nil {
return errors.New("content policy store unavailable")
}
rows, err := e.pool.Query(ctx, `SELECT id::text,name,description,action,priority,paths,models,api_key_ids::text[],rules,enabled,revision,created_at,updated_at FROM gateway.content_policies WHERE enabled ORDER BY priority DESC,name,id`)
if err != nil {
return fmt.Errorf("load content policies: %w", err)
}
defer rows.Close()
loaded := make([]compiledPolicy, 0)
for rows.Next() {
var policy Policy
var raw []byte
if err := rows.Scan(&policy.ID, &policy.Name, &policy.Description, &policy.Action, &policy.Priority, &policy.Paths, &policy.Models, &policy.APIKeyIDs, &raw, &policy.Enabled, &policy.Revision, &policy.CreatedAt, &policy.UpdatedAt); err != nil {
return err
}
if err := json.Unmarshal(raw, &policy.Rules); err != nil {
return fmt.Errorf("decode policy %s: %w", policy.ID, err)
}
compiled, err := compilePolicy(policy)
if err != nil {
return fmt.Errorf("compile policy %s: %w", policy.ID, err)
}
loaded = append(loaded, compiled)
}
if err := rows.Err(); err != nil {
return err
}
e.current.Store(&snapshot{policies: loaded})
return nil
}
func compilePolicy(policy Policy) (compiledPolicy, error) {
if policy.Action != "audit" && policy.Action != "block" && policy.Action != "redact" {
return compiledPolicy{}, errors.New("unsupported action")
}
if len(policy.Rules) == 0 || len(policy.Rules) > 20 {
return compiledPolicy{}, errors.New("rules must contain 1 to 20 entries")
}
result := compiledPolicy{Policy: policy, rules: make([]compiledRule, 0, len(policy.Rules))}
for _, rule := range policy.Rules {
if strings.TrimSpace(rule.Name) == "" || len(rule.Pattern) == 0 || len(rule.Pattern) > 512 {
return compiledPolicy{}, errors.New("invalid rule name or pattern length")
}
expression, err := regexp.Compile(rule.Pattern)
if err != nil {
return compiledPolicy{}, err
}
replacement := rule.Replacement
if replacement == "" {
replacement = "[REDACTED]"
}
result.rules = append(result.rules, compiledRule{name: rule.Name, replacement: replacement, expression: expression})
}
return result, nil
}
func Validate(policy Policy) error { _, err := compilePolicy(policy); return err }
func (e *Engine) Apply(request *http.Request, maxBody int64, apiKeyID string) (Result, error) {
if e == nil || request.Body == nil || request.Body == http.NoBody || request.Method == http.MethodGet {
return Result{}, nil
}
body, err := io.ReadAll(io.LimitReader(request.Body, maxBody+1))
if err != nil {
return Result{}, err
}
_ = request.Body.Close()
if int64(len(body)) > maxBody {
return Result{}, ErrBodyTooLarge
}
restoreBody(request, body)
var document any
decoder := json.NewDecoder(bytes.NewReader(body))
decoder.UseNumber()
if decoder.Decode(&document) != nil {
return Result{}, nil
}
model := findModel(document)
result := Result{}
for _, policy := range e.current.Load().policies {
if !policyApplies(policy, request.URL.Path, model, apiKeyID) {
continue
}
matched := make([]string, 0)
changed := false
walkText(&document, false, func(value string) string {
for _, rule := range policy.rules {
if rule.expression.MatchString(value) {
if !slices.Contains(matched, rule.name) {
matched = append(matched, rule.name)
}
if policy.Action == "redact" {
next := rule.expression.ReplaceAllString(value, rule.replacement)
changed = changed || next != value
value = next
}
}
}
return value
})
if len(matched) == 0 {
continue
}
result.Matches = append(result.Matches, Match{PolicyID: policy.ID, PolicyName: policy.Name, Action: policy.Action, Rules: matched})
if policy.Action == "block" {
result.Blocked = true
break
}
result.Redacted = result.Redacted || changed
}
if result.Redacted && !result.Blocked {
encoded, err := json.Marshal(document)
if err != nil {
return Result{}, err
}
restoreBody(request, encoded)
}
return result, nil
}
func policyApplies(policy compiledPolicy, path, model, apiKeyID string) bool {
return matches(policy.Paths, path) && matches(policy.Models, model) && matches(policy.APIKeyIDs, apiKeyID)
}
func matches(values []string, candidate string) bool {
if len(values) == 0 {
return true
}
for _, value := range values {
if value == "*" || value == candidate {
return true
}
}
return false
}
var textualFields = map[string]bool{"content": true, "text": true, "input": true, "prompt": true, "instructions": true}
func walkText(value *any, selected bool, transform func(string) string) {
switch current := (*value).(type) {
case map[string]any:
for key, child := range current {
local := child
walkText(&local, selected || textualFields[strings.ToLower(key)], transform)
current[key] = local
}
case []any:
for index, child := range current {
local := child
walkText(&local, selected, transform)
current[index] = local
}
case string:
if selected {
*value = transform(current)
}
}
}
func findModel(document any) string {
if object, ok := document.(map[string]any); ok {
if model, ok := object["model"].(string); ok {
return model
}
}
return ""
}
func restoreBody(request *http.Request, body []byte) {
request.Body = io.NopCloser(bytes.NewReader(body))
request.ContentLength = int64(len(body))
request.GetBody = func() (io.ReadCloser, error) { return io.NopCloser(bytes.NewReader(body)), nil }
}
+44
View File
@@ -0,0 +1,44 @@
package contentpolicy
import (
"io"
"net/http"
"strings"
"testing"
)
func testEngine(t *testing.T, action string) *Engine {
t.Helper()
policy := Policy{ID: "p1", Name: "secret", Action: action, Rules: []Rule{{Name: "token", Pattern: `sk-[A-Za-z0-9]{8,}`, Replacement: "[MASKED]"}}, Enabled: true}
compiled, err := compilePolicy(policy)
if err != nil {
t.Fatal(err)
}
engine := &Engine{}
engine.current.Store(&snapshot{policies: []compiledPolicy{compiled}})
return engine
}
func TestApplyRedactsOnlyTextualPromptFields(t *testing.T) {
request, _ := http.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{"model":"gpt-5","messages":[{"content":"use sk-12345678"}],"tools":[{"description":"keep sk-abcdefgh"}]}`))
result, err := testEngine(t, "redact").Apply(request, 1<<20, "key")
if err != nil || !result.Redacted || result.Blocked {
t.Fatalf("unexpected result %#v: %v", result, err)
}
body, _ := io.ReadAll(request.Body)
text := string(body)
if strings.Contains(text, "sk-12345678") || !strings.Contains(text, "keep sk-abcdefgh") {
t.Fatalf("unexpected redaction: %s", text)
}
}
func TestApplyBlocksWithoutReturningMatchedSecret(t *testing.T) {
request, _ := http.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"input":"sk-12345678"}`))
result, err := testEngine(t, "block").Apply(request, 1<<20, "")
if err != nil || !result.Blocked || len(result.Matches) != 1 {
t.Fatalf("unexpected result %#v: %v", result, err)
}
if strings.Contains(result.Matches[0].Rules[0], "12345678") {
t.Fatal("match metadata leaked secret")
}
}
+156
View File
@@ -0,0 +1,156 @@
package contentpolicy
import (
"context"
"encoding/json"
"errors"
"fmt"
"strings"
platformid "aigateway.local/core/internal/platform/id"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgtype"
"github.com/jackc/pgx/v5/pgxpool"
)
var ErrNotFound = errors.New("content policy not found")
type Store struct{ pool *pgxpool.Pool }
func NewStore(pool *pgxpool.Pool) *Store { return &Store{pool: pool} }
func (s *Store) List(ctx context.Context) ([]Policy, error) {
rows, err := s.pool.Query(ctx, `SELECT id::text,name,description,action,priority,paths,models,api_key_ids::text[],rules,enabled,revision,created_at,updated_at FROM gateway.content_policies ORDER BY priority DESC,name`)
if err != nil {
return nil, err
}
defer rows.Close()
result := []Policy{}
for rows.Next() {
var p Policy
var raw []byte
if err := rows.Scan(&p.ID, &p.Name, &p.Description, &p.Action, &p.Priority, &p.Paths, &p.Models, &p.APIKeyIDs, &raw, &p.Enabled, &p.Revision, &p.CreatedAt, &p.UpdatedAt); err != nil {
return nil, err
}
if err := json.Unmarshal(raw, &p.Rules); err != nil {
return nil, err
}
result = append(result, p)
}
return result, rows.Err()
}
func (s *Store) Save(ctx context.Context, p Policy, actorID string, create bool) (Policy, error) {
if err := Validate(p); err != nil {
return Policy{}, err
}
rules, _ := json.Marshal(p.Rules)
eventID, _ := platformid.NewUUID()
tx, err := s.pool.Begin(ctx)
if err != nil {
return Policy{}, err
}
defer func() { _ = tx.Rollback(ctx) }()
if create {
p.ID, _ = platformid.NewUUID()
err = tx.QueryRow(ctx, `INSERT INTO gateway.content_policies(id,name,description,action,priority,paths,models,api_key_ids,rules,enabled,created_by) VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11) RETURNING revision,created_at,updated_at`, p.ID, p.Name, p.Description, p.Action, p.Priority, p.Paths, p.Models, p.APIKeyIDs, rules, p.Enabled, actorID).Scan(&p.Revision, &p.CreatedAt, &p.UpdatedAt)
} else {
err = tx.QueryRow(ctx, `UPDATE gateway.content_policies SET name=$2,description=$3,action=$4,priority=$5,paths=$6,models=$7,api_key_ids=$8,rules=$9,enabled=$10,revision=revision+1,updated_at=clock_timestamp() WHERE id=$1 RETURNING revision,created_at,updated_at`, p.ID, p.Name, p.Description, p.Action, p.Priority, p.Paths, p.Models, p.APIKeyIDs, rules, p.Enabled).Scan(&p.Revision, &p.CreatedAt, &p.UpdatedAt)
}
if errors.Is(err, pgx.ErrNoRows) {
return Policy{}, ErrNotFound
}
if err != nil {
return Policy{}, err
}
eventType := "content_policy.updated"
if create {
eventType = "content_policy.created"
}
payload, _ := json.Marshal(map[string]any{"content_policy_id": p.ID, "actor_id": actorID, "revision": p.Revision})
_, err = tx.Exec(ctx, `INSERT INTO gateway.outbox_events(event_id,event_type,event_version,aggregate_type,aggregate_id,payload) VALUES($1,$2,1,'content_policy',$3,$4)`, eventID, eventType, p.ID, payload)
if err != nil {
return Policy{}, err
}
return p, tx.Commit(ctx)
}
func (s *Store) Delete(ctx context.Context, id, actorID string) error {
eventID, _ := platformid.NewUUID()
tx, err := s.pool.Begin(ctx)
if err != nil {
return err
}
defer func() { _ = tx.Rollback(ctx) }()
tag, err := tx.Exec(ctx, `DELETE FROM gateway.content_policies WHERE id=$1`, id)
if err != nil {
return err
}
if tag.RowsAffected() == 0 {
return ErrNotFound
}
payload, _ := json.Marshal(map[string]string{"content_policy_id": id, "actor_id": actorID})
if _, err = tx.Exec(ctx, `INSERT INTO gateway.outbox_events(event_id,event_type,event_version,aggregate_type,aggregate_id,payload) VALUES($1,'content_policy.deleted',1,'content_policy',$2,$3)`, eventID, id, payload); err != nil {
return err
}
return tx.Commit(ctx)
}
func Normalize(p *Policy) {
p.Name = strings.TrimSpace(p.Name)
p.Description = strings.TrimSpace(p.Description)
p.Action = strings.ToLower(strings.TrimSpace(p.Action))
p.Paths = normalizeStrings(p.Paths)
p.Models = normalizeStrings(p.Models)
p.APIKeyIDs = normalizeStrings(p.APIKeyIDs)
}
func ValidateInput(p Policy) error {
if p.Name == "" || len(p.Name) > 128 {
return fmt.Errorf("name is required and must not exceed 128 characters")
}
if len(p.Description) > 1000 {
return fmt.Errorf("description is too long")
}
if p.Priority < -100000 || p.Priority > 100000 {
return fmt.Errorf("priority 必须在 -100000 到 100000 之间")
}
if len(p.Paths) > 20 || len(p.Models) > 100 || len(p.APIKeyIDs) > 100 {
return fmt.Errorf("策略范围条目过多")
}
allowedPaths := map[string]bool{"*": true, "/v1/chat/completions": true, "/v1/responses": true, "/v1/embeddings": true, "/v1/messages": true}
for _, value := range p.Paths {
if !allowedPaths[value] {
return fmt.Errorf("不支持的端点范围 %s", value)
}
}
for _, value := range p.Models {
if len(value) > 512 {
return fmt.Errorf("模型名过长")
}
}
for _, value := range p.APIKeyIDs {
var id pgtype.UUID
if id.Scan(value) != nil || !id.Valid {
return fmt.Errorf("API Key ID 格式无效")
}
}
for _, rule := range p.Rules {
if len(rule.Name) > 128 || len(rule.Replacement) > 1024 {
return fmt.Errorf("规则名称或替换文本过长")
}
}
return Validate(p)
}
func normalizeStrings(values []string) []string {
result := make([]string, 0, len(values))
seen := map[string]bool{}
for _, value := range values {
value = strings.TrimSpace(value)
if value != "" && !seen[value] {
seen[value] = true
result = append(result, value)
}
}
return result
}
+154
View File
@@ -0,0 +1,154 @@
package factcheck
import (
"aigateway.local/core/internal/identity"
"aigateway.local/core/internal/platform/apiresponse"
"encoding/json"
"errors"
"net/http"
"strconv"
"strings"
)
type AdminHTTPHandler struct {
service *Service
identity *identity.Service
mux *http.ServeMux
}
func NewAdminHTTPHandler(service *Service, identityService *identity.Service) *AdminHTTPHandler {
h := &AdminHTTPHandler{service: service, identity: identityService, mux: http.NewServeMux()}
h.mux.HandleFunc("GET /api/v1/admin/fact-check/settings", h.getSettings)
h.mux.HandleFunc("PUT /api/v1/admin/fact-check/settings", h.saveSettings)
h.mux.HandleFunc("GET /api/v1/admin/fact-check/policies", h.policies)
h.mux.HandleFunc("POST /api/v1/admin/fact-check/policies", h.savePolicy)
h.mux.HandleFunc("PUT /api/v1/admin/fact-check/policies/{id}", h.savePolicy)
h.mux.HandleFunc("DELETE /api/v1/admin/fact-check/policies/{id}", h.deletePolicy)
h.mux.HandleFunc("GET /api/v1/admin/fact-check/events", h.events)
h.mux.HandleFunc("GET /api/v1/admin/fact-check/events/{id}", h.event)
return h
}
func (h *AdminHTTPHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { h.mux.ServeHTTP(w, r) }
func (h *AdminHTTPHandler) account(w http.ResponseWriter, r *http.Request, manage bool) (identity.Account, bool) {
a, err := h.identity.Authenticate(r.Context(), identity.KindAdmin, r.Header.Get("Authorization"))
if err != nil {
apiresponse.Error(w, 401, "登录状态无效")
return a, false
}
permission := identity.PermissionKnowledgeRead
if manage {
permission = identity.PermissionKnowledgeManage
}
if !identity.HasPermission(a, permission) {
apiresponse.Error(w, 403, "缺少事实核验管理权限")
return a, false
}
return a, true
}
func decodeFact(w http.ResponseWriter, r *http.Request, target any) bool {
d := json.NewDecoder(http.MaxBytesReader(w, r.Body, 1<<20))
d.DisallowUnknownFields()
if err := d.Decode(target); err != nil {
apiresponse.Error(w, 400, "请求格式无效")
return false
}
return true
}
func factError(w http.ResponseWriter, err error) {
if errors.Is(err, ErrNotFound) {
apiresponse.Error(w, 404, "资源不存在")
return
}
apiresponse.Error(w, 400, err.Error())
}
func (h *AdminHTTPHandler) getSettings(w http.ResponseWriter, r *http.Request) {
if _, ok := h.account(w, r, false); !ok {
return
}
x, err := h.service.Settings(r.Context())
if err != nil {
factError(w, err)
return
}
apiresponse.OK(w, x)
}
func (h *AdminHTTPHandler) saveSettings(w http.ResponseWriter, r *http.Request) {
a, ok := h.account(w, r, true)
if !ok {
return
}
var x Settings
if !decodeFact(w, r, &x) {
return
}
x, err := h.service.SaveSettings(r.Context(), x, a.ID)
if err != nil {
factError(w, err)
return
}
apiresponse.OK(w, x)
}
func (h *AdminHTTPHandler) policies(w http.ResponseWriter, r *http.Request) {
if _, ok := h.account(w, r, false); !ok {
return
}
x, err := h.service.Policies(r.Context())
if err != nil {
factError(w, err)
return
}
apiresponse.OK(w, map[string]any{"items": x})
}
func (h *AdminHTTPHandler) savePolicy(w http.ResponseWriter, r *http.Request) {
a, ok := h.account(w, r, true)
if !ok {
return
}
var x Policy
if !decodeFact(w, r, &x) {
return
}
if id := r.PathValue("id"); id != "" {
x.ID = id
}
x, err := h.service.SavePolicy(r.Context(), x, a.ID)
if err != nil {
factError(w, err)
return
}
apiresponse.OK(w, x)
}
func (h *AdminHTTPHandler) deletePolicy(w http.ResponseWriter, r *http.Request) {
if _, ok := h.account(w, r, true); !ok {
return
}
if err := h.service.DeletePolicy(r.Context(), r.PathValue("id")); err != nil {
factError(w, err)
return
}
apiresponse.OK(w, map[string]bool{"deleted": true})
}
func (h *AdminHTTPHandler) events(w http.ResponseWriter, r *http.Request) {
if _, ok := h.account(w, r, false); !ok {
return
}
limit, _ := strconv.Atoi(r.URL.Query().Get("limit"))
verdict := strings.TrimSpace(r.URL.Query().Get("verdict"))
x, err := h.service.Events(r.Context(), verdict, limit)
if err != nil {
factError(w, err)
return
}
apiresponse.OK(w, map[string]any{"items": x})
}
func (h *AdminHTTPHandler) event(w http.ResponseWriter, r *http.Request) {
if _, ok := h.account(w, r, false); !ok {
return
}
x, err := h.service.Event(r.Context(), r.PathValue("id"))
if err != nil {
factError(w, err)
return
}
apiresponse.OK(w, x)
}
+298
View File
@@ -0,0 +1,298 @@
package factcheck
import (
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"log/slog"
"strings"
"time"
platformid "aigateway.local/core/internal/platform/id"
"github.com/jackc/pgx/v5/pgxpool"
)
// EvidenceHit is a single retrieved knowledge-base excerpt supplied to the
// verifier. The title/content shape mirrors the workbench retriever output.
type EvidenceHit struct {
DocumentTitle string `json:"document_title"`
Content string `json:"content"`
}
// EvidenceRetriever fetches supporting excerpts from a knowledge base. The
// workbench PostgreSQLRetriever satisfies it (through a small adapter), so
// fact-check reuses the exact retrieval path application prompts already use.
type EvidenceRetriever interface {
Search(context.Context, string, string, int) ([]EvidenceHit, error)
}
// Verifier performs the model-backed fact-check call and returns the raw model
// output text. The engine is responsible for parsing it into claims and a
// verdict. The workbench runtime implements it by routing a non-streaming chat
// completion through the same governed gateway, reusing the caller's own
// credential headers.
type Verifier interface {
Verify(context.Context, string, string, string, time.Duration) (string, error)
}
// VerifierFunc adapts a function to the Verifier interface.
type VerifierFunc func(context.Context, string, string, string, time.Duration) (string, error)
func (f VerifierFunc) Verify(ctx context.Context, model, system, user string, timeout time.Duration) (string, error) {
return f(ctx, model, system, user, timeout)
}
// Engine executes fact-check policies against assistant answers and records the
// outcome in fact_check_events. It is deliberately side-effect safe: Check
// never returns an error a caller must propagate — callers treat any failure as
// "fact-check skipped" and keep serving the chat.
type Engine struct {
pool *pgxpool.Pool
retriever EvidenceRetriever
logger *slog.Logger
}
func NewEngine(pool *pgxpool.Pool, retriever EvidenceRetriever, logger *slog.Logger) *Engine {
return &Engine{pool: pool, retriever: retriever, logger: logger}
}
// Check verifies one assistant answer against the configured knowledge bases
// and persists a fact_check_events row. The returned Event has a zero ID when
// fact-checking is not configured or no policy applies; callers should skip
// quietly in that case.
func (e *Engine) Check(ctx context.Context, requestID, question, answer string, verifier Verifier) (Event, error) {
if e == nil || e.retriever == nil || verifier == nil || strings.TrimSpace(answer) == "" {
return Event{}, nil
}
settings, err := e.checkSettings(ctx)
if err != nil {
return Event{}, fmt.Errorf("read fact-check settings: %w", err)
}
if strings.TrimSpace(settings.Model) == "" {
return Event{}, nil // not configured; skip without noise
}
policy, err := e.enabledPolicy(ctx)
if err != nil {
return Event{}, err
}
if policy.ID == "" {
return Event{}, nil
}
started := time.Now()
evidence := e.gatherEvidence(ctx, policy, question)
if len(evidence) == 0 {
// Nothing to check against; record an uncertain event so the admin can
// see retrieval produced no evidence rather than silently passing.
return e.record(ctx, Event{
PolicyID: &policy.ID, RequestID: requestID, Model: settings.Model,
Mode: policy.Mode, Action: policy.Action, Verdict: "uncertain",
LatencyMS: int(time.Since(started).Milliseconds()), Question: question, Answer: answer,
Claims: json.RawMessage("[]"), Evidence: evidenceJSON(evidence),
})
}
system, user := buildVerifyPrompt(question, answer, evidence, policy.MaxClaims)
raw, err := verifier.Verify(ctx, settings.Model, system, user, time.Duration(settings.TimeoutSeconds)*time.Second)
if err != nil {
return e.record(ctx, Event{
PolicyID: &policy.ID, RequestID: requestID, Model: settings.Model,
Mode: policy.Mode, Action: policy.Action, Verdict: "error",
LatencyMS: int(time.Since(started).Milliseconds()), Question: question, Answer: answer,
Claims: json.RawMessage("[]"), Evidence: evidenceJSON(evidence), Error: err.Error(),
})
}
verdict, score, claims := parseVerdict(raw, policy.SupportThreshold)
return e.record(ctx, Event{
PolicyID: &policy.ID, RequestID: requestID, Model: settings.Model,
Mode: policy.Mode, Action: policy.Action, Verdict: verdict, SupportScore: score,
LatencyMS: int(time.Since(started).Milliseconds()), Question: question, Answer: answer,
Claims: claimsJSON(claims), Evidence: evidenceJSON(evidence),
})
}
func (e *Engine) checkSettings(ctx context.Context) (Settings, error) {
var x Settings
err := e.pool.QueryRow(ctx, `SELECT f.provider_id::text,coalesce(p.code,''),f.model,f.timeout_seconds,f.updated_at FROM gateway.fact_check_settings f LEFT JOIN gateway.providers p ON p.id=f.provider_id WHERE singleton`).Scan(&x.ProviderID, &x.ProviderCode, &x.Model, &x.TimeoutSeconds, &x.UpdatedAt)
return x, err
}
func (e *Engine) enabledPolicy(ctx context.Context) (Policy, error) {
policy, err := scanPolicy(e.pool.QueryRow(ctx, policySelect+` WHERE enabled ORDER BY scope LIMIT 1`))
if errors.Is(err, ErrNotFound) {
return Policy{}, nil
}
return policy, err
}
// gatherEvidence retrieves up to a bounded number of excerpts across the
// policy's knowledge bases. Retrieval failures never abort the check — a
// broken knowledge base just contributes no evidence.
func (e *Engine) gatherEvidence(ctx context.Context, policy Policy, question string) []EvidenceHit {
maxHits := policy.MaxClaims * 2
if maxHits < 4 {
maxHits = 4
}
if maxHits > 20 {
maxHits = 20
}
evidence := []EvidenceHit{}
for _, kbID := range policy.KnowledgeBaseIDs {
hits, err := e.retriever.Search(ctx, kbID, question, policy.TopK)
if err != nil {
if e.logger != nil {
e.logger.Warn("fact-check evidence retrieval failed", "knowledge_base_id", kbID, "error", err)
}
continue
}
for _, hit := range hits {
if len(evidence) >= maxHits {
break
}
evidence = append(evidence, EvidenceHit{DocumentTitle: hit.DocumentTitle, Content: hit.Content})
}
if len(evidence) >= maxHits {
break
}
}
return evidence
}
func (e *Engine) record(ctx context.Context, event Event) (Event, error) {
if event.ID == "" {
id, err := platformid.NewUUID()
if err != nil {
return event, err
}
event.ID = id
}
_, err := e.pool.Exec(ctx, `INSERT INTO gateway.fact_check_events(id,policy_id,request_id,model,mode,action,verdict,support_score,latency_ms,question,answer,claims,evidence,error) VALUES($1,nullif($2,'')::uuid,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12::jsonb,$13::jsonb,$14)`,
event.ID, ptrString(event.PolicyID), event.RequestID, event.Model, event.Mode, event.Action, event.Verdict, event.SupportScore, event.LatencyMS, event.Question, event.Answer, event.Claims, event.Evidence, event.Error)
return event, err
}
func ptrString(value *string) string {
if value == nil {
return ""
}
return *value
}
type verifyClaim struct {
Claim string `json:"claim"`
Verdict string `json:"verdict"`
EvidenceIndex []int `json:"evidence_index"`
}
type verifyResult struct {
Verdict string `json:"verdict"`
SupportScore float64 `json:"support_score"`
Claims []verifyClaim `json:"claims"`
}
func buildVerifyPrompt(question, answer string, evidence []EvidenceHit, maxClaims int) (string, string) {
if maxClaims < 1 {
maxClaims = 8
}
var buffer bytes.Buffer
buffer.WriteString("问题:\n")
buffer.WriteString(question)
buffer.WriteString("\n\n回答:\n")
buffer.WriteString(answer)
buffer.WriteString("\n\n参考资料:\n")
for i, hit := range evidence {
fmt.Fprintf(&buffer, "[资料%d] %s\n%s\n\n", i, hit.DocumentTitle, hit.Content)
}
system := `你是企业知识库事实核查引擎。你需要逐条判断"回答"中的关键陈述是否能被"参考资料"支持。
输出必须是严格 JSON,不要输出任何其他内容,格式:
{"verdict":"supported|unsupported|uncertain","support_score":0-100,"claims":[{"claim":"...","verdict":"supported|unsupported|uncertain","evidence_index":[0,1]}]}
- verdict:所有关键陈述均被参考资料支持→supported;存在明确被资料否定或资料完全无法支撑的关键陈述→unsupported;资料不足无法判断→uncertain。
- support_score:被支持的陈述占比(0-100)。
- claims:从回答中提取的关键陈述,最多 ` + fmt.Sprint(maxClaims) + ` 条。evidence_index 列出支撑该陈述的资料编号(从0开始);无支撑填[]。
- 只依据参考资料判断,不要使用你自己的世界知识。
- 回答为空或不包含可核查陈述时,verdict 输出 uncertainclaims 输出 []。`
return system, buffer.String()
}
func parseVerdict(raw string, threshold int) (string, *int, []verifyClaim) {
var result verifyResult
if extracted := extractJSON(raw); json.Unmarshal([]byte(extracted), &result) != nil {
return "error", nil, []verifyClaim{}
}
claims := result.Claims
if claims == nil {
claims = []verifyClaim{}
}
verdict := strings.ToLower(strings.TrimSpace(result.Verdict))
var score *int
if result.SupportScore > 0 {
value := int(result.SupportScore)
if value > 100 {
value = 100
}
score = &value
}
total := len(claims)
supported := 0
for _, c := range claims {
switch strings.ToLower(strings.TrimSpace(c.Verdict)) {
case "supported":
supported++
case "unsupported", "uncertain":
default:
c.Verdict = "uncertain"
}
}
if total > 0 {
percent := supported * 100 / total
if score == nil {
score = &percent
}
if verdict != "supported" && verdict != "unsupported" && verdict != "uncertain" {
switch {
case percent >= threshold:
verdict = "supported"
case supported > 0:
verdict = "uncertain"
default:
verdict = "unsupported"
}
}
} else if verdict != "supported" && verdict != "unsupported" && verdict != "uncertain" {
verdict = "uncertain"
}
return verdict, score, claims
}
// extractJSON returns the text between the first '{' and the last '}', stripping
// markdown code fences models sometimes wrap around their JSON output.
func extractJSON(raw string) string {
start := strings.Index(raw, "{")
end := strings.LastIndex(raw, "}")
if start < 0 || end < start {
return ""
}
return raw[start : end+1]
}
func claimsJSON(claims []verifyClaim) json.RawMessage {
if claims == nil {
return json.RawMessage("[]")
}
encoded, err := json.Marshal(claims)
if err != nil {
return json.RawMessage("[]")
}
return encoded
}
func evidenceJSON(evidence []EvidenceHit) json.RawMessage {
if evidence == nil {
return json.RawMessage("[]")
}
encoded, err := json.Marshal(evidence)
if err != nil {
return json.RawMessage("[]")
}
return encoded
}
+56
View File
@@ -0,0 +1,56 @@
package factcheck
import "testing"
func TestParseVerdictSupported(t *testing.T) {
raw := "```json\n{\"verdict\":\"supported\",\"support_score\":88,\"claims\":[{\"claim\":\"a\",\"verdict\":\"supported\",\"evidence_index\":[0]}]}\n```"
verdict, score, claims := parseVerdict(raw, 70)
if verdict != "supported" {
t.Fatalf("expected supported, got %q", verdict)
}
if score == nil || *score != 88 {
t.Fatalf("expected score 88, got %v", score)
}
if len(claims) != 1 {
t.Fatalf("expected 1 claim, got %d", len(claims))
}
}
func TestParseVerdictDerivesFromClaims(t *testing.T) {
// Model returns no top-level verdict; the engine must derive it from claims.
raw := `{"claims":[{"claim":"x","verdict":"unsupported","evidence_index":[]}]}`
verdict, _, _ := parseVerdict(raw, 70)
if verdict != "unsupported" {
t.Fatalf("expected unsupported, got %q", verdict)
}
}
func TestParseVerdictEmptyClaimsUncertain(t *testing.T) {
raw := `{"verdict":"","support_score":0,"claims":[]}`
verdict, _, _ := parseVerdict(raw, 70)
if verdict != "uncertain" {
t.Fatalf("expected uncertain, got %q", verdict)
}
}
func TestParseVerdictGarbageIsError(t *testing.T) {
verdict, _, _ := parseVerdict("this is not json at all", 70)
if verdict != "error" {
t.Fatalf("expected error, got %q", verdict)
}
}
func TestParseVerdictClampsScore(t *testing.T) {
raw := `{"verdict":"supported","support_score":150,"claims":[{"claim":"a","verdict":"supported","evidence_index":[0]}]}`
_, score, _ := parseVerdict(raw, 70)
if score == nil || *score > 100 {
t.Fatalf("score should be clamped to 100, got %v", score)
}
}
func TestExtractJSONStripsFences(t *testing.T) {
raw := "以下是结果:\n```json\n{\"a\":1}\n```\n完"
if got := extractJSON(raw); got != `{"a":1}` {
t.Fatalf("expected {\"a\":1}, got %q", got)
}
}
+235
View File
@@ -0,0 +1,235 @@
package factcheck
import (
"context"
"encoding/json"
"errors"
"strconv"
"strings"
"time"
platformid "aigateway.local/core/internal/platform/id"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgxpool"
)
var ErrNotFound = errors.New("fact-check resource not found")
type Service struct{ pool *pgxpool.Pool }
func NewService(pool *pgxpool.Pool) *Service { return &Service{pool: pool} }
type Settings struct {
ProviderID *string `json:"provider_id"`
ProviderCode string `json:"provider_code"`
Model string `json:"model"`
TimeoutSeconds int `json:"timeout_seconds"`
UpdatedAt time.Time `json:"updated_at"`
}
func (s *Service) Settings(ctx context.Context) (Settings, error) {
var x Settings
err := s.pool.QueryRow(ctx, `SELECT f.provider_id::text,coalesce(p.code,''),f.model,f.timeout_seconds,f.updated_at FROM gateway.fact_check_settings f LEFT JOIN gateway.providers p ON p.id=f.provider_id WHERE singleton`).Scan(&x.ProviderID, &x.ProviderCode, &x.Model, &x.TimeoutSeconds, &x.UpdatedAt)
return x, err
}
func (s *Service) SaveSettings(ctx context.Context, x Settings, actor string) (Settings, error) {
x.Model = strings.TrimSpace(x.Model)
if x.TimeoutSeconds < 3 || x.TimeoutSeconds > 60 {
return Settings{}, errors.New("超时必须在 3-60 秒之间")
}
if x.ProviderID != nil && strings.TrimSpace(*x.ProviderID) == "" {
x.ProviderID = nil
}
if x.ProviderID != nil && x.Model == "" {
return Settings{}, errors.New("配置供应商时模型不能为空")
}
_, err := s.pool.Exec(ctx, `UPDATE gateway.fact_check_settings SET provider_id=$1,model=$2,timeout_seconds=$3,updated_by=$4,updated_at=clock_timestamp() WHERE singleton`, x.ProviderID, x.Model, x.TimeoutSeconds, actor)
if err != nil {
return Settings{}, err
}
return s.Settings(ctx)
}
type Policy struct {
ID string `json:"id"`
Scope string `json:"scope"`
Enabled bool `json:"enabled"`
Mode string `json:"mode"`
Action string `json:"action"`
KnowledgeBaseIDs []string `json:"knowledge_base_ids"`
SupportThreshold int `json:"support_threshold"`
EvidenceThreshold float64 `json:"evidence_threshold"`
TopK int `json:"top_k"`
MaxClaims int `json:"max_claims"`
Revision int64 `json:"revision"`
CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"`
}
func validatePolicy(x *Policy) error {
x.Scope = strings.ToLower(strings.TrimSpace(x.Scope))
if x.Scope == "" {
x.Scope = "global"
}
if x.Mode == "" {
x.Mode = "async"
}
if x.Action == "" {
x.Action = "observe"
}
if x.Mode != "async" && x.Mode != "sync" {
return errors.New("模式无效")
}
if x.Action != "observe" && x.Action != "annotate" && x.Action != "block" {
return errors.New("处置动作无效")
}
if x.Mode == "async" && x.Action != "observe" {
return errors.New("异步模式只能观察")
}
if x.Enabled && len(x.KnowledgeBaseIDs) == 0 {
return errors.New("启用策略前至少选择一个知识库")
}
if x.TopK == 0 {
x.TopK = 4
}
if x.MaxClaims == 0 {
x.MaxClaims = 8
}
if x.SupportThreshold == 0 {
x.SupportThreshold = 70
}
if x.EvidenceThreshold == 0 {
x.EvidenceThreshold = .35
}
if x.TopK < 1 || x.TopK > 10 || x.MaxClaims < 1 || x.MaxClaims > 20 || x.SupportThreshold < 0 || x.SupportThreshold > 100 || x.EvidenceThreshold < 0 || x.EvidenceThreshold > 1 {
return errors.New("事实核验阈值无效")
}
return nil
}
const policySelect = `SELECT id::text,scope,enabled,mode,action,knowledge_base_ids::text[],support_threshold,evidence_threshold,top_k,max_claims,revision,created_at,updated_at FROM gateway.fact_check_policies`
func scanPolicy(row pgx.Row) (Policy, error) {
var x Policy
err := row.Scan(&x.ID, &x.Scope, &x.Enabled, &x.Mode, &x.Action, &x.KnowledgeBaseIDs, &x.SupportThreshold, &x.EvidenceThreshold, &x.TopK, &x.MaxClaims, &x.Revision, &x.CreatedAt, &x.UpdatedAt)
if errors.Is(err, pgx.ErrNoRows) {
err = ErrNotFound
}
return x, err
}
func (s *Service) Policies(ctx context.Context) ([]Policy, error) {
rows, err := s.pool.Query(ctx, policySelect+` ORDER BY scope`)
if err != nil {
return nil, err
}
defer rows.Close()
items := []Policy{}
for rows.Next() {
x, err := scanPolicy(rows)
if err != nil {
return nil, err
}
items = append(items, x)
}
return items, rows.Err()
}
func (s *Service) SavePolicy(ctx context.Context, x Policy, actor string) (Policy, error) {
if err := validatePolicy(&x); err != nil {
return Policy{}, err
}
if x.KnowledgeBaseIDs == nil {
x.KnowledgeBaseIDs = []string{}
}
if x.ID == "" {
x.ID, _ = platformid.NewUUID()
_, err := s.pool.Exec(ctx, `INSERT INTO gateway.fact_check_policies(id,scope,enabled,mode,action,knowledge_base_ids,support_threshold,evidence_threshold,top_k,max_claims,created_by) VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11)`, x.ID, x.Scope, x.Enabled, x.Mode, x.Action, x.KnowledgeBaseIDs, x.SupportThreshold, x.EvidenceThreshold, x.TopK, x.MaxClaims, actor)
if err != nil {
return Policy{}, err
}
} else {
tag, err := s.pool.Exec(ctx, `UPDATE gateway.fact_check_policies SET scope=$2,enabled=$3,mode=$4,action=$5,knowledge_base_ids=$6,support_threshold=$7,evidence_threshold=$8,top_k=$9,max_claims=$10,revision=revision+1,updated_at=clock_timestamp() WHERE id=$1`, x.ID, x.Scope, x.Enabled, x.Mode, x.Action, x.KnowledgeBaseIDs, x.SupportThreshold, x.EvidenceThreshold, x.TopK, x.MaxClaims)
if err != nil {
return Policy{}, err
}
if tag.RowsAffected() == 0 {
return Policy{}, ErrNotFound
}
}
return scanPolicy(s.pool.QueryRow(ctx, policySelect+` WHERE id=$1`, x.ID))
}
func (s *Service) DeletePolicy(ctx context.Context, id string) error {
tag, err := s.pool.Exec(ctx, `DELETE FROM gateway.fact_check_policies WHERE id=$1`, id)
if err != nil {
return err
}
if tag.RowsAffected() == 0 {
return ErrNotFound
}
return nil
}
type Event struct {
ID string `json:"id"`
PolicyID *string `json:"policy_id"`
RequestID string `json:"request_id"`
Model string `json:"model"`
Mode string `json:"mode"`
Action string `json:"action"`
Verdict string `json:"verdict"`
SupportScore *int `json:"support_score"`
LatencyMS int `json:"latency_ms"`
Question string `json:"question,omitempty"`
Answer string `json:"answer,omitempty"`
Claims json.RawMessage `json:"claims,omitempty"`
Evidence json.RawMessage `json:"evidence,omitempty"`
Error string `json:"error"`
CreatedAt time.Time `json:"created_at"`
}
const eventSelect = `SELECT id::text,policy_id::text,request_id,model,mode,action,verdict,support_score,latency_ms,question,answer,claims,evidence,error,created_at FROM gateway.fact_check_events`
func scanEvent(row pgx.Row) (Event, error) {
var x Event
err := row.Scan(&x.ID, &x.PolicyID, &x.RequestID, &x.Model, &x.Mode, &x.Action, &x.Verdict, &x.SupportScore, &x.LatencyMS, &x.Question, &x.Answer, &x.Claims, &x.Evidence, &x.Error, &x.CreatedAt)
if errors.Is(err, pgx.ErrNoRows) {
err = ErrNotFound
}
return x, err
}
func (s *Service) Events(ctx context.Context, verdict string, limit int) ([]Event, error) {
if limit < 1 {
limit = 50
}
if limit > 200 {
limit = 200
}
query := eventSelect
args := []any{}
if verdict != "" {
args = append(args, verdict)
query += ` WHERE verdict=$1`
}
args = append(args, limit)
query += ` ORDER BY created_at DESC LIMIT $` + strconv.Itoa(len(args))
rows, err := s.pool.Query(ctx, query, args...)
if err != nil {
return nil, err
}
defer rows.Close()
items := []Event{}
for rows.Next() {
x, err := scanEvent(rows)
if err != nil {
return nil, err
}
x.Question = ""
x.Answer = ""
x.Claims = nil
x.Evidence = nil
items = append(items, x)
}
return items, rows.Err()
}
func (s *Service) Event(ctx context.Context, id string) (Event, error) {
return scanEvent(s.pool.QueryRow(ctx, eventSelect+` WHERE id=$1`, id))
}
@@ -0,0 +1,100 @@
package factcheck
import (
"context"
"encoding/json"
"os"
"testing"
"aigateway.local/core/internal/platform/config"
"aigateway.local/core/internal/platform/database"
)
func TestFactCheckPostgreSQLLifecycle(t *testing.T) {
databaseURL := os.Getenv("FACTCHECK_TEST_DATABASE_URL")
if databaseURL == "" {
t.Skip("FACTCHECK_TEST_DATABASE_URL is not set")
}
ctx := context.Background()
pool, err := database.Open(ctx, config.Database{URL: databaseURL, MaxConns: 4})
if err != nil {
t.Fatal(err)
}
defer pool.Close()
const actorID = "33333333-3333-4333-8333-333333333333"
const knowledgeBaseID = "44444444-4444-4444-8444-444444444444"
const eventID = "55555555-5555-4555-8555-555555555555"
_, err = pool.Exec(ctx, `INSERT INTO gateway.admin_accounts(id,username,password_hash,role)
VALUES($1,'factcheck-test-admin','test','superadmin') ON CONFLICT(id) DO NOTHING`, actorID)
if err != nil {
t.Fatal(err)
}
cleanup := func() {
_, _ = pool.Exec(ctx, `DELETE FROM gateway.fact_check_events WHERE id=$1`, eventID)
_, _ = pool.Exec(ctx, `DELETE FROM gateway.fact_check_policies WHERE scope='global'`)
_, _ = pool.Exec(ctx, `DELETE FROM gateway.knowledge_bases WHERE id=$1`, knowledgeBaseID)
_, _ = pool.Exec(ctx, `UPDATE gateway.fact_check_settings SET provider_id=NULL,model='',updated_by=NULL`)
_, _ = pool.Exec(ctx, `DELETE FROM gateway.admin_accounts WHERE id=$1`, actorID)
}
cleanup()
_, err = pool.Exec(ctx, `INSERT INTO gateway.admin_accounts(id,username,password_hash,role)
VALUES($1,'factcheck-test-admin','test','superadmin') ON CONFLICT(id) DO NOTHING`, actorID)
if err != nil {
t.Fatal(err)
}
defer cleanup()
_, err = pool.Exec(ctx, `INSERT INTO gateway.knowledge_bases(id,name,retrieval_mode,chunk_size,chunk_overlap,created_by)
VALUES($1,'factcheck-test-kb','postgres_fts',800,100,$2)`, knowledgeBaseID, actorID)
if err != nil {
t.Fatal(err)
}
service := NewService(pool)
settings, err := service.SaveSettings(ctx, Settings{TimeoutSeconds: 17}, actorID)
if err != nil || settings.TimeoutSeconds != 17 || settings.ProviderID != nil {
t.Fatalf("settings=%#v err=%v", settings, err)
}
policy, err := service.SavePolicy(ctx, Policy{
Scope: "global",
Enabled: true,
Mode: "sync",
Action: "annotate",
KnowledgeBaseIDs: []string{knowledgeBaseID},
}, actorID)
if err != nil {
t.Fatal(err)
}
if policy.Revision != 1 || policy.TopK != 4 || policy.SupportThreshold != 70 {
t.Fatalf("unexpected policy defaults: %#v", policy)
}
policy.Action = "block"
updated, err := service.SavePolicy(ctx, policy, actorID)
if err != nil || updated.Revision != 2 || updated.Action != "block" {
t.Fatalf("updated=%#v err=%v", updated, err)
}
claims := json.RawMessage(`[{"text":"Go gateway"}]`)
evidence := json.RawMessage(`[{"source":"kb"}]`)
_, err = pool.Exec(ctx, `INSERT INTO gateway.fact_check_events
(id,policy_id,request_id,model,mode,action,verdict,support_score,latency_ms,question,answer,claims,evidence)
VALUES($1,$2,'factcheck-request','test-model','sync','block','supported',92,12,'question','answer',$3,$4)`,
eventID, policy.ID, claims, evidence)
if err != nil {
t.Fatal(err)
}
events, err := service.Events(ctx, "supported", 10)
if err != nil || len(events) != 1 {
t.Fatalf("events=%#v err=%v", events, err)
}
if events[0].Question != "" || events[0].Answer != "" || events[0].Claims != nil {
t.Fatalf("event list leaked payload: %#v", events[0])
}
detail, err := service.Event(ctx, eventID)
if err != nil || detail.Question != "question" || len(detail.Evidence) == 0 {
t.Fatalf("detail=%#v err=%v", detail, err)
}
if err = service.DeletePolicy(ctx, policy.ID); err != nil {
t.Fatal(err)
}
}
+139
View File
@@ -0,0 +1,139 @@
package gateway
import (
"context"
"errors"
"fmt"
"strconv"
"time"
"aigateway.local/core/internal/apikey"
"github.com/redis/go-redis/v9"
)
var ErrAdmissionUnavailable = errors.New("admission control unavailable")
type AdmissionReason string
const (
AdmissionAllowed AdmissionReason = ""
AdmissionRateLimit AdmissionReason = "rate_limit"
AdmissionMonthlyQuota AdmissionReason = "monthly_quota"
)
type AdmissionDecision struct {
Allowed bool
Reason AdmissionReason
Limit int64
Remaining int64
ResetAt time.Time
RetryAfter time.Duration
}
type AdmissionController interface {
Allow(context.Context, apikey.Principal, time.Time) (AdmissionDecision, error)
}
type RedisAdmissionController struct {
client *redis.Client
script *redis.Script
}
func NewRedisAdmissionController(client *redis.Client) *RedisAdmissionController {
return &RedisAdmissionController{client: client, script: redis.NewScript(admissionScript)}
}
func (c *RedisAdmissionController) Allow(ctx context.Context, principal apikey.Principal, now time.Time) (AdmissionDecision, error) {
if principal.RequestsPerMinute == 0 && principal.MonthlyRequestQuota == 0 {
return AdmissionDecision{Allowed: true}, nil
}
if c == nil || c.client == nil || principal.APIKeyID == "" {
return AdmissionDecision{}, ErrAdmissionUnavailable
}
now = now.UTC()
minuteReset := now.Truncate(time.Minute).Add(time.Minute)
monthReset := time.Date(now.Year(), now.Month()+1, 1, 0, 0, 0, 0, time.UTC)
minuteKey := fmt.Sprintf("gateway:limit:api-key:%s:minute:%d", principal.APIKeyID, now.Unix()/60)
monthKey := fmt.Sprintf("gateway:limit:api-key:%s:month:%s", principal.APIKeyID, now.Format("200601"))
result, err := c.script.Run(ctx, c.client, []string{minuteKey, monthKey},
principal.RequestsPerMinute, principal.MonthlyRequestQuota,
int64(minuteReset.Sub(now).Seconds())+2, int64(monthReset.Sub(now).Seconds())+86400,
).Slice()
if err != nil {
return AdmissionDecision{}, fmt.Errorf("%w: %v", ErrAdmissionUnavailable, err)
}
if len(result) != 3 {
return AdmissionDecision{}, ErrAdmissionUnavailable
}
code, err := redisInteger(result[0])
if err != nil {
return AdmissionDecision{}, ErrAdmissionUnavailable
}
minuteCount, err := redisInteger(result[1])
if err != nil {
return AdmissionDecision{}, ErrAdmissionUnavailable
}
monthCount, err := redisInteger(result[2])
if err != nil {
return AdmissionDecision{}, ErrAdmissionUnavailable
}
switch code {
case 0:
decision := AdmissionDecision{Allowed: true}
if principal.RequestsPerMinute > 0 {
decision.Limit = int64(principal.RequestsPerMinute)
decision.Remaining = max(decision.Limit-minuteCount, 0)
decision.ResetAt = minuteReset
}
return decision, nil
case 1:
return AdmissionDecision{
Allowed: false, Reason: AdmissionRateLimit, Limit: int64(principal.RequestsPerMinute),
Remaining: 0, ResetAt: minuteReset, RetryAfter: minuteReset.Sub(now),
}, nil
case 2:
return AdmissionDecision{
Allowed: false, Reason: AdmissionMonthlyQuota, Limit: principal.MonthlyRequestQuota,
Remaining: max(principal.MonthlyRequestQuota-monthCount, 0), ResetAt: monthReset, RetryAfter: monthReset.Sub(now),
}, nil
default:
return AdmissionDecision{}, ErrAdmissionUnavailable
}
}
func redisInteger(value any) (int64, error) {
switch number := value.(type) {
case int64:
return number, nil
case string:
return strconv.ParseInt(number, 10, 64)
case []byte:
return strconv.ParseInt(string(number), 10, 64)
default:
return 0, fmt.Errorf("unexpected redis integer %T", value)
}
}
const admissionScript = `
local rpm = tonumber(ARGV[1])
local monthly_quota = tonumber(ARGV[2])
local minute_count = tonumber(redis.call('GET', KEYS[1]) or '0')
local month_count = tonumber(redis.call('GET', KEYS[2]) or '0')
if monthly_quota > 0 and month_count >= monthly_quota then
return {2, minute_count, month_count}
end
if rpm > 0 and minute_count >= rpm then
return {1, minute_count, month_count}
end
if rpm > 0 then
minute_count = redis.call('INCR', KEYS[1])
if minute_count == 1 then redis.call('EXPIRE', KEYS[1], tonumber(ARGV[3])) end
end
if monthly_quota > 0 then
month_count = redis.call('INCR', KEYS[2])
if month_count == 1 then redis.call('EXPIRE', KEYS[2], tonumber(ARGV[4])) end
end
return {0, minute_count, month_count}
`
+30
View File
@@ -0,0 +1,30 @@
package gateway
import (
"testing"
"time"
)
func TestRedisInteger(t *testing.T) {
for _, value := range []any{int64(7), "7", []byte("7")} {
got, err := redisInteger(value)
if err != nil || got != 7 {
t.Fatalf("redisInteger(%T) = %d, %v", value, got, err)
}
}
if _, err := redisInteger(7.0); err == nil {
t.Fatal("unexpected redis number type accepted")
}
}
func TestAdmissionResetBoundaries(t *testing.T) {
now := time.Date(2026, time.August, 10, 13, 25, 40, 0, time.UTC)
minuteReset := now.Truncate(time.Minute).Add(time.Minute)
monthReset := time.Date(now.Year(), now.Month()+1, 1, 0, 0, 0, 0, time.UTC)
if minuteReset != time.Date(2026, time.August, 10, 13, 26, 0, 0, time.UTC) {
t.Fatal("minute boundary is incorrect")
}
if monthReset != time.Date(2026, time.September, 1, 0, 0, 0, 0, time.UTC) {
t.Fatal("month boundary is incorrect")
}
}
+221
View File
@@ -0,0 +1,221 @@
package gateway
import (
"bufio"
"encoding/json"
"io"
"net"
"net/http"
"sync"
"time"
"aigateway.local/core/internal/apikey"
auditpkg "aigateway.local/core/internal/audit"
"aigateway.local/core/internal/contentpolicy"
"aigateway.local/core/internal/pricing"
)
const auditRequestCaptureBytes = 64 << 10
type AuditRecorder interface {
Record(auditpkg.Event) bool
}
type auditSpan struct {
recorder AuditRecorder
started time.Time
event auditpkg.Event
mu sync.Mutex
capture *captureReadCloser
target string
pricing *pricing.Service
policies []contentpolicy.Match
redacted bool
}
func newAuditSpan(recorder AuditRecorder, principal apikey.Principal, request *http.Request, started time.Time) *auditSpan {
if recorder == nil {
return nil
}
event := auditpkg.Event{TenantID: principal.TenantID, RequestID: RequestID(request.Context()), Protocol: request.URL.Path, RecordedAt: time.Now().UTC()}
if principal.APIKeyID != "" {
id := principal.APIKeyID
event.APIKeyID = &id
}
return &auditSpan{recorder: recorder, started: started, event: event}
}
func (s *auditSpan) setModel(model string) {
if s != nil && model != "" {
s.mu.Lock()
s.event.Model = model
s.mu.Unlock()
}
}
func (s *auditSpan) setRoute(providerCode, targetModel string) {
if s == nil {
return
}
s.mu.Lock()
s.event.ProviderCode = providerCode
s.target = targetModel
s.mu.Unlock()
}
func (s *auditSpan) setUsage(usage TokenUsage) {
if s == nil {
return
}
s.mu.Lock()
s.event.PromptTokens = usage.Input
s.event.CompletionTokens = usage.Output
s.calculateCostLocked()
s.mu.Unlock()
}
func (s *auditSpan) setContentPolicy(result contentpolicy.Result) {
if s == nil {
return
}
s.mu.Lock()
s.policies = append(s.policies, result.Matches...)
s.redacted = result.Redacted
s.mu.Unlock()
}
func (s *auditSpan) captureBody(body io.ReadCloser) io.ReadCloser {
if s == nil || body == nil || body == http.NoBody {
return body
}
capture := &captureReadCloser{ReadCloser: body, limit: auditRequestCaptureBytes}
s.capture = capture
return capture
}
func (s *auditSpan) finish(status int) {
if s == nil {
return
}
s.mu.Lock()
if s.event.Model == "" && s.capture != nil {
var payload map[string]json.RawMessage
if json.Unmarshal(s.capture.buffer, &payload) == nil {
_ = json.Unmarshal(payload["model"], &s.event.Model)
}
}
s.calculateCostLocked()
s.event.StatusCode = status
s.event.LatencyMS = int(time.Since(s.started).Milliseconds())
s.event.RecordedAt = time.Now().UTC()
labels := map[string]any{}
if s.target != "" && s.target != s.event.Model {
labels["target_model"] = s.target
labels["routed"] = true
}
if len(s.policies) > 0 {
labels["content_policies"] = s.policies
}
if s.redacted {
labels["content_redacted"] = true
}
if s.event.PriceID != "" {
labels["price_id"] = s.event.PriceID
labels["currency"] = s.event.Currency
}
s.event.Labels = labels
event := s.event
s.mu.Unlock()
s.recorder.Record(event)
}
func (s *auditSpan) calculateCostLocked() {
if s.pricing == nil || s.event.ProviderCode == "" || s.event.Model == "" {
return
}
model := s.event.Model
if s.target != "" {
model = s.target
}
cost := s.pricing.Calculate(s.event.ProviderCode, model, s.event.PromptTokens, s.event.CompletionTokens, s.started.UTC())
s.event.CostMicrounits = cost.Microunits
s.event.PriceID = cost.PriceID
s.event.Currency = cost.Currency
}
type captureReadCloser struct {
io.ReadCloser
buffer []byte
limit int
}
func (r *captureReadCloser) Read(buffer []byte) (int, error) {
n, err := r.ReadCloser.Read(buffer)
if n > 0 && len(r.buffer) < r.limit {
remaining := r.limit - len(r.buffer)
r.buffer = append(r.buffer, buffer[:min(n, remaining)]...)
}
return n, err
}
type statusResponseWriter struct {
http.ResponseWriter
status int
}
func (w *statusResponseWriter) WriteHeader(status int) {
if w.status != 0 {
return
}
w.status = status
w.ResponseWriter.WriteHeader(status)
}
func (w *statusResponseWriter) Write(buffer []byte) (int, error) {
if w.status == 0 {
w.WriteHeader(http.StatusOK)
}
return w.ResponseWriter.Write(buffer)
}
func (w *statusResponseWriter) Status() int {
if w.status == 0 {
return http.StatusOK
}
return w.status
}
func (w *statusResponseWriter) Flush() {
if w.status == 0 {
w.WriteHeader(http.StatusOK)
}
if flusher, ok := w.ResponseWriter.(http.Flusher); ok {
flusher.Flush()
}
}
func (w *statusResponseWriter) Unwrap() http.ResponseWriter { return w.ResponseWriter }
func (w *statusResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
if hijacker, ok := w.ResponseWriter.(http.Hijacker); ok {
return hijacker.Hijack()
}
return nil, nil, http.ErrNotSupported
}
func (w *statusResponseWriter) Push(target string, options *http.PushOptions) error {
if pusher, ok := w.ResponseWriter.(http.Pusher); ok {
return pusher.Push(target, options)
}
return http.ErrNotSupported
}
func (w *statusResponseWriter) ReadFrom(reader io.Reader) (int64, error) {
if w.status == 0 {
w.WriteHeader(http.StatusOK)
}
if readerFrom, ok := w.ResponseWriter.(io.ReaderFrom); ok {
return readerFrom.ReadFrom(reader)
}
return io.Copy(struct{ io.Writer }{w.ResponseWriter}, reader)
}
@@ -0,0 +1,141 @@
package gateway
import (
"context"
"io"
"log/slog"
"net/http"
"net/http/httptest"
"os"
"strings"
"sync/atomic"
"testing"
"time"
"aigateway.local/core/internal/audit"
"aigateway.local/core/internal/contentpolicy"
"aigateway.local/core/internal/pricing"
provideropenai "aigateway.local/core/internal/provider/openai"
"github.com/jackc/pgx/v5/pgxpool"
)
func TestContentPolicyAndPricingIntegration(t *testing.T) {
databaseURL := os.Getenv("CONTENT_PRICING_TEST_DATABASE_URL")
if databaseURL == "" {
t.Skip("CONTENT_PRICING_TEST_DATABASE_URL is not set")
}
ctx := context.Background()
pool, err := pgxpool.New(ctx, databaseURL)
if err != nil {
t.Fatal(err)
}
defer pool.Close()
const blockID = "00000000-0000-4000-8000-000000001601"
const priceID = "00000000-0000-4000-8000-000000001602"
_, err = pool.Exec(ctx, `INSERT INTO gateway.content_policies(id,name,action,priority,rules,enabled) VALUES($1,'integration block','block',2000,'[{"name":"forbidden","pattern":"forbidden","replacement":"[BLOCKED]"}]',true) ON CONFLICT(id) DO UPDATE SET enabled=true`, blockID)
if err != nil {
t.Fatal(err)
}
_, err = pool.Exec(ctx, `INSERT INTO gateway.model_prices(id,provider_code,model_pattern,input_microunits_per_million,output_microunits_per_million,currency,effective_from,enabled) VALUES($1,'environment','gpt-test',2000000,8000000,'USD',clock_timestamp()-interval '1 hour',true) ON CONFLICT(id) DO UPDATE SET enabled=true`, priceID)
if err != nil {
t.Fatal(err)
}
defer func() {
_, _ = pool.Exec(context.Background(), `DELETE FROM gateway.outbox_events WHERE event_type='content_policy.matched' AND aggregate_id IN ('m3-content-price-redact','m3-content-price-block')`)
_, _ = pool.Exec(context.Background(), `DELETE FROM gateway.content_policies WHERE id=$1`, blockID)
_, _ = pool.Exec(context.Background(), `DELETE FROM gateway.model_prices WHERE id=$1`, priceID)
}()
var upstreamCalls atomic.Int64
bodySeen := make(chan string, 1)
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
upstreamCalls.Add(1)
body, _ := io.ReadAll(r.Body)
bodySeen <- string(body)
w.Header().Set("Content-Type", "application/json")
_, _ = io.WriteString(w, `{"id":"ok","usage":{"prompt_tokens":1000,"completion_tokens":500,"total_tokens":1500}}`)
}))
defer upstream.Close()
adapter, _ := provideropenai.New(upstream.URL, "upstream-key")
engine := contentpolicy.NewEngine(pool, time.Minute, slog.Default())
if err := engine.Reload(ctx); err != nil {
t.Fatal(err)
}
prices := pricing.NewService(pool, time.Minute, slog.Default())
if err := prices.Reload(ctx); err != nil {
t.Fatal(err)
}
recorder := audit.NewRecorder(pool, slog.Default(), 100, 1, 10*time.Millisecond)
recordCtx, cancel := context.WithCancel(ctx)
stopped := make(chan struct{})
go func() { recorder.Run(recordCtx); close(stopped) }()
defer func() { cancel(); <-stopped }()
proxy := NewProxy(adapter, "test-key", 1<<20, slog.Default())
proxy.SetAuditRecorder(recorder)
proxy.SetContentPolicyEngine(engine)
proxy.SetPricingService(prices)
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
proxy.ServeHTTP(w, r.WithContext(WithRequestID(r.Context(), r.Header.Get("X-Request-ID"))))
}))
defer server.Close()
req, _ := http.NewRequest(http.MethodPost, server.URL+"/v1/chat/completions", strings.NewReader(`{"model":"gpt-test","messages":[{"role":"user","content":"token sk-1234567890123456"}]}`))
req.Header.Set("Authorization", "Bearer test-key")
req.Header.Set("X-Request-ID", "m3-content-price-redact")
response, err := http.DefaultClient.Do(req)
if err != nil {
t.Fatal(err)
}
_, _ = io.ReadAll(response.Body)
_ = response.Body.Close()
if response.StatusCode != 200 || response.Header.Get("X-Gateway-Content-Redacted") != "true" {
t.Fatalf("unexpected response %d headers=%v", response.StatusCode, response.Header)
}
select {
case body := <-bodySeen:
if strings.Contains(body, "sk-1234567890123456") || !strings.Contains(body, "[REDACTED_API_KEY]") {
t.Fatalf("upstream saw unsafe body %s", body)
}
case <-time.After(time.Second):
t.Fatal("upstream did not receive request")
}
blocked, _ := http.NewRequest(http.MethodPost, server.URL+"/v1/responses", strings.NewReader(`{"model":"gpt-test","input":"forbidden"}`))
blocked.Header.Set("Authorization", "Bearer test-key")
blocked.Header.Set("X-Request-ID", "m3-content-price-block")
blockedResponse, err := http.DefaultClient.Do(blocked)
if err != nil {
t.Fatal(err)
}
_, _ = io.ReadAll(blockedResponse.Body)
_ = blockedResponse.Body.Close()
if blockedResponse.StatusCode != http.StatusUnprocessableEntity || upstreamCalls.Load() != 1 {
t.Fatalf("block failed status=%d calls=%d", blockedResponse.StatusCode, upstreamCalls.Load())
}
deadline := time.Now().Add(2 * time.Second)
for {
var cost *int64
var labels string
err = pool.QueryRow(ctx, `SELECT cost_microunits,labels::text FROM gateway.audit_events WHERE request_id='m3-content-price-redact' ORDER BY recorded_at DESC LIMIT 1`).Scan(&cost, &labels)
if err == nil && cost != nil && *cost == 6000 && strings.Contains(labels, "content_redacted") {
break
}
if time.Now().After(deadline) {
t.Fatalf("audit cost/redaction missing cost=%v labels=%s err=%v", cost, labels, err)
}
time.Sleep(20 * time.Millisecond)
}
deadline = time.Now().Add(2 * time.Second)
for {
var count int
err = pool.QueryRow(ctx, `SELECT count(*) FROM gateway.outbox_events WHERE event_type='content_policy.matched' AND aggregate_id IN ('m3-content-price-redact','m3-content-price-block')`).Scan(&count)
if err == nil && count == 2 {
break
}
if time.Now().After(deadline) {
t.Fatalf("content policy notification outbox missing count=%d err=%v", count, err)
}
time.Sleep(20 * time.Millisecond)
}
}
+14
View File
@@ -0,0 +1,14 @@
package gateway
import "context"
type requestIDKey struct{}
func WithRequestID(ctx context.Context, requestID string) context.Context {
return context.WithValue(ctx, requestIDKey{}, requestID)
}
func RequestID(ctx context.Context) string {
value, _ := ctx.Value(requestIDKey{}).(string)
return value
}
+430
View File
@@ -0,0 +1,430 @@
package gateway
import (
"context"
"crypto/subtle"
"encoding/json"
"errors"
"fmt"
"log/slog"
"net"
"net/http"
"net/http/httputil"
"strconv"
"strings"
"sync"
"time"
"aigateway.local/core/internal/apikey"
"aigateway.local/core/internal/contentpolicy"
"aigateway.local/core/internal/pricing"
"aigateway.local/core/internal/provider"
)
type Proxy struct {
resolver AdapterResolver
auth apikey.KeyAuthenticator
maxBody int64
logger *slog.Logger
transport *http.Transport
proxies sync.Map
circuits sync.Map
admission AdmissionController
tokenQuota TokenQuotaController
resilience ResiliencePolicy
audit AuditRecorder
policies *contentpolicy.Engine
pricing *pricing.Service
}
type cachedProxy struct {
key string
proxy *httputil.ReverseProxy
}
var (
ErrProviderNotFound = errors.New("requested provider is not available")
ErrProviderUnavailable = errors.New("provider configuration is unavailable")
)
type ResolvedAdapter struct {
Code string
Revision int64
Adapter provider.Adapter
Capabilities map[provider.Capability]bool
}
type AdapterResolver interface {
Resolve(providerCode string) (ResolvedAdapter, error)
}
func NewProxy(adapter provider.Adapter, apiKey string, maxBody int64, logger *slog.Logger) *Proxy {
return NewProxyWithAuthenticator(adapter, staticKeyAuthenticator(apiKey), maxBody, logger)
}
func NewProxyWithAuthenticator(adapter provider.Adapter, authenticator apikey.KeyAuthenticator, maxBody int64, logger *slog.Logger) *Proxy {
return NewDynamicProxy(staticAdapterResolver{adapter: adapter}, authenticator, maxBody, logger)
}
func NewDynamicProxy(resolver AdapterResolver, authenticator apikey.KeyAuthenticator, maxBody int64, logger *slog.Logger) *Proxy {
transport := &http.Transport{
Proxy: http.ProxyFromEnvironment,
DialContext: (&net.Dialer{Timeout: 5 * time.Second, KeepAlive: 30 * time.Second}).DialContext,
ForceAttemptHTTP2: true,
MaxIdleConns: 512,
MaxIdleConnsPerHost: 256,
IdleConnTimeout: 90 * time.Second,
TLSHandshakeTimeout: 5 * time.Second,
ResponseHeaderTimeout: 60 * time.Second,
ExpectContinueTimeout: time.Second,
}
return &Proxy{resolver: resolver, auth: authenticator, maxBody: maxBody, logger: logger, transport: transport, resilience: DefaultResiliencePolicy()}
}
func (p *Proxy) SetAdmissionController(controller AdmissionController) {
p.admission = controller
}
func (p *Proxy) SetTokenQuotaController(controller TokenQuotaController) {
p.tokenQuota = controller
}
func (p *Proxy) SetResiliencePolicy(policy ResiliencePolicy) {
p.resilience = policy
p.transport.ResponseHeaderTimeout = policy.ResponseHeaderTimeout
}
func (p *Proxy) SetAuditRecorder(recorder AuditRecorder) { p.audit = recorder }
func (p *Proxy) SetContentPolicyEngine(engine *contentpolicy.Engine) { p.policies = engine }
func (p *Proxy) SetPricingService(service *pricing.Service) { p.pricing = service }
func (p *Proxy) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
started := time.Now()
if !isSupportedPath(request.URL.Path, request.Method) {
writeOpenAIError(writer, http.StatusNotFound, "invalid_request_error", "unsupported gateway endpoint")
return
}
principal, err := p.authorized(request)
if err != nil {
if errors.Is(err, apikey.ErrStore) {
writeOpenAIError(writer, http.StatusServiceUnavailable, "authentication_unavailable", "API key service is unavailable")
return
}
writer.Header().Set("WWW-Authenticate", "Bearer")
writeOpenAIError(writer, http.StatusUnauthorized, "authentication_error", "invalid API key")
return
}
span := newAuditSpan(p.audit, principal, request, started)
if span != nil {
span.pricing = p.pricing
}
if span != nil {
statusWriter := &statusResponseWriter{ResponseWriter: writer}
writer = statusWriter
defer func() { span.finish(statusWriter.Status()) }()
}
if p.admission != nil && principal.APIKeyID != "" {
decision, err := p.admission.Allow(request.Context(), principal, time.Now())
if err != nil {
writeOpenAIError(writer, http.StatusServiceUnavailable, "rate_limit_unavailable", "rate limit service is unavailable")
return
}
writeAdmissionHeaders(writer.Header(), decision)
if !decision.Allowed {
writer.Header().Set("Retry-After", strconv.FormatInt(max(int64(decision.RetryAfter.Seconds()), 1), 10))
if decision.Reason == AdmissionMonthlyQuota {
writeOpenAIError(writer, http.StatusTooManyRequests, "insufficient_quota", "monthly request quota exceeded")
} else {
writeOpenAIError(writer, http.StatusTooManyRequests, "rate_limit_error", "request rate limit exceeded")
}
return
}
}
if request.ContentLength > p.maxBody {
writeOpenAIError(writer, http.StatusRequestEntityTooLarge, "invalid_request_error", "request body is too large")
return
}
if p.policies != nil {
policyResult, policyErr := p.policies.Apply(request, p.maxBody, principal.APIKeyID)
if errors.Is(policyErr, contentpolicy.ErrBodyTooLarge) {
writeOpenAIError(writer, http.StatusRequestEntityTooLarge, "invalid_request_error", "request body is too large")
return
}
if policyErr != nil {
writeOpenAIError(writer, http.StatusBadRequest, "invalid_request_error", "request body could not be inspected")
return
}
span.setContentPolicy(policyResult)
if policyResult.Blocked {
writeOpenAIError(writer, http.StatusUnprocessableEntity, "content_policy_violation", "request was blocked by content policy")
return
}
if policyResult.Redacted {
writer.Header().Set("X-Gateway-Content-Redacted", "true")
}
}
routeResolver, routingEnabled := p.resolver.(ModelRouteResolver)
routingEnabled = routingEnabled && routeResolver.ModelRoutingEnabled()
var modelPayload modelRequest
if routingEnabled {
modelPayload, err = readModelRequest(request, p.maxBody)
if errors.Is(err, errRequestBodyTooLarge) {
writeOpenAIError(writer, http.StatusRequestEntityTooLarge, "invalid_request_error", "request body is too large")
return
}
if err != nil {
writeOpenAIError(writer, http.StatusBadRequest, "invalid_request_error", "request body could not be read")
return
}
}
var usage *usageSession
if p.tokenQuota != nil && principal.APIKeyID != "" {
estimate := int64(0)
if principal.MonthlyTokenQuota > 0 {
estimate, err = prepareTokenBudget(request, p.maxBody)
if errors.Is(err, errRequestBodyTooLarge) {
writeOpenAIError(writer, http.StatusRequestEntityTooLarge, "invalid_request_error", "request body is too large")
return
}
if err != nil {
writeOpenAIError(writer, http.StatusBadRequest, "invalid_request_error", "request body could not be read")
return
}
}
reservation, reserveErr := p.tokenQuota.Reserve(request.Context(), principal, estimate, time.Now())
if reserveErr != nil && principal.MonthlyTokenQuota > 0 {
writeOpenAIError(writer, http.StatusServiceUnavailable, "token_quota_unavailable", "token quota service is unavailable")
return
}
if reserveErr == nil && !reservation.Allowed {
writeTokenQuotaHeaders(writer.Header(), reservation)
writer.Header().Set("Retry-After", strconv.FormatInt(max(int64(time.Until(reservation.ResetAt).Seconds()), 1), 10))
writeOpenAIError(writer, http.StatusTooManyRequests, "insufficient_quota", "monthly token quota exceeded")
return
}
if reserveErr == nil && reservation.CounterKey != "" {
writeTokenQuotaHeaders(writer.Header(), reservation)
usage = &usageSession{controller: p.tokenQuota, reservation: reservation, log: p.logger}
request = withUsageSession(request, usage)
defer usage.finish(TokenUsage{})
}
}
if span != nil {
if usage == nil {
usage = &usageSession{log: p.logger}
request = withUsageSession(request, usage)
defer usage.finish(TokenUsage{})
}
usage.onFinish = span.setUsage
}
if request.Body != nil && request.Method != http.MethodGet {
if request.Header.Get("Idempotency-Key") != "" && request.GetBody == nil {
if _, err := prepareTokenBudget(request, p.maxBody); err != nil {
if errors.Is(err, errRequestBodyTooLarge) {
writeOpenAIError(writer, http.StatusRequestEntityTooLarge, "invalid_request_error", "request body is too large")
} else {
writeOpenAIError(writer, http.StatusBadRequest, "invalid_request_error", "request body could not be read")
}
return
}
}
request.Body = http.MaxBytesReader(writer, request.Body, p.maxBody)
}
providerCode := strings.ToLower(strings.TrimSpace(request.Header.Get("X-Gateway-Provider")))
var resolved ResolvedAdapter
if routingEnabled && modelPayload.model != "" {
span.setModel(modelPayload.model)
tenantID := ""
if principal.TenantID != nil {
tenantID = *principal.TenantID
}
route, routeErr := routeResolver.ResolveModelRoute(ModelRouteQuery{
ProviderCode: providerCode, Model: modelPayload.model, Endpoint: request.URL.Path,
APIKeyID: principal.APIKeyID, TenantID: tenantID, Seed: RequestID(request.Context()),
})
if routeErr != nil {
writeOpenAIError(writer, http.StatusServiceUnavailable, "provider_unavailable", "model routing configuration is unavailable")
return
}
if route.Known && !route.Matched {
writeOpenAIError(writer, http.StatusBadRequest, "invalid_request_error", "model route is not available for this request")
return
}
if route.Matched {
resolved = route.ResolvedAdapter
if err := modelPayload.rewrite(request, route.TargetModel); err != nil {
writeOpenAIError(writer, http.StatusBadRequest, "invalid_request_error", "request model could not be rewritten")
return
}
writer.Header().Set("X-Gateway-Model", route.TargetModel)
}
}
if resolved.Adapter == nil {
resolved, err = p.resolver.Resolve(providerCode)
}
if err != nil {
if errors.Is(err, ErrProviderNotFound) {
writeOpenAIError(writer, http.StatusBadRequest, "invalid_request_error", "requested provider is not available")
return
}
writeOpenAIError(writer, http.StatusServiceUnavailable, "provider_unavailable", "provider configuration is unavailable")
return
}
capability := capabilityForPath(request.URL.Path)
if capability != "" && !resolved.Capabilities[capability] {
writeOpenAIError(writer, http.StatusBadRequest, "invalid_request_error", "provider does not support this endpoint")
return
}
request.Header.Del("X-Gateway-Provider")
writer.Header().Set("X-Gateway-Provider", resolved.Code)
span.setRoute(resolved.Code, writer.Header().Get("X-Gateway-Model"))
request.Body = span.captureBody(request.Body)
p.proxyFor(resolved).ServeHTTP(writer, request)
}
func (p *Proxy) proxyFor(resolved ResolvedAdapter) *httputil.ReverseProxy {
target := resolved.Adapter.Target()
key := fmt.Sprintf("%s:%d:%s", resolved.Code, resolved.Revision, target.String())
if cached, ok := p.proxies.Load(resolved.Code); ok {
entry := cached.(cachedProxy)
if entry.key == key {
return entry.proxy
}
}
reverseProxy := httputil.NewSingleHostReverseProxy(target)
originalDirector := reverseProxy.Director
reverseProxy.Director = func(request *http.Request) {
originalDirector(request)
request.Host = target.Host
resolved.Adapter.Prepare(request)
}
reverseProxy.FlushInterval = -1
circuitValue, _ := p.circuits.LoadOrStore(resolved.Code, newCircuitBreaker(p.resilience))
reverseProxy.Transport = &resilientTransport{
base: p.transport, circuit: circuitValue.(*circuitBreaker), maxRetries: p.resilience.MaxRetries, backoff: p.resilience.RetryBackoff,
}
reverseProxy.ModifyResponse = func(response *http.Response) error {
if session := usageSessionFrom(response.Request); session != nil {
if response.StatusCode >= http.StatusOK && response.StatusCode < http.StatusMultipleChoices {
session.fallback = session.reservation.Reserved
response.Body = newUsageReadCloser(response.Body, response.Header.Get("Content-Type"), session)
} else {
session.fallback = 0
session.finish(TokenUsage{})
}
}
return nil
}
reverseProxy.ErrorHandler = func(writer http.ResponseWriter, request *http.Request, err error) {
if session := usageSessionFrom(request); session != nil {
session.fallback = 0
session.finish(TokenUsage{})
}
p.logger.Error("upstream request failed", "request_id", RequestID(request.Context()), "provider", resolved.Code, "error", err)
if errors.Is(err, ErrCircuitOpen) {
writer.Header().Set("Retry-After", "30")
writeOpenAIError(writer, http.StatusServiceUnavailable, "provider_unavailable", "provider circuit is temporarily open")
return
}
writeOpenAIError(writer, http.StatusBadGateway, "upstream_error", "upstream service is unavailable")
}
p.proxies.Store(resolved.Code, cachedProxy{key: key, proxy: reverseProxy})
return reverseProxy
}
func (p *Proxy) authorized(request *http.Request) (apikey.Principal, error) {
presented := strings.TrimSpace(request.Header.Get("X-Gateway-API-Key"))
if presented == "" {
authorization := strings.TrimSpace(request.Header.Get("Authorization"))
if len(authorization) > len("Bearer ") && strings.EqualFold(authorization[:len("Bearer ")], "Bearer ") {
presented = strings.TrimSpace(authorization[len("Bearer "):])
}
}
if p.auth == nil {
return apikey.Principal{}, apikey.ErrInvalid
}
if authenticator, ok := p.auth.(apikey.PrincipalAuthenticator); ok {
return authenticator.AuthenticatePrincipal(request.Context(), presented)
}
return apikey.Principal{}, p.auth.Authenticate(request.Context(), presented)
}
func writeAdmissionHeaders(header http.Header, decision AdmissionDecision) {
if decision.Limit <= 0 || decision.ResetAt.IsZero() {
return
}
header.Set("X-RateLimit-Limit", strconv.FormatInt(decision.Limit, 10))
header.Set("X-RateLimit-Remaining", strconv.FormatInt(max(decision.Remaining, 0), 10))
header.Set("X-RateLimit-Reset", strconv.FormatInt(decision.ResetAt.Unix(), 10))
}
type staticKeyAuthenticator string
func (a staticKeyAuthenticator) Authenticate(_ context.Context, presented string) error {
key := string(a)
if key == "" {
return nil
}
if len(presented) != len(key) || subtle.ConstantTimeCompare([]byte(presented), []byte(key)) != 1 {
return apikey.ErrInvalid
}
return nil
}
type staticAdapterResolver struct{ adapter provider.Adapter }
func (r staticAdapterResolver) Resolve(code string) (ResolvedAdapter, error) {
if r.adapter == nil || code != "" && code != "environment" {
return ResolvedAdapter{}, ErrProviderNotFound
}
capabilities := make(map[provider.Capability]bool)
for _, capability := range r.adapter.Capabilities() {
capabilities[capability] = true
}
return ResolvedAdapter{Code: "environment", Adapter: r.adapter, Capabilities: capabilities}, nil
}
func capabilityForPath(path string) provider.Capability {
switch path {
case "/v1/models":
return provider.CapabilityModels
case "/v1/chat/completions":
return provider.CapabilityChat
case "/v1/responses":
return provider.CapabilityResponses
case "/v1/embeddings":
return provider.CapabilityEmbeddings
case "/v1/messages":
return provider.CapabilityMessages
default:
return ""
}
}
func isSupportedPath(path, method string) bool {
switch path {
case "/v1/models":
return method == http.MethodGet
case "/v1/chat/completions", "/v1/responses", "/v1/embeddings", "/v1/messages":
return method == http.MethodPost
default:
return false
}
}
func writeOpenAIError(writer http.ResponseWriter, status int, errorType, message string) {
writer.Header().Set("Content-Type", "application/json")
writer.WriteHeader(status)
_ = json.NewEncoder(writer).Encode(map[string]any{
"error": map[string]any{"message": message, "type": errorType, "param": nil, "code": nil},
})
}
var errInvalidAdapter = errors.New("invalid provider adapter")
func ValidateAdapter(adapter provider.Adapter) error {
if adapter == nil || adapter.Target() == nil || adapter.Name() == "" {
return errInvalidAdapter
}
return nil
}
+252
View File
@@ -0,0 +1,252 @@
package gateway
import (
"context"
"io"
"log/slog"
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
"aigateway.local/core/internal/apikey"
auditpkg "aigateway.local/core/internal/audit"
"aigateway.local/core/internal/provider"
provideropenai "aigateway.local/core/internal/provider/openai"
)
type principalAuthenticator struct{ principal apikey.Principal }
func (a principalAuthenticator) Authenticate(context.Context, string) error { return nil }
func (a principalAuthenticator) AuthenticatePrincipal(context.Context, string) (apikey.Principal, error) {
return a.principal, nil
}
type fixedAdmission struct{ decision AdmissionDecision }
func (a fixedAdmission) Allow(context.Context, apikey.Principal, time.Time) (AdmissionDecision, error) {
return a.decision, nil
}
type recordingTokenQuota struct {
reservation TokenReservation
estimate int64
actual int64
}
type fixedRoutingResolver struct {
adapter ResolvedAdapter
}
type recordingAudit struct{ event auditpkg.Event }
func (r *recordingAudit) Record(event auditpkg.Event) bool { r.event = event; return true }
func (r fixedRoutingResolver) Resolve(string) (ResolvedAdapter, error) { return r.adapter, nil }
func (r fixedRoutingResolver) ModelRoutingEnabled() bool { return true }
func (r fixedRoutingResolver) ResolveModelRoute(query ModelRouteQuery) (ModelRouteResult, error) {
if query.Model == "public-chat" {
return ModelRouteResult{ResolvedAdapter: r.adapter, TargetModel: "upstream-chat", Matched: true, Known: true}, nil
}
return ModelRouteResult{}, nil
}
func (q *recordingTokenQuota) Reserve(_ context.Context, _ apikey.Principal, estimate int64, _ time.Time) (TokenReservation, error) {
q.estimate = estimate
return q.reservation, nil
}
func (q *recordingTokenQuota) Commit(_ context.Context, _ TokenReservation, actual int64) error {
q.actual = actual
return nil
}
func TestProxyRejectsInvalidKey(t *testing.T) {
upstream := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
t.Fatal("upstream must not be called")
}))
defer upstream.Close()
adapter, err := provideropenai.New(upstream.URL, "upstream-secret")
if err != nil {
t.Fatal(err)
}
proxy := NewProxy(adapter, "gateway-secret", 1024, slog.New(slog.NewTextHandler(io.Discard, nil)))
request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
request.Header.Set("Authorization", "Bearer wrong")
response := httptest.NewRecorder()
proxy.ServeHTTP(response, request)
if response.Code != http.StatusUnauthorized {
t.Fatalf("got %d, want %d", response.Code, http.StatusUnauthorized)
}
}
func TestProxyRejectsKnownOversizedBody(t *testing.T) {
upstream := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
t.Fatal("upstream must not be called")
}))
defer upstream.Close()
adapter, _ := provideropenai.New(upstream.URL, "upstream-secret")
proxy := NewProxy(adapter, "gateway-secret", 4, slog.New(slog.NewTextHandler(io.Discard, nil)))
request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader("12345"))
request.Header.Set("Authorization", "Bearer gateway-secret")
response := httptest.NewRecorder()
proxy.ServeHTTP(response, request)
if response.Code != http.StatusRequestEntityTooLarge {
t.Fatalf("got %d, want %d", response.Code, http.StatusRequestEntityTooLarge)
}
}
func TestProxyReplacesClientAuthorization(t *testing.T) {
upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
if got := request.Header.Get("Authorization"); got != "Bearer upstream-secret" {
t.Fatalf("unexpected upstream authorization: %q", got)
}
writer.Header().Set("Content-Type", "application/json")
_, _ = writer.Write([]byte(`{"ok":true}`))
}))
defer upstream.Close()
adapter, _ := provideropenai.New(upstream.URL, "upstream-secret")
proxy := NewProxy(adapter, "gateway-secret", 1024, slog.New(slog.NewTextHandler(io.Discard, nil)))
request := httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
request.Header.Set("Authorization", "Bearer gateway-secret")
response := httptest.NewRecorder()
proxy.ServeHTTP(response, request)
if response.Code != http.StatusOK {
t.Fatalf("got %d, want %d", response.Code, http.StatusOK)
}
}
func TestProxyRejectsRateLimitedPrincipalBeforeUpstream(t *testing.T) {
upstream := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
t.Fatal("upstream must not be called")
}))
defer upstream.Close()
adapter, _ := provideropenai.New(upstream.URL, "upstream-secret")
proxy := NewProxyWithAuthenticator(adapter, principalAuthenticator{principal: apikey.Principal{
APIKeyID: "key-1", Scopes: []string{"gateway:invoke"}, RequestsPerMinute: 2,
}}, 1024, slog.New(slog.NewTextHandler(io.Discard, nil)))
proxy.SetAdmissionController(fixedAdmission{decision: AdmissionDecision{
Allowed: false, Reason: AdmissionRateLimit, Limit: 2, Remaining: 0,
ResetAt: time.Now().Add(time.Minute), RetryAfter: 30 * time.Second,
}})
request := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
request.Header.Set("Authorization", "Bearer gateway-secret")
response := httptest.NewRecorder()
proxy.ServeHTTP(response, request)
if response.Code != http.StatusTooManyRequests {
t.Fatalf("got %d, want %d", response.Code, http.StatusTooManyRequests)
}
if response.Header().Get("X-RateLimit-Limit") != "2" || response.Header().Get("Retry-After") == "" {
t.Fatalf("missing rate limit headers: %#v", response.Header())
}
if !strings.Contains(response.Body.String(), `"type":"rate_limit_error"`) {
t.Fatalf("unexpected body: %s", response.Body.String())
}
}
func TestProxyReconcilesReservedTokensWithUpstreamUsage(t *testing.T) {
upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
writer.Header().Set("Content-Type", "application/json")
_, _ = writer.Write([]byte(`{"choices":[],"usage":{"prompt_tokens":7,"completion_tokens":5,"total_tokens":12}}`))
}))
defer upstream.Close()
adapter, _ := provideropenai.New(upstream.URL, "upstream-secret")
quota := &recordingTokenQuota{reservation: TokenReservation{
Allowed: true, CounterKey: "tokens", Reserved: 24, Limit: 1000, Remaining: 976, ResetAt: time.Now().Add(time.Hour),
}}
proxy := NewProxyWithAuthenticator(adapter, principalAuthenticator{principal: apikey.Principal{
APIKeyID: "key-1", Scopes: []string{"gateway:invoke"}, MonthlyTokenQuota: 1000,
}}, 1024, slog.New(slog.NewTextHandler(io.Discard, nil)))
proxy.SetTokenQuotaController(quota)
request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{"model":"test","max_tokens":20}`))
request.Header.Set("Authorization", "Bearer gateway-secret")
response := httptest.NewRecorder()
proxy.ServeHTTP(response, request)
if response.Code != http.StatusOK {
t.Fatalf("got %d, want %d: %s", response.Code, http.StatusOK, response.Body.String())
}
if quota.estimate <= 20 {
t.Fatalf("request budget did not include input tokens: %d", quota.estimate)
}
if quota.actual != 12 {
t.Fatalf("committed %d tokens, want 12", quota.actual)
}
}
func TestProxyRejectsTokenReservationBeforeUpstream(t *testing.T) {
upstream := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
t.Fatal("upstream must not be called")
}))
defer upstream.Close()
adapter, _ := provideropenai.New(upstream.URL, "upstream-secret")
quota := &recordingTokenQuota{reservation: TokenReservation{
Allowed: false, Limit: 100, Remaining: 4, ResetAt: time.Now().Add(time.Hour),
}}
proxy := NewProxyWithAuthenticator(adapter, principalAuthenticator{principal: apikey.Principal{
APIKeyID: "key-1", Scopes: []string{"gateway:invoke"}, MonthlyTokenQuota: 100,
}}, 1024, slog.New(slog.NewTextHandler(io.Discard, nil)))
proxy.SetTokenQuotaController(quota)
request := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"test","max_output_tokens":20}`))
request.Header.Set("Authorization", "Bearer gateway-secret")
response := httptest.NewRecorder()
proxy.ServeHTTP(response, request)
if response.Code != http.StatusTooManyRequests || !strings.Contains(response.Body.String(), `"type":"insufficient_quota"`) {
t.Fatalf("unexpected response %d: %s", response.Code, response.Body.String())
}
if response.Header().Get("X-TokenLimit-Limit") != "100" {
t.Fatalf("missing token quota headers: %#v", response.Header())
}
}
func TestProxyRewritesModelAliasBeforeUpstream(t *testing.T) {
upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
body, _ := io.ReadAll(request.Body)
if !strings.Contains(string(body), `"model":"upstream-chat"`) {
t.Fatalf("model alias was not rewritten: %s", body)
}
writer.Header().Set("Content-Type", "application/json")
_, _ = writer.Write([]byte(`{"choices":[]}`))
}))
defer upstream.Close()
adapter, _ := provideropenai.New(upstream.URL, "upstream-secret")
resolved := ResolvedAdapter{Code: "routed", Adapter: adapter, Capabilities: map[provider.Capability]bool{provider.CapabilityChat: true}}
proxy := NewDynamicProxy(fixedRoutingResolver{adapter: resolved}, principalAuthenticator{principal: apikey.Principal{Scopes: []string{"gateway:invoke"}}}, 1024, slog.New(slog.NewTextHandler(io.Discard, nil)))
request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{"model":"public-chat","messages":[]}`))
request.Header.Set("Authorization", "Bearer gateway-secret")
response := httptest.NewRecorder()
proxy.ServeHTTP(response, request)
if response.Code != http.StatusOK || response.Header().Get("X-Gateway-Model") != "upstream-chat" || response.Header().Get("X-Gateway-Provider") != "routed" {
t.Fatalf("unexpected routed response %d %#v: %s", response.Code, response.Header(), response.Body.String())
}
}
func TestProxyRecordsAuditWithoutBufferingWholeResponse(t *testing.T) {
upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
writer.Header().Set("Content-Type", "application/json")
_, _ = writer.Write([]byte(`{"choices":[],"usage":{"prompt_tokens":9,"completion_tokens":4,"total_tokens":13}}`))
}))
defer upstream.Close()
adapter, _ := provideropenai.New(upstream.URL, "upstream-secret")
recorder := &recordingAudit{}
proxy := NewProxyWithAuthenticator(adapter, principalAuthenticator{principal: apikey.Principal{
APIKeyID: "11111111-1111-4111-8111-111111111111", Scopes: []string{"gateway:invoke"},
}}, 1024, slog.New(slog.NewTextHandler(io.Discard, nil)))
proxy.SetAuditRecorder(recorder)
request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{"model":"audit-model","messages":[]}`))
request.Header.Set("Authorization", "Bearer gateway-secret")
response := httptest.NewRecorder()
proxy.ServeHTTP(response, request)
if response.Code != http.StatusOK {
t.Fatalf("unexpected response %d: %s", response.Code, response.Body.String())
}
if recorder.event.Model != "audit-model" || recorder.event.ProviderCode != "environment" || recorder.event.StatusCode != 200 || recorder.event.PromptTokens != 9 || recorder.event.CompletionTokens != 4 {
t.Fatalf("unexpected audit event: %#v", recorder.event)
}
}
+46
View File
@@ -0,0 +1,46 @@
package gateway
import (
"bytes"
"encoding/json"
"errors"
"io"
"net/http"
)
var errRequestBodyTooLarge = errors.New("request body is too large")
func prepareTokenBudget(request *http.Request, maxBody int64) (int64, error) {
if request.Body == nil || request.Method == http.MethodGet {
return 0, nil
}
limited := io.LimitReader(request.Body, maxBody+1)
body, err := io.ReadAll(limited)
if err != nil {
return 0, err
}
_ = request.Body.Close()
if int64(len(body)) > maxBody {
return 0, errRequestBodyTooLarge
}
restoreRequestBody(request, body)
inputEstimate := max(int64((len(body)+3)/4), 1)
if request.URL.Path == "/v1/embeddings" {
return inputEstimate, nil
}
outputBudget := int64(1024)
decoder := json.NewDecoder(bytes.NewReader(body))
decoder.UseNumber()
var payload map[string]any
if decoder.Decode(&payload) == nil {
for _, field := range []string{"max_output_tokens", "max_completion_tokens", "max_tokens"} {
if value := jsonInt64(payload[field]); value > 0 {
outputBudget = value
break
}
}
}
// A per-request ceiling prevents malicious payloads from reserving unbounded Redis counters.
return min(inputEstimate+outputBudget, int64(100_000_000)), nil
}
+149
View File
@@ -0,0 +1,149 @@
package gateway
import (
"context"
"errors"
"io"
"net/http"
"sync"
"time"
)
var ErrCircuitOpen = errors.New("provider circuit is open")
type ResiliencePolicy struct {
ResponseHeaderTimeout time.Duration
MaxRetries int
RetryBackoff time.Duration
CircuitThreshold int
CircuitOpenDuration time.Duration
}
func DefaultResiliencePolicy() ResiliencePolicy {
return ResiliencePolicy{
ResponseHeaderTimeout: 60 * time.Second, MaxRetries: 2, RetryBackoff: 50 * time.Millisecond,
CircuitThreshold: 5, CircuitOpenDuration: 30 * time.Second,
}
}
type circuitBreaker struct {
mu sync.Mutex
failures int
threshold int
openFor time.Duration
openUntil time.Time
halfOpenRun bool
}
func newCircuitBreaker(policy ...ResiliencePolicy) *circuitBreaker {
settings := DefaultResiliencePolicy()
if len(policy) > 0 {
settings = policy[0]
}
return &circuitBreaker{threshold: settings.CircuitThreshold, openFor: settings.CircuitOpenDuration}
}
func (c *circuitBreaker) allow(now time.Time) bool {
c.mu.Lock()
defer c.mu.Unlock()
if c.openUntil.IsZero() {
return true
}
if now.Before(c.openUntil) || c.halfOpenRun {
return false
}
c.halfOpenRun = true
return true
}
func (c *circuitBreaker) success() {
c.mu.Lock()
c.failures = 0
c.openUntil = time.Time{}
c.halfOpenRun = false
c.mu.Unlock()
}
func (c *circuitBreaker) failure(now time.Time) {
c.mu.Lock()
defer c.mu.Unlock()
c.halfOpenRun = false
c.failures++
if c.failures >= c.threshold || !c.openUntil.IsZero() {
c.openUntil = now.Add(c.openFor)
}
}
type resilientTransport struct {
base http.RoundTripper
circuit *circuitBreaker
maxRetries int
backoff time.Duration
}
func (t *resilientTransport) RoundTrip(request *http.Request) (*http.Response, error) {
if !t.circuit.allow(time.Now()) {
return nil, ErrCircuitOpen
}
replayable := request.Method == http.MethodGet || request.Method == http.MethodHead ||
request.Header.Get("Idempotency-Key") != "" && request.GetBody != nil
attempts := 1
if replayable {
attempts += t.maxRetries
}
var response *http.Response
var err error
for attempt := 0; attempt < attempts; attempt++ {
current := request
if attempt > 0 {
if waitErr := waitBackoff(request.Context(), t.backoff*time.Duration(attempt)); waitErr != nil {
err = waitErr
break
}
current = request.Clone(request.Context())
if request.Body != nil && request.Body != http.NoBody {
body, bodyErr := request.GetBody()
if bodyErr != nil {
err = bodyErr
break
}
current.Body = body
}
}
response, err = t.base.RoundTrip(current)
if !retryableResult(response, err) || attempt == attempts-1 {
break
}
if response != nil {
_, _ = io.CopyN(io.Discard, response.Body, 4096)
_ = response.Body.Close()
}
}
if retryableResult(response, err) {
t.circuit.failure(time.Now())
} else {
t.circuit.success()
}
return response, err
}
func retryableResult(response *http.Response, err error) bool {
if err != nil {
return !errors.Is(err, context.Canceled) && !errors.Is(err, context.DeadlineExceeded)
}
return response != nil && (response.StatusCode == http.StatusBadGateway || response.StatusCode == http.StatusServiceUnavailable || response.StatusCode == http.StatusGatewayTimeout)
}
func waitBackoff(ctx context.Context, duration time.Duration) error {
if duration <= 0 {
return nil
}
timer := time.NewTimer(duration)
defer timer.Stop()
select {
case <-ctx.Done():
return ctx.Err()
case <-timer.C:
return nil
}
}
+61
View File
@@ -0,0 +1,61 @@
package gateway
import (
"errors"
"io"
"net/http"
"strings"
"testing"
"time"
)
type roundTripFunc func(*http.Request) (*http.Response, error)
func (f roundTripFunc) RoundTrip(request *http.Request) (*http.Response, error) { return f(request) }
func TestResilientTransportRetriesReplayableRequest(t *testing.T) {
attempts := 0
transport := &resilientTransport{base: roundTripFunc(func(*http.Request) (*http.Response, error) {
attempts++
if attempts < 3 {
return &http.Response{StatusCode: http.StatusServiceUnavailable, Body: io.NopCloser(strings.NewReader("busy")), Header: make(http.Header)}, nil
}
return &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader("ok")), Header: make(http.Header)}, nil
}), circuit: newCircuitBreaker(), maxRetries: 2}
request, _ := http.NewRequest(http.MethodGet, "https://provider.example/v1/models", nil)
response, err := transport.RoundTrip(request)
if err != nil || response.StatusCode != http.StatusOK || attempts != 3 {
t.Fatalf("response=%v err=%v attempts=%d", response, err, attempts)
}
}
func TestResilientTransportDoesNotRetryUnsafePost(t *testing.T) {
attempts := 0
transport := &resilientTransport{base: roundTripFunc(func(*http.Request) (*http.Response, error) {
attempts++
return nil, errors.New("network failure")
}), circuit: newCircuitBreaker(), maxRetries: 2}
request, _ := http.NewRequest(http.MethodPost, "https://provider.example/v1/chat/completions", strings.NewReader("{}"))
_, _ = transport.RoundTrip(request)
if attempts != 1 {
t.Fatalf("unsafe POST attempted %d times", attempts)
}
}
func TestCircuitOpensAndAllowsSingleProbe(t *testing.T) {
circuit := newCircuitBreaker()
circuit.threshold = 2
circuit.openFor = time.Millisecond
now := time.Now()
circuit.failure(now)
circuit.failure(now)
if circuit.allow(now) {
t.Fatal("open circuit allowed request")
}
if !circuit.allow(now.Add(2 * time.Millisecond)) {
t.Fatal("circuit did not allow half-open probe")
}
if circuit.allow(now.Add(2 * time.Millisecond)) {
t.Fatal("circuit allowed concurrent half-open probe")
}
}
+76
View File
@@ -0,0 +1,76 @@
package gateway
import (
"bytes"
"encoding/json"
"io"
"net/http"
)
type ModelRouteQuery struct {
ProviderCode string
Model string
Endpoint string
APIKeyID string
TenantID string
Seed string
}
type ModelRouteResult struct {
ResolvedAdapter
TargetModel string
Matched bool
Known bool
}
type ModelRouteResolver interface {
ModelRoutingEnabled() bool
ResolveModelRoute(ModelRouteQuery) (ModelRouteResult, error)
}
type modelRequest struct {
body map[string]json.RawMessage
model string
}
func readModelRequest(request *http.Request, maxBody int64) (modelRequest, error) {
if request.Body == nil || request.Method == http.MethodGet {
return modelRequest{}, nil
}
body, err := io.ReadAll(io.LimitReader(request.Body, maxBody+1))
if err != nil {
return modelRequest{}, err
}
_ = request.Body.Close()
if int64(len(body)) > maxBody {
return modelRequest{}, errRequestBodyTooLarge
}
restoreRequestBody(request, body)
var payload map[string]json.RawMessage
if json.Unmarshal(body, &payload) != nil {
return modelRequest{}, nil
}
var model string
_ = json.Unmarshal(payload["model"], &model)
return modelRequest{body: payload, model: model}, nil
}
func (m modelRequest) rewrite(request *http.Request, target string) error {
if len(m.body) == 0 || target == "" || target == m.model {
return nil
}
encodedModel, _ := json.Marshal(target)
m.body["model"] = encodedModel
body, err := json.Marshal(m.body)
if err != nil {
return err
}
restoreRequestBody(request, body)
return nil
}
func restoreRequestBody(request *http.Request, body []byte) {
request.Body = io.NopCloser(bytes.NewReader(body))
request.ContentLength = int64(len(body))
request.GetBody = func() (io.ReadCloser, error) { return io.NopCloser(bytes.NewReader(body)), nil }
}
+129
View File
@@ -0,0 +1,129 @@
package gateway
import (
"context"
"errors"
"fmt"
"strconv"
"time"
"aigateway.local/core/internal/apikey"
"github.com/redis/go-redis/v9"
)
var ErrTokenQuotaUnavailable = errors.New("token quota unavailable")
type TokenReservation struct {
Allowed bool
APIKeyID string
CounterKey string
Reserved int64
Limit int64
Remaining int64
ResetAt time.Time
}
type TokenQuotaController interface {
Reserve(context.Context, apikey.Principal, int64, time.Time) (TokenReservation, error)
Commit(context.Context, TokenReservation, int64) error
}
type RedisTokenQuotaController struct {
client *redis.Client
reserveScript *redis.Script
commitScript *redis.Script
}
func NewRedisTokenQuotaController(client *redis.Client) *RedisTokenQuotaController {
return &RedisTokenQuotaController{
client: client, reserveScript: redis.NewScript(tokenReserveScript), commitScript: redis.NewScript(tokenCommitScript),
}
}
func (c *RedisTokenQuotaController) Reserve(ctx context.Context, principal apikey.Principal, estimate int64, now time.Time) (TokenReservation, error) {
if principal.APIKeyID == "" {
return TokenReservation{Allowed: true}, nil
}
if estimate < 0 {
estimate = 0
}
// Accounts without a monthly token quota never need a reservation and must
// not create a pointless monthly counter key in Redis.
if principal.MonthlyTokenQuota == 0 {
return TokenReservation{Allowed: true}, nil
}
if c == nil || c.client == nil {
return TokenReservation{}, ErrTokenQuotaUnavailable
}
now = now.UTC()
reset := time.Date(now.Year(), now.Month()+1, 1, 0, 0, 0, 0, time.UTC)
key := apikey.MonthlyTokenUsageKey(principal.APIKeyID, now)
result, err := c.reserveScript.Run(ctx, c.client, []string{key}, estimate, principal.MonthlyTokenQuota, int64(reset.Sub(now).Seconds())+86400).Slice()
if err != nil || len(result) != 2 {
if principal.MonthlyTokenQuota == 0 {
return TokenReservation{Allowed: true}, nil
}
return TokenReservation{}, fmt.Errorf("%w: %v", ErrTokenQuotaUnavailable, err)
}
allowed, err := redisInteger(result[0])
if err != nil {
return TokenReservation{}, ErrTokenQuotaUnavailable
}
current, err := redisInteger(result[1])
if err != nil {
return TokenReservation{}, ErrTokenQuotaUnavailable
}
return TokenReservation{
Allowed: allowed == 1, APIKeyID: principal.APIKeyID, CounterKey: key, Reserved: estimate,
Limit: principal.MonthlyTokenQuota, Remaining: max(principal.MonthlyTokenQuota-current, 0), ResetAt: reset,
}, nil
}
func (c *RedisTokenQuotaController) Commit(ctx context.Context, reservation TokenReservation, actual int64) error {
if c == nil || c.client == nil || reservation.CounterKey == "" || !reservation.Allowed {
return nil
}
if actual < 0 {
actual = 0
}
if _, err := c.commitScript.Run(ctx, c.client, []string{reservation.CounterKey}, actual-reservation.Reserved).Result(); err != nil {
return fmt.Errorf("%w: %v", ErrTokenQuotaUnavailable, err)
}
return nil
}
const tokenReserveScript = `
local estimate = tonumber(ARGV[1])
local quota = tonumber(ARGV[2])
local current = tonumber(redis.call('GET', KEYS[1]) or '0')
if quota > 0 and current + estimate > quota then
return {0, current}
end
if estimate > 0 then
current = redis.call('INCRBY', KEYS[1], estimate)
if current == estimate then redis.call('EXPIRE', KEYS[1], tonumber(ARGV[3])) end
elseif redis.call('EXISTS', KEYS[1]) == 0 then
redis.call('SET', KEYS[1], 0, 'EX', tonumber(ARGV[3]))
end
return {1, current}
`
const tokenCommitScript = `
local delta = tonumber(ARGV[1])
local current = tonumber(redis.call('GET', KEYS[1]) or '0')
local updated = current + delta
if updated < 0 then updated = 0 end
redis.call('SET', KEYS[1], updated, 'KEEPTTL')
return updated
`
func writeTokenQuotaHeaders(header mapHeader, reservation TokenReservation) {
if reservation.Limit <= 0 || reservation.ResetAt.IsZero() {
return
}
header.Set("X-TokenLimit-Limit", strconv.FormatInt(reservation.Limit, 10))
header.Set("X-TokenLimit-Remaining", strconv.FormatInt(max(reservation.Remaining, 0), 10))
header.Set("X-TokenLimit-Reset", strconv.FormatInt(reservation.ResetAt.Unix(), 10))
}
type mapHeader interface{ Set(string, string) }
+261
View File
@@ -0,0 +1,261 @@
package gateway
import (
"bytes"
"context"
"encoding/json"
"io"
"log/slog"
"mime"
"net/http"
"strings"
"sync"
"time"
)
// maxUsageDocumentBytes bounds the buffered tail of a non-streaming response.
// The window is kept from the end of the body (where providers place the
// "usage" object), so token accounting stays correct for responses far larger
// than this bound while memory use stays bounded per in-flight response.
const maxUsageDocumentBytes = 2 << 20
type usageSession struct {
controller TokenQuotaController
reservation TokenReservation
fallback int64
onFinish func(TokenUsage)
once sync.Once
log *slog.Logger
}
type TokenUsage struct {
Input int64
Output int64
Total int64
}
type usageSessionContextKey struct{}
func withUsageSession(request *http.Request, session *usageSession) *http.Request {
return request.WithContext(context.WithValue(request.Context(), usageSessionContextKey{}, session))
}
func usageSessionFrom(request *http.Request) *usageSession {
session, _ := request.Context().Value(usageSessionContextKey{}).(*usageSession)
return session
}
func (s *usageSession) finish(usage TokenUsage) {
if s == nil {
return
}
s.once.Do(func() {
usage.Total = max(usage.Total, usage.Input+usage.Output)
if usage.Total <= 0 {
usage.Total = s.fallback
}
if s.controller != nil {
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
if err := s.controller.Commit(ctx, s.reservation, usage.Total); err != nil && s.log != nil {
// The reservation already counted towards the quota, so a
// failed reconciliation silently leaves the counter slightly
// off. Surface it instead of dropping it.
s.log.Warn("token quota commit failed", "error", err)
}
cancel()
}
if s.onFinish != nil {
s.onFinish(usage)
}
})
}
type usageReadCloser struct {
io.ReadCloser
collector *usageCollector
session *usageSession
}
func newUsageReadCloser(body io.ReadCloser, contentType string, session *usageSession) io.ReadCloser {
mediaType, _, _ := mime.ParseMediaType(contentType)
return &usageReadCloser{ReadCloser: body, collector: &usageCollector{sse: mediaType == "text/event-stream"}, session: session}
}
func (r *usageReadCloser) Read(buffer []byte) (int, error) {
n, err := r.ReadCloser.Read(buffer)
if n > 0 {
r.collector.feed(buffer[:n])
}
if err == io.EOF {
r.session.finish(r.collector.usage())
}
return n, err
}
func (r *usageReadCloser) Close() error {
r.session.finish(r.collector.usage())
return r.ReadCloser.Close()
}
type usageCollector struct {
sse bool
pending []byte
doc []byte
input int64
output int64
total int64
}
func (c *usageCollector) feed(chunk []byte) {
if !c.sse {
c.doc = append(c.doc, chunk...)
// Keep only a bounded tail window. The "usage" member lives at the end
// of a non-streaming response, so dropping the head (never the tail)
// preserves accounting for arbitrarily large bodies at a fixed memory
// cost instead of truncating usage away.
if len(c.doc) > maxUsageDocumentBytes {
c.doc = append([]byte(nil), c.doc[len(c.doc)-maxUsageDocumentBytes:]...)
}
return
}
c.pending = append(c.pending, chunk...)
for {
index := bytes.IndexByte(c.pending, '\n')
if index < 0 {
if len(c.pending) > maxUsageDocumentBytes {
c.pending = c.pending[:0]
}
return
}
line := strings.TrimSpace(string(c.pending[:index]))
c.pending = c.pending[index+1:]
if strings.HasPrefix(line, "data:") {
payload := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
if payload != "" && payload != "[DONE]" {
c.consumeJSON([]byte(payload))
}
}
}
}
func (c *usageCollector) tokens() int64 {
return c.usage().Total
}
func (c *usageCollector) usage() TokenUsage {
if c.sse && len(c.pending) > 0 {
line := strings.TrimSpace(string(c.pending))
if strings.HasPrefix(line, "data:") {
c.consumeJSON([]byte(strings.TrimSpace(strings.TrimPrefix(line, "data:"))))
}
c.pending = nil
}
if !c.sse && len(c.doc) > 0 {
c.consumeJSON(c.doc) // fast path: whole valid JSON object
c.consumeUsageObject(c.doc) // tail extraction: covers truncated bodies
}
return TokenUsage{Input: c.input, Output: c.output, Total: max(c.total, c.input+c.output)}
}
// consumeUsageObject extracts the trailing `"usage"` object from a possibly
// truncated response document. It scans for the final `"usage"` member and
// decodes the JSON object that immediately follows — where OpenAI-compatible
// providers place token accounting in non-streaming responses — so usage is
// still counted when the buffered window starts mid-object.
func (c *usageCollector) consumeUsageObject(doc []byte) {
const key = `"usage"`
index := bytes.LastIndex(doc, []byte(key))
if index < 0 {
return
}
rest := doc[index+len(key):]
colon := bytes.IndexByte(rest, ':')
if colon < 0 {
return
}
rest = bytes.TrimSpace(rest[colon+1:])
if len(rest) == 0 || rest[0] != '{' {
return
}
depth := 0
end := -1
inString := false
escaped := false
for i := 0; i < len(rest); i++ {
ch := rest[i]
if inString {
if escaped {
escaped = false
} else if ch == '\\' {
escaped = true
} else if ch == '"' {
inString = false
}
continue
}
switch ch {
case '"':
inString = true
case '{':
depth++
case '}':
depth--
if depth == 0 {
end = i + 1
}
}
if end > 0 {
break
}
}
if end > 0 {
c.consumeJSON(rest[:end])
}
}
func (c *usageCollector) consumeJSON(payload []byte) {
decoder := json.NewDecoder(bytes.NewReader(payload))
decoder.UseNumber()
var value any
if decoder.Decode(&value) == nil {
c.walk(value, false)
}
}
func (c *usageCollector) walk(value any, inUsage bool) {
switch typed := value.(type) {
case map[string]any:
for key, child := range typed {
usage := inUsage || key == "usage"
if usage {
switch key {
case "input_tokens", "prompt_tokens":
c.input = max(c.input, jsonInt64(child))
case "output_tokens", "completion_tokens":
c.output = max(c.output, jsonInt64(child))
case "total_tokens":
c.total = max(c.total, jsonInt64(child))
}
}
c.walk(child, usage)
}
case []any:
for _, child := range typed {
c.walk(child, inUsage)
}
}
}
func jsonInt64(value any) int64 {
switch number := value.(type) {
case json.Number:
parsed, _ := number.Int64()
return max(parsed, 0)
case float64:
return max(int64(number), 0)
case int64:
return max(number, 0)
default:
return 0
}
}
+21
View File
@@ -0,0 +1,21 @@
package gateway
import "testing"
func TestUsageCollectorReadsOpenAIJSON(t *testing.T) {
collector := &usageCollector{}
collector.feed([]byte(`{"usage":{"prompt_tokens":11,"completion_tokens":7,"total_tokens":18}}`))
if got := collector.tokens(); got != 18 {
t.Fatalf("got %d, want 18", got)
}
}
func TestUsageCollectorReadsChunkedSSEAndAnthropicUsage(t *testing.T) {
collector := &usageCollector{sse: true}
collector.feed([]byte("event: message_start\ndata: {\"message\":{\"usage\":{\"input_tokens\":13}}}\n\n"))
collector.feed([]byte("event: message_delta\ndata: {\"usage\":{\"output_"))
collector.feed([]byte("tokens\":9}}\n\ndata: [DONE]\n\n"))
if got := collector.tokens(); got != 22 {
t.Fatalf("got %d, want 22", got)
}
}
+128
View File
@@ -0,0 +1,128 @@
package identity
import (
"errors"
"sort"
"time"
)
var (
ErrNotFound = errors.New("identity not found")
ErrUnavailable = errors.New("identity service unavailable")
)
type Kind string
const (
KindAdmin Kind = "admin"
KindPortal Kind = "portal"
)
type Account struct {
ID string
Kind Kind
Login string
DisplayName string
Role string
Permissions []string
PasswordHash string
AuthSource string
Active bool
FailedLogins int
LockedUntil *time.Time
TOTPEnabled bool
EncryptedTOTPSecret []byte
TOTPKekVersion *int
TOTPLastStep *int64
TOTPBackupCodes []byte
DepartmentID *string
DepartmentName string
CreatedAt time.Time
UpdatedAt time.Time
}
const (
PermissionIdentityManage = "identity:manage"
PermissionProviderRead = "provider:read"
PermissionProviderManage = "provider:manage"
PermissionAPIKeyRead = "api_key:read"
PermissionAPIKeyManage = "api_key:manage"
PermissionAuditRead = "audit:read"
PermissionUsageRead = "usage:read"
PermissionOutboxRead = "outbox:read"
PermissionOutboxManage = "outbox:manage"
PermissionContentPolicyRead = "content_policy:read"
PermissionContentPolicyManage = "content_policy:manage"
PermissionPricingRead = "pricing:read"
PermissionPricingManage = "pricing:manage"
PermissionPromptRead = "prompt:read"
PermissionPromptManage = "prompt:manage"
PermissionKnowledgeRead = "knowledge:read"
PermissionKnowledgeManage = "knowledge:manage"
PermissionToolRead = "tool:read"
PermissionToolManage = "tool:manage"
PermissionApplicationRead = "application:read"
PermissionApplicationManage = "application:manage"
PermissionNotificationRead = "notification:read"
PermissionNotificationManage = "notification:manage"
PermissionMCPServerRead = "mcp_server:read"
PermissionMCPServerManage = "mcp_server:manage"
PermissionSkillRead = "skill:read"
PermissionSkillManage = "skill:manage"
PermissionDigitalEmployeeRead = "digital_employee:read"
PermissionDigitalEmployeeManage = "digital_employee:manage"
PermissionMarketplaceRead = "marketplace:read"
PermissionMarketplaceManage = "marketplace:manage"
)
var rolePermissions = map[string][]string{
"superadmin": {"*"},
"operator": {
PermissionProviderRead, PermissionProviderManage,
PermissionAPIKeyRead, PermissionAPIKeyManage,
PermissionUsageRead,
PermissionOutboxRead, PermissionOutboxManage,
PermissionContentPolicyRead, PermissionContentPolicyManage,
PermissionPricingRead, PermissionPricingManage,
PermissionPromptRead, PermissionPromptManage,
PermissionKnowledgeRead, PermissionKnowledgeManage,
PermissionToolRead, PermissionToolManage,
PermissionApplicationRead, PermissionApplicationManage,
PermissionNotificationRead, PermissionNotificationManage,
PermissionMCPServerRead, PermissionMCPServerManage,
PermissionSkillRead, PermissionSkillManage,
PermissionDigitalEmployeeRead, PermissionDigitalEmployeeManage,
PermissionMarketplaceRead, PermissionMarketplaceManage,
},
"auditor": {PermissionProviderRead, PermissionAPIKeyRead, PermissionAuditRead, PermissionUsageRead, PermissionOutboxRead, PermissionContentPolicyRead, PermissionPricingRead, PermissionPromptRead, PermissionKnowledgeRead, PermissionToolRead, PermissionApplicationRead, PermissionNotificationRead, PermissionMCPServerRead, PermissionSkillRead, PermissionDigitalEmployeeRead, PermissionMarketplaceRead},
"member": {},
}
func EffectivePermissions(account Account) []string {
seen := make(map[string]struct{})
result := make([]string, 0, len(account.Permissions)+4)
for _, permissions := range [][]string{rolePermissions[account.Role], account.Permissions} {
for _, permission := range permissions {
if _, exists := seen[permission]; exists {
continue
}
seen[permission] = struct{}{}
result = append(result, permission)
}
}
sort.Strings(result)
return result
}
func HasPermission(account Account, required string) bool {
for _, permission := range EffectivePermissions(account) {
if permission == "*" || permission == required {
return true
}
}
return false
}
func (a Account) Locked(now time.Time) bool {
return a.LockedUntil != nil && a.LockedUntil.After(now)
}
+297
View File
@@ -0,0 +1,297 @@
package identity
import (
"context"
"encoding/json"
"errors"
"fmt"
"net/http"
"regexp"
"strings"
"time"
"aigateway.local/core/internal/platform/apiresponse"
platformid "aigateway.local/core/internal/platform/id"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgconn"
)
var (
ErrDepartmentConflict = errors.New("department already exists")
ErrDepartmentInUse = errors.New("department is in use")
ErrDepartmentCycle = errors.New("department hierarchy cycle")
departmentCodePattern = regexp.MustCompile(`^[a-z][a-z0-9_-]{1,63}$`)
)
type Department struct {
ID string `json:"id"`
Code string `json:"code"`
Name string `json:"name"`
Description string `json:"description"`
ParentID *string `json:"parent_id"`
ParentName string `json:"parent_name,omitempty"`
Active bool `json:"active"`
UserCount int `json:"user_count"`
CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"`
}
type departmentInput struct {
Code string `json:"code"`
Name string `json:"name"`
Description string `json:"description"`
ParentID *string `json:"parent_id"`
Active *bool `json:"active"`
}
func (h *ManagementHTTPHandler) listDepartments(writer http.ResponseWriter, request *http.Request) {
if _, ok := h.requirePermission(writer, request); !ok {
return
}
departments, err := h.service.repository.ListDepartments(request.Context())
if err != nil {
h.writeDepartmentError(writer, err)
return
}
apiresponse.OK(writer, departments)
}
func (h *ManagementHTTPHandler) createDepartment(writer http.ResponseWriter, request *http.Request) {
actor, ok := h.requirePermission(writer, request)
if !ok {
return
}
input, department, ok := decodeDepartment(writer, request)
if !ok {
return
}
_ = input
created, err := h.service.repository.CreateDepartment(request.Context(), department, actor.ID)
if err != nil {
h.writeDepartmentError(writer, err)
return
}
apiresponse.OK(writer, created)
}
func (h *ManagementHTTPHandler) updateDepartment(writer http.ResponseWriter, request *http.Request) {
actor, ok := h.requirePermission(writer, request)
if !ok {
return
}
_, department, ok := decodeDepartment(writer, request)
if !ok {
return
}
department.ID = request.PathValue("department_id")
updated, err := h.service.repository.UpdateDepartment(request.Context(), department, actor.ID)
if err != nil {
h.writeDepartmentError(writer, err)
return
}
apiresponse.OK(writer, updated)
}
func decodeDepartment(writer http.ResponseWriter, request *http.Request) (departmentInput, Department, bool) {
var input departmentInput
decoder := json.NewDecoder(http.MaxBytesReader(writer, request.Body, 1<<20))
decoder.DisallowUnknownFields()
if err := decoder.Decode(&input); err != nil {
apiresponse.Error(writer, http.StatusBadRequest, "请求格式无效")
return input, Department{}, false
}
input.Code = strings.ToLower(strings.TrimSpace(input.Code))
input.Name = strings.TrimSpace(input.Name)
input.Description = strings.TrimSpace(input.Description)
if !departmentCodePattern.MatchString(input.Code) || input.Name == "" || len(input.Name) > 128 || len(input.Description) > 1024 {
apiresponse.Error(writer, http.StatusBadRequest, "部门代码、名称或描述格式无效")
return input, Department{}, false
}
var parentID *string
if input.ParentID != nil && strings.TrimSpace(*input.ParentID) != "" {
value := strings.TrimSpace(*input.ParentID)
parentID = &value
}
active := true
if input.Active != nil {
active = *input.Active
}
return input, Department{Code: input.Code, Name: input.Name, Description: input.Description, ParentID: parentID, Active: active}, true
}
func (h *ManagementHTTPHandler) writeDepartmentError(writer http.ResponseWriter, err error) {
switch {
case errors.Is(err, ErrNotFound):
apiresponse.Error(writer, http.StatusNotFound, "部门不存在")
case errors.Is(err, ErrDepartmentConflict):
apiresponse.Error(writer, http.StatusConflict, "部门代码已存在")
case errors.Is(err, ErrDepartmentCycle):
apiresponse.Error(writer, http.StatusConflict, "部门层级不能形成循环")
case errors.Is(err, ErrDepartmentInUse):
apiresponse.Error(writer, http.StatusConflict, "部门仍包含启用用户或启用子部门,不能停用")
case errors.Is(err, ErrUnavailable):
apiresponse.Error(writer, http.StatusServiceUnavailable, "部门服务暂不可用")
default:
apiresponse.Error(writer, http.StatusBadRequest, "部门操作失败")
}
}
func (r *Repository) ListDepartments(ctx context.Context) ([]Department, error) {
if r.pool == nil {
return nil, ErrUnavailable
}
rows, err := r.pool.Query(ctx, `
SELECT d.id::text, d.code, d.name, d.description, d.parent_id::text,
COALESCE(p.name, ''), d.active,
count(u.id) FILTER (WHERE u.active), d.created_at, d.updated_at
FROM gateway.departments d
LEFT JOIN gateway.departments p ON p.id = d.parent_id
LEFT JOIN gateway.portal_users u ON u.department_id = d.id
GROUP BY d.id, p.name
ORDER BY d.code`)
if err != nil {
return nil, fmt.Errorf("%w: %v", ErrUnavailable, err)
}
defer rows.Close()
departments := make([]Department, 0)
for rows.Next() {
var department Department
if err := rows.Scan(&department.ID, &department.Code, &department.Name, &department.Description,
&department.ParentID, &department.ParentName, &department.Active, &department.UserCount,
&department.CreatedAt, &department.UpdatedAt); err != nil {
return nil, fmt.Errorf("%w: %v", ErrUnavailable, err)
}
departments = append(departments, department)
}
return departments, mapRepositoryError(rows.Err())
}
func (r *Repository) GetDepartment(ctx context.Context, id string) (Department, error) {
if r.pool == nil {
return Department{}, ErrUnavailable
}
var department Department
err := r.pool.QueryRow(ctx, `
SELECT id::text, code, name, description, parent_id::text, active, created_at, updated_at
FROM gateway.departments WHERE id = $1`, id).Scan(
&department.ID, &department.Code, &department.Name, &department.Description,
&department.ParentID, &department.Active, &department.CreatedAt, &department.UpdatedAt,
)
return department, mapRepositoryError(err)
}
func (r *Repository) CreateDepartment(ctx context.Context, department Department, actorID string) (Department, error) {
id, err := platformid.NewUUID()
if err != nil {
return Department{}, err
}
department.ID = id
return r.storeDepartment(ctx, department, actorID, true)
}
func (r *Repository) UpdateDepartment(ctx context.Context, department Department, actorID string) (Department, error) {
return r.storeDepartment(ctx, department, actorID, false)
}
func (r *Repository) storeDepartment(ctx context.Context, department Department, actorID string, creating bool) (Department, error) {
if r.pool == nil {
return Department{}, ErrUnavailable
}
tx, err := r.pool.Begin(ctx)
if err != nil {
return Department{}, fmt.Errorf("%w: %v", ErrUnavailable, err)
}
defer func() { _ = tx.Rollback(ctx) }()
if department.ParentID != nil {
var parentActive bool
if err := tx.QueryRow(ctx, `SELECT active FROM gateway.departments WHERE id = $1`, *department.ParentID).Scan(&parentActive); err != nil {
return Department{}, mapRepositoryError(err)
}
if !parentActive {
return Department{}, ErrDepartmentInUse
}
}
if !creating && department.ParentID != nil {
var cycle bool
if err := tx.QueryRow(ctx, `
WITH RECURSIVE descendants AS (
SELECT id FROM gateway.departments WHERE parent_id = $1
UNION ALL
SELECT d.id FROM gateway.departments d JOIN descendants x ON d.parent_id = x.id
)
SELECT $2::uuid = $1::uuid OR EXISTS (SELECT 1 FROM descendants WHERE id = $2)`,
department.ID, *department.ParentID).Scan(&cycle); err != nil {
return Department{}, fmt.Errorf("%w: %v", ErrUnavailable, err)
}
if cycle {
return Department{}, ErrDepartmentCycle
}
}
if !creating && !department.Active {
var inUse bool
if err := tx.QueryRow(ctx, `
SELECT EXISTS (SELECT 1 FROM gateway.portal_users WHERE department_id = $1 AND active)
OR EXISTS (SELECT 1 FROM gateway.departments WHERE parent_id = $1 AND active)`, department.ID).Scan(&inUse); err != nil {
return Department{}, fmt.Errorf("%w: %v", ErrUnavailable, err)
}
if inUse {
return Department{}, ErrDepartmentInUse
}
}
if creating {
err = tx.QueryRow(ctx, `
INSERT INTO gateway.departments (id, code, name, description, parent_id, active)
VALUES ($1, $2, $3, $4, $5, $6)
RETURNING created_at, updated_at`, department.ID, department.Code, department.Name,
department.Description, department.ParentID, department.Active).Scan(&department.CreatedAt, &department.UpdatedAt)
} else {
err = tx.QueryRow(ctx, `
UPDATE gateway.departments
SET code = $2, name = $3, description = $4, parent_id = $5,
active = $6, updated_at = clock_timestamp()
WHERE id = $1
RETURNING created_at, updated_at`, department.ID, department.Code, department.Name,
department.Description, department.ParentID, department.Active).Scan(&department.CreatedAt, &department.UpdatedAt)
}
if err != nil {
return Department{}, mapDepartmentError(err)
}
eventID, err := platformid.NewUUID()
if err != nil {
return Department{}, err
}
eventType := "department.updated"
if creating {
eventType = "department.created"
}
payload, _ := json.Marshal(map[string]any{"department_id": department.ID, "code": department.Code, "actor_id": actorID})
if _, err := tx.Exec(ctx, `
INSERT INTO gateway.outbox_events
(event_id, event_type, event_version, aggregate_type, aggregate_id, payload)
VALUES ($1, $2, 1, 'department', $3, $4)`, eventID, eventType, department.ID, payload); err != nil {
return Department{}, fmt.Errorf("%w: %v", ErrUnavailable, err)
}
if err := tx.Commit(ctx); err != nil {
return Department{}, fmt.Errorf("%w: %v", ErrUnavailable, err)
}
return department, nil
}
func mapDepartmentError(err error) error {
if errors.Is(err, pgx.ErrNoRows) {
return ErrNotFound
}
var pgError *pgconn.PgError
if errors.As(err, &pgError) {
switch pgError.Code {
case "23505":
return ErrDepartmentConflict
case "23503", "23514":
return ErrDepartmentCycle
}
}
if err != nil {
return fmt.Errorf("%w: %v", ErrUnavailable, err)
}
return nil
}
+429
View File
@@ -0,0 +1,429 @@
package identity
import (
"encoding/json"
"errors"
"net/http"
"strings"
"time"
"aigateway.local/core/internal/platform/apiresponse"
"aigateway.local/core/internal/platform/cryptox"
)
type HTTPHandler struct {
service *Service
mux *http.ServeMux
}
type loginRequest struct {
UserName string `json:"userName"`
Username string `json:"username"`
Account string `json:"account"`
Password string `json:"password"`
}
type totpLoginRequest struct {
TempToken string `json:"temp_token"`
Code string `json:"code"`
BackupCode string `json:"backup_code"`
}
type passwordRequest struct {
Password string `json:"password"`
}
type factorRequest struct {
Password string `json:"password"`
Code string `json:"code"`
BackupCode string `json:"backup_code"`
}
func NewHTTPHandler(service *Service) *HTTPHandler {
handler := &HTTPHandler{service: service, mux: http.NewServeMux()}
handler.mux.HandleFunc("POST /api/v1/admin/login", handler.login(KindAdmin))
handler.registerTOTP(KindAdmin, "/api/v1/admin")
handler.mux.HandleFunc("GET /api/v1/admin/whoami", handler.whoami(KindAdmin))
handler.mux.HandleFunc("POST /api/v1/admin/password", handler.changePassword(KindAdmin))
handler.mux.HandleFunc("POST /api/v1/admin/logout", handler.logout)
handler.mux.HandleFunc("GET /api/v1/admin/menus", handler.menus(KindAdmin))
handler.mux.HandleFunc("POST /api/v1/portal/login", handler.login(KindPortal))
handler.registerTOTP(KindPortal, "/api/v1/portal")
handler.mux.HandleFunc("GET /api/v1/portal/me", handler.whoami(KindPortal))
handler.mux.HandleFunc("POST /api/v1/portal/logout", handler.logout)
handler.mux.HandleFunc("GET /api/v1/portal/menus", handler.menus(KindPortal))
handler.registerOIDC()
return handler
}
func (h *HTTPHandler) changePassword(kind Kind) http.HandlerFunc {
return func(writer http.ResponseWriter, request *http.Request) {
account, ok := h.requireAccount(writer, request, kind)
if !ok {
return
}
var input struct {
OldPassword string `json:"old_password"`
NewPassword string `json:"new_password"`
}
if !decodeJSON(writer, request, &input) {
apiresponse.Error(writer, http.StatusBadRequest, "请求格式无效")
return
}
if err := h.service.ChangePassword(request.Context(), account, input.OldPassword, input.NewPassword); err != nil {
h.writeIdentityError(writer, err)
return
}
apiresponse.OK(writer, map[string]bool{"changed": true})
}
}
func (h *HTTPHandler) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
h.mux.ServeHTTP(writer, request)
}
func (h *HTTPHandler) login(kind Kind) http.HandlerFunc {
return func(writer http.ResponseWriter, request *http.Request) {
// 防爆破:按 IP 的滑动窗口限流,超限返回 429(与账号锁定叠加)。
if !h.service.AllowLogin(request.Context(), ClientIP(request)) {
apiresponse.Error(writer, http.StatusTooManyRequests, "登录尝试过于频繁,请稍后再试")
return
}
var input loginRequest
decoder := json.NewDecoder(http.MaxBytesReader(writer, request.Body, 1<<20))
decoder.DisallowUnknownFields()
if err := decoder.Decode(&input); err != nil {
apiresponse.Error(writer, http.StatusBadRequest, "请求格式无效")
return
}
login := strings.TrimSpace(input.UserName)
if login == "" {
login = strings.TrimSpace(input.Username)
}
if login == "" {
login = strings.TrimSpace(input.Account)
}
if login == "" || len(login) > 128 || len(input.Password) < 1 || len(input.Password) > 1024 {
apiresponse.Error(writer, http.StatusBadRequest, "账号或口令格式无效")
return
}
result, err := h.service.Login(request.Context(), kind, login, input.Password)
if err != nil {
h.writeIdentityError(writer, err)
return
}
apiresponse.OK(writer, map[string]any{
"token": result.Token, "refreshToken": "", "require_totp": result.RequireTOTP,
"temp_token": result.TempToken,
})
}
}
func (h *HTTPHandler) registerTOTP(kind Kind, prefix string) {
h.mux.HandleFunc("POST "+prefix+"/login/totp", h.completeTOTPLogin(kind))
h.mux.HandleFunc("GET "+prefix+"/totp/status", h.totpStatus(kind))
h.mux.HandleFunc("POST "+prefix+"/totp/setup", h.setupTOTP(kind))
h.mux.HandleFunc("POST "+prefix+"/totp/confirm", h.confirmTOTP(kind))
h.mux.HandleFunc("POST "+prefix+"/totp/disable", h.disableTOTP(kind))
h.mux.HandleFunc("POST "+prefix+"/totp/backup-codes/regenerate", h.regenerateBackupCodes(kind))
}
func (h *HTTPHandler) completeTOTPLogin(kind Kind) http.HandlerFunc {
return func(writer http.ResponseWriter, request *http.Request) {
// 防爆破:TOTP 完成端点同样按 IP 限流。
if !h.service.AllowLogin(request.Context(), ClientIP(request)) {
apiresponse.Error(writer, http.StatusTooManyRequests, "登录尝试过于频繁,请稍后再试")
return
}
var input totpLoginRequest
if !decodeJSON(writer, request, &input) || strings.TrimSpace(input.TempToken) == "" || (strings.TrimSpace(input.Code) == "" && strings.TrimSpace(input.BackupCode) == "") {
apiresponse.Error(writer, http.StatusBadRequest, "请输入动态验证码或备用码")
return
}
result, err := h.service.CompleteTOTPLogin(request.Context(), kind, input.TempToken, input.Code, input.BackupCode)
if err != nil {
h.writeIdentityError(writer, err)
return
}
apiresponse.OK(writer, map[string]any{"token": result.Token, "refreshToken": "", "require_totp": false})
}
}
func (h *HTTPHandler) totpStatus(kind Kind) http.HandlerFunc {
return func(writer http.ResponseWriter, request *http.Request) {
account, ok := h.requireAccount(writer, request, kind)
if !ok {
return
}
apiresponse.OK(writer, map[string]bool{"enabled": account.TOTPEnabled, "setup_pending": !account.TOTPEnabled && len(account.EncryptedTOTPSecret) > 0})
}
}
func (h *HTTPHandler) setupTOTP(kind Kind) http.HandlerFunc {
return func(writer http.ResponseWriter, request *http.Request) {
account, ok := h.requireAccount(writer, request, kind)
if !ok {
return
}
var input passwordRequest
if !decodeJSON(writer, request, &input) || input.Password == "" {
apiresponse.Error(writer, http.StatusBadRequest, "请输入当前口令")
return
}
result, err := h.service.SetupTOTP(request.Context(), account, input.Password)
if err != nil {
h.writeIdentityError(writer, err)
return
}
apiresponse.OK(writer, map[string]string{"secret": result.Secret, "provisioning_uri": result.ProvisioningURI})
}
}
func (h *HTTPHandler) confirmTOTP(kind Kind) http.HandlerFunc {
return func(writer http.ResponseWriter, request *http.Request) {
account, ok := h.requireAccount(writer, request, kind)
if !ok {
return
}
var input factorRequest
if !decodeJSON(writer, request, &input) || strings.TrimSpace(input.Code) == "" {
apiresponse.Error(writer, http.StatusBadRequest, "请输入动态验证码")
return
}
codes, err := h.service.ConfirmTOTP(request.Context(), account, input.Code)
if err != nil {
h.writeIdentityError(writer, err)
return
}
apiresponse.OK(writer, map[string]any{"enabled": true, "backup_codes": codes})
}
}
func (h *HTTPHandler) disableTOTP(kind Kind) http.HandlerFunc {
return func(writer http.ResponseWriter, request *http.Request) {
account, ok := h.requireAccount(writer, request, kind)
if !ok {
return
}
var input factorRequest
if !decodeJSON(writer, request, &input) || input.Password == "" {
apiresponse.Error(writer, http.StatusBadRequest, "当前口令和验证因子不能为空")
return
}
if err := h.service.DisableTOTP(request.Context(), account, input.Password, input.Code, input.BackupCode); err != nil {
h.writeIdentityError(writer, err)
return
}
apiresponse.OK(writer, map[string]bool{"enabled": false})
}
}
func (h *HTTPHandler) regenerateBackupCodes(kind Kind) http.HandlerFunc {
return func(writer http.ResponseWriter, request *http.Request) {
account, ok := h.requireAccount(writer, request, kind)
if !ok {
return
}
var input factorRequest
if !decodeJSON(writer, request, &input) || input.Password == "" {
apiresponse.Error(writer, http.StatusBadRequest, "当前口令和验证因子不能为空")
return
}
codes, err := h.service.RegenerateBackupCodes(request.Context(), account, input.Password, input.Code, input.BackupCode)
if err != nil {
h.writeIdentityError(writer, err)
return
}
apiresponse.OK(writer, map[string]any{"backup_codes": codes})
}
}
func (h *HTTPHandler) requireAccount(writer http.ResponseWriter, request *http.Request, kind Kind) (Account, bool) {
account, err := h.service.Authenticate(request.Context(), kind, request.Header.Get("Authorization"))
if err != nil {
h.writeIdentityError(writer, err)
return Account{}, false
}
return account, true
}
func decodeJSON(writer http.ResponseWriter, request *http.Request, target any) bool {
decoder := json.NewDecoder(http.MaxBytesReader(writer, request.Body, 1<<20))
decoder.DisallowUnknownFields()
return decoder.Decode(target) == nil
}
func (h *HTTPHandler) whoami(kind Kind) http.HandlerFunc {
return func(writer http.ResponseWriter, request *http.Request) {
account, err := h.service.Authenticate(request.Context(), kind, request.Header.Get("Authorization"))
if err != nil {
h.writeIdentityError(writer, err)
return
}
roles := []string{"R_USER"}
if kind == KindAdmin {
roles = []string{"R_ADMIN"}
if account.Role == "superadmin" {
roles = []string{"R_SUPER"}
}
}
permissions := EffectivePermissions(account)
apiresponse.OK(writer, map[string]any{
"userId": account.ID, "userName": account.Login,
"displayName": account.DisplayName, "email": "",
"roles": roles, "buttons": permissions, "permissions": permissions, "role": account.Role,
})
}
}
func (h *HTTPHandler) logout(writer http.ResponseWriter, request *http.Request) {
if err := h.service.Logout(request.Context(), request.Header.Get("Authorization")); err != nil && !errors.Is(err, ErrInvalidSession) {
h.writeIdentityError(writer, err)
return
}
apiresponse.OK(writer, map[string]bool{"ok": true})
}
func (h *HTTPHandler) menus(kind Kind) http.HandlerFunc {
return func(writer http.ResponseWriter, request *http.Request) {
account, err := h.service.Authenticate(request.Context(), kind, request.Header.Get("Authorization"))
if err != nil {
h.writeIdentityError(writer, err)
return
}
if kind == KindAdmin {
apiresponse.OK(writer, adminMenus(account))
return
}
apiresponse.OK(writer, portalMenus())
}
}
func (h *HTTPHandler) writeIdentityError(writer http.ResponseWriter, err error) {
var locked LockedError
switch {
case errors.As(err, &locked):
minutes := int(time.Until(locked.Until).Minutes()) + 1
apiresponse.Error(writer, http.StatusTooManyRequests, "账号已锁定,请在 "+(time.Duration(minutes)*time.Minute).String()+" 后重试")
case errors.Is(err, ErrInvalidCredentials):
apiresponse.Error(writer, http.StatusUnauthorized, "账号或口令错误")
case errors.Is(err, ErrInvalidSession), errors.Is(err, ErrNotFound):
apiresponse.Error(writer, http.StatusUnauthorized, "登录状态无效或已过期")
case errors.Is(err, ErrAccountDisabled):
apiresponse.Error(writer, http.StatusForbidden, "账号已被停用")
case errors.Is(err, ErrInvalidTOTP):
apiresponse.Error(writer, http.StatusUnauthorized, "动态验证码无效、已使用或备用码无效")
case errors.Is(err, ErrTOTPAlreadyEnabled):
apiresponse.Error(writer, http.StatusConflict, "两步验证已经启用")
case errors.Is(err, ErrTOTPNotEnabled), errors.Is(err, ErrTOTPSetupRequired):
apiresponse.Error(writer, http.StatusConflict, "两步验证尚未完成配置")
case errors.Is(err, cryptox.ErrKeyUnavailable):
apiresponse.Error(writer, http.StatusServiceUnavailable, "两步验证加密密钥不可用")
case errors.Is(err, ErrUnavailable):
apiresponse.Error(writer, http.StatusServiceUnavailable, "身份服务暂不可用")
default:
apiresponse.Error(writer, http.StatusInternalServerError, "身份服务处理失败")
}
}
func adminMenus(account Account) []map[string]any {
// 运行概览:首页仪表盘(叶子菜单,避免出现「运行概览>运行概览」同级冗余)。
menus := []map[string]any{
{"name": "Dashboard", "path": "/dashboard/console", "component": "/dashboard/console", "meta": map[string]any{"title": "运行概览", "icon": "ri:pie-chart-line", "fixedTab": true}},
}
// 网关接入:上游供应商、路由与凭据。
gatewayChildren := make([]map[string]any, 0, 4)
if HasPermission(account, PermissionProviderRead) || HasPermission(account, PermissionProviderManage) {
gatewayChildren = append(gatewayChildren, map[string]any{"name": "Providers", "path": "providers", "component": "/gateway/providers", "meta": map[string]any{"title": "模型供应商"}})
gatewayChildren = append(gatewayChildren, map[string]any{"name": "ModelRoutes", "path": "model-routes", "component": "/gateway/model-routes", "meta": map[string]any{"title": "模型路由"}})
}
if HasPermission(account, PermissionAPIKeyRead) || HasPermission(account, PermissionAPIKeyManage) {
gatewayChildren = append(gatewayChildren, map[string]any{"name": "APIKeys", "path": "api-keys", "component": "/gateway/api-keys", "meta": map[string]any{"title": "API Key"}})
}
if HasPermission(account, PermissionPricingRead) || HasPermission(account, PermissionPricingManage) {
gatewayChildren = append(gatewayChildren, map[string]any{"name": "ModelPrices", "path": "model-prices", "component": "/gateway/model-prices", "meta": map[string]any{"title": "模型价格"}})
}
if len(gatewayChildren) > 0 {
menus = append(menus, map[string]any{"name": "Gateway", "path": "/gateway", "component": "/index/index", "meta": map[string]any{"title": "网关接入", "icon": "ri:router-line"}, "children": gatewayChildren})
}
// 安全与审计:审计用量、内容策略与模型治理。
securityChildren := make([]map[string]any, 0, 3)
if HasPermission(account, PermissionAuditRead) || HasPermission(account, PermissionUsageRead) {
securityChildren = append(securityChildren, map[string]any{"name": "AuditUsage", "path": "audit-usage", "component": "/gateway/audit-usage", "meta": map[string]any{"title": "审计与用量"}})
}
if HasPermission(account, PermissionContentPolicyRead) || HasPermission(account, PermissionContentPolicyManage) {
securityChildren = append(securityChildren, map[string]any{"name": "ContentPolicies", "path": "content-policies", "component": "/gateway/content-policies", "meta": map[string]any{"title": "内容策略"}})
}
if HasPermission(account, PermissionKnowledgeRead) || HasPermission(account, PermissionKnowledgeManage) {
securityChildren = append(securityChildren, map[string]any{"name": "Governance", "path": "governance", "component": "/gateway/governance", "meta": map[string]any{"title": "模型治理"}})
}
if len(securityChildren) > 0 {
menus = append(menus, map[string]any{"name": "Security", "path": "/security", "component": "/index/index", "meta": map[string]any{"title": "安全与审计", "icon": "ri:shield-check-line"}, "children": securityChildren})
}
// AI 资产:Prompt、知识库、工具与 AI 应用。
assetsChildren := make([]map[string]any, 0, 4)
if HasPermission(account, PermissionPromptRead) || HasPermission(account, PermissionPromptManage) {
assetsChildren = append(assetsChildren, map[string]any{"name": "Prompts", "path": "prompts", "component": "/gateway/prompts", "meta": map[string]any{"title": "Prompt 资产"}})
}
if HasPermission(account, PermissionKnowledgeRead) || HasPermission(account, PermissionKnowledgeManage) {
assetsChildren = append(assetsChildren, map[string]any{"name": "Knowledge", "path": "knowledge", "component": "/gateway/knowledge", "meta": map[string]any{"title": "知识库"}})
}
if HasPermission(account, PermissionToolRead) || HasPermission(account, PermissionToolManage) {
assetsChildren = append(assetsChildren, map[string]any{"name": "Tools", "path": "tools", "component": "/gateway/tools", "meta": map[string]any{"title": "工具中心"}})
}
if HasPermission(account, PermissionApplicationRead) || HasPermission(account, PermissionApplicationManage) {
assetsChildren = append(assetsChildren, map[string]any{"name": "Applications", "path": "applications", "component": "/gateway/applications", "meta": map[string]any{"title": "AI 应用"}})
}
if len(assetsChildren) > 0 {
menus = append(menus, map[string]any{"name": "Assets", "path": "/assets", "component": "/index/index", "meta": map[string]any{"title": "AI 资产", "icon": "ri:box-3-line"}, "children": assetsChildren})
}
// 资源市场:MCP 服务器、Skills 与数字员工(旗舰版资源市场)。
marketChildren := make([]map[string]any, 0, 4)
if HasPermission(account, PermissionMarketplaceRead) || HasPermission(account, PermissionMarketplaceManage) {
marketChildren = append(marketChildren, map[string]any{"name": "Marketplace", "path": "overview", "component": "/gateway/marketplace", "meta": map[string]any{"title": "市场总览"}})
}
if HasPermission(account, PermissionMCPServerRead) || HasPermission(account, PermissionMCPServerManage) {
marketChildren = append(marketChildren, map[string]any{"name": "MCPServers", "path": "mcp-servers", "component": "/gateway/mcp-servers", "meta": map[string]any{"title": "MCP 服务器"}})
}
if HasPermission(account, PermissionSkillRead) || HasPermission(account, PermissionSkillManage) {
marketChildren = append(marketChildren, map[string]any{"name": "Skills", "path": "skills", "component": "/gateway/skills", "meta": map[string]any{"title": "Skills 技能"}})
}
if HasPermission(account, PermissionDigitalEmployeeRead) || HasPermission(account, PermissionDigitalEmployeeManage) {
marketChildren = append(marketChildren, map[string]any{"name": "DigitalEmployees", "path": "digital-employees", "component": "/gateway/digital-employees", "meta": map[string]any{"title": "数字员工"}})
}
if len(marketChildren) > 0 {
menus = append(menus, map[string]any{"name": "ResourceMarket", "path": "/resource-market", "component": "/index/index", "meta": map[string]any{"title": "资源市场", "icon": "ri:store-3-line"}, "children": marketChildren})
}
// 系统管理:账号权限、事件投递与通知。
systemChildren := make([]map[string]any, 0, 3)
if HasPermission(account, PermissionIdentityManage) {
systemChildren = append(systemChildren, map[string]any{"name": "User", "path": "user", "component": "/system/user", "meta": map[string]any{"title": "账号与权限"}})
}
if HasPermission(account, PermissionOutboxRead) || HasPermission(account, PermissionOutboxManage) {
systemChildren = append(systemChildren, map[string]any{"name": "Outbox", "path": "outbox", "component": "/gateway/outbox", "meta": map[string]any{"title": "事件投递"}})
}
if HasPermission(account, PermissionNotificationRead) || HasPermission(account, PermissionNotificationManage) {
systemChildren = append(systemChildren, map[string]any{"name": "Notifications", "path": "notifications", "component": "/gateway/notifications", "meta": map[string]any{"title": "通知中心"}})
}
if len(systemChildren) > 0 {
menus = append(menus, map[string]any{"name": "System", "path": "/system", "component": "/index/index", "meta": map[string]any{"title": "系统管理", "icon": "ri:user-3-line"}, "children": systemChildren})
}
return menus
}
func portalMenus() []map[string]any {
return []map[string]any{
{"name": "Portal", "path": "/portal", "component": "/index/index", "meta": map[string]any{"title": "AI 工作台", "icon": "ri:sparkling-line"}, "children": []map[string]any{
{"name": "PortalCatalog", "path": "catalog", "component": "/portal/catalog", "meta": map[string]any{"title": "资产目录", "fixedTab": true}},
{"name": "PortalMarketplace", "path": "marketplace", "component": "/portal/marketplace", "meta": map[string]any{"title": "资源市场"}},
{"name": "PortalPrompts", "path": "prompts", "component": "/portal/prompts", "meta": map[string]any{"title": "Prompt 广场"}},
{"name": "PortalUsage", "path": "usage", "component": "/portal/usage", "meta": map[string]any{"title": "我的用量"}},
{"name": "PortalAccess", "path": "access", "component": "/portal/access", "meta": map[string]any{"title": "模型权限"}},
}},
}
}
+415
View File
@@ -0,0 +1,415 @@
package identity
import (
"context"
"encoding/json"
"errors"
"fmt"
"net/http"
"regexp"
"sort"
"strings"
"aigateway.local/core/internal/platform/apiresponse"
platformid "aigateway.local/core/internal/platform/id"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgconn"
)
var (
ErrIdentityConflict = errors.New("identity already exists")
permissionPattern = regexp.MustCompile(`^[a-z][a-z0-9_.:-]{2,127}$`)
)
type ManagementHTTPHandler struct {
service *Service
mux *http.ServeMux
}
type identityInput struct {
Login string `json:"login"`
DisplayName string `json:"display_name"`
Role string `json:"role"`
Password *string `json:"password"`
Permissions []string `json:"permissions"`
Active *bool `json:"active"`
DepartmentID *string `json:"department_id"`
}
func NewManagementHTTPHandler(service *Service) *ManagementHTTPHandler {
h := &ManagementHTTPHandler{service: service, mux: http.NewServeMux()}
h.mux.HandleFunc("GET /api/v1/admin/identities/admins", h.list(KindAdmin))
h.mux.HandleFunc("POST /api/v1/admin/identities/admins", h.create(KindAdmin))
h.mux.HandleFunc("PUT /api/v1/admin/identities/admins/{identity_id}", h.update(KindAdmin))
h.mux.HandleFunc("GET /api/v1/admin/identities/portal-users", h.list(KindPortal))
h.mux.HandleFunc("POST /api/v1/admin/identities/portal-users", h.create(KindPortal))
h.mux.HandleFunc("PUT /api/v1/admin/identities/portal-users/{identity_id}", h.update(KindPortal))
h.mux.HandleFunc("GET /api/v1/admin/departments", h.listDepartments)
h.mux.HandleFunc("POST /api/v1/admin/departments", h.createDepartment)
h.mux.HandleFunc("PUT /api/v1/admin/departments/{department_id}", h.updateDepartment)
h.registerOIDC()
return h
}
func (h *ManagementHTTPHandler) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
h.mux.ServeHTTP(writer, request)
}
func (h *ManagementHTTPHandler) list(kind Kind) http.HandlerFunc {
return func(writer http.ResponseWriter, request *http.Request) {
if _, ok := h.requirePermission(writer, request); !ok {
return
}
accounts, err := h.service.repository.ListIdentities(request.Context(), kind)
if err != nil {
h.writeError(writer, err)
return
}
items := make([]map[string]any, 0, len(accounts))
for _, account := range accounts {
items = append(items, managementView(account))
}
apiresponse.OK(writer, items)
}
}
func (h *ManagementHTTPHandler) create(kind Kind) http.HandlerFunc {
return func(writer http.ResponseWriter, request *http.Request) {
actor, ok := h.requirePermission(writer, request)
if !ok {
return
}
input, account, password, ok := h.decode(writer, request, kind, true)
if !ok {
return
}
_ = input
passwordHash, err := h.service.hasher.Hash(password)
if err != nil {
h.writeError(writer, err)
return
}
account.PasswordHash = passwordHash
created, err := h.service.repository.CreateIdentity(request.Context(), account, actor.ID)
if err != nil {
h.writeError(writer, err)
return
}
apiresponse.OK(writer, managementView(created))
}
}
func (h *ManagementHTTPHandler) update(kind Kind) http.HandlerFunc {
return func(writer http.ResponseWriter, request *http.Request) {
actor, ok := h.requirePermission(writer, request)
if !ok {
return
}
_, account, password, ok := h.decode(writer, request, kind, false)
if !ok {
return
}
account.ID = request.PathValue("identity_id")
current, err := h.service.findByID(request.Context(), kind, account.ID)
if err != nil {
h.writeError(writer, err)
return
}
if kind == KindAdmin && actor.ID == current.ID && (account.Role != current.Role || !account.Active) {
apiresponse.Error(writer, http.StatusConflict, "不能停用自身账号或修改自身角色")
return
}
var passwordHash *string
if password != "" {
hash, err := h.service.hasher.Hash(password)
if err != nil {
h.writeError(writer, err)
return
}
passwordHash = &hash
}
updated, err := h.service.repository.UpdateIdentity(request.Context(), account, passwordHash, actor.ID)
if err != nil {
h.writeError(writer, err)
return
}
apiresponse.OK(writer, managementView(updated))
}
}
func (h *ManagementHTTPHandler) decode(writer http.ResponseWriter, request *http.Request, kind Kind, creating bool) (identityInput, Account, string, bool) {
var input identityInput
decoder := json.NewDecoder(http.MaxBytesReader(writer, request.Body, 1<<20))
decoder.DisallowUnknownFields()
if err := decoder.Decode(&input); err != nil {
apiresponse.Error(writer, http.StatusBadRequest, "请求格式无效")
return input, Account{}, "", false
}
input.Login = strings.ToLower(strings.TrimSpace(input.Login))
input.DisplayName = strings.TrimSpace(input.DisplayName)
input.Role = strings.ToLower(strings.TrimSpace(input.Role))
if input.Role == "" {
if kind == KindPortal {
input.Role = "member"
} else {
input.Role = "operator"
}
}
if len(input.Login) < 2 || len(input.Login) > 128 || len(input.DisplayName) > 64 {
apiresponse.Error(writer, http.StatusBadRequest, "账号或显示名称格式无效")
return input, Account{}, "", false
}
if kind == KindAdmin && input.Role != "superadmin" && input.Role != "operator" && input.Role != "auditor" {
apiresponse.Error(writer, http.StatusBadRequest, "管理员角色无效")
return input, Account{}, "", false
}
if kind == KindPortal && input.Role != "member" {
apiresponse.Error(writer, http.StatusBadRequest, "门户角色无效")
return input, Account{}, "", false
}
permissions, err := normalizePermissions(input.Permissions)
if err != nil {
apiresponse.Error(writer, http.StatusBadRequest, err.Error())
return input, Account{}, "", false
}
password := ""
if input.Password != nil {
password = *input.Password
}
if creating && len(password) < 12 || password != "" && len(password) < 12 || len(password) > 1024 {
apiresponse.Error(writer, http.StatusBadRequest, "口令长度必须为 12 至 1024 个字符")
return input, Account{}, "", false
}
active := true
if input.Active != nil {
active = *input.Active
}
var departmentID *string
if input.DepartmentID != nil && strings.TrimSpace(*input.DepartmentID) != "" {
value := strings.TrimSpace(*input.DepartmentID)
department, err := h.service.repository.GetDepartment(request.Context(), value)
if err != nil || !department.Active {
apiresponse.Error(writer, http.StatusBadRequest, "所选部门不存在或已停用")
return input, Account{}, "", false
}
departmentID = &value
}
return input, Account{
Kind: kind, Login: input.Login, DisplayName: input.DisplayName,
Role: input.Role, Permissions: permissions, Active: active, DepartmentID: departmentID,
}, password, true
}
func (h *ManagementHTTPHandler) requirePermission(writer http.ResponseWriter, request *http.Request) (Account, bool) {
account, err := h.service.Authenticate(request.Context(), KindAdmin, request.Header.Get("Authorization"))
if err != nil {
apiresponse.Error(writer, http.StatusUnauthorized, "登录状态无效或已过期")
return Account{}, false
}
if !HasPermission(account, PermissionIdentityManage) {
apiresponse.Error(writer, http.StatusForbidden, "缺少身份管理权限")
return Account{}, false
}
return account, true
}
func (h *ManagementHTTPHandler) writeError(writer http.ResponseWriter, err error) {
switch {
case errors.Is(err, ErrNotFound):
apiresponse.Error(writer, http.StatusNotFound, "账号不存在")
case errors.Is(err, ErrIdentityConflict):
apiresponse.Error(writer, http.StatusConflict, "账号已存在")
case errors.Is(err, ErrUnavailable):
apiresponse.Error(writer, http.StatusServiceUnavailable, "身份管理服务暂不可用")
default:
apiresponse.Error(writer, http.StatusInternalServerError, "身份管理操作失败")
}
}
func normalizePermissions(values []string) ([]string, error) {
seen := make(map[string]struct{})
result := make([]string, 0, len(values))
for _, value := range values {
value = strings.ToLower(strings.TrimSpace(value))
if value != "*" && !permissionPattern.MatchString(value) {
return nil, fmt.Errorf("权限字符串 %q 格式无效", value)
}
if _, exists := seen[value]; exists {
continue
}
seen[value] = struct{}{}
result = append(result, value)
}
sort.Strings(result)
return result, nil
}
func managementView(account Account) map[string]any {
return map[string]any{
"id": account.ID, "kind": account.Kind, "login": account.Login,
"display_name": account.DisplayName, "role": account.Role,
"permissions": account.Permissions, "effective_permissions": EffectivePermissions(account),
"active": account.Active, "auth_source": account.AuthSource,
"totp_enabled": account.TOTPEnabled, "locked_until": account.LockedUntil,
"created_at": account.CreatedAt, "updated_at": account.UpdatedAt,
"department_id": account.DepartmentID, "department_name": account.DepartmentName,
}
}
func (r *Repository) ListIdentities(ctx context.Context, kind Kind) ([]Account, error) {
if r.pool == nil {
return nil, ErrUnavailable
}
query := `
SELECT id::text, username, display_name, role, permissions, active,
totp_enabled, locked_until, 'local', created_at, updated_at
FROM gateway.admin_accounts ORDER BY lower(username)`
if kind == KindPortal {
query = `
SELECT u.id::text, u.account, u.name, u.role, u.permissions, u.active,
u.totp_enabled, u.locked_until, u.auth_source, u.created_at, u.updated_at,
u.department_id::text, COALESCE(d.name, '')
FROM gateway.portal_users u
LEFT JOIN gateway.departments d ON d.id = u.department_id
ORDER BY lower(u.account)`
}
rows, err := r.pool.Query(ctx, query)
if err != nil {
return nil, fmt.Errorf("%w: %v", ErrUnavailable, err)
}
defer rows.Close()
accounts := make([]Account, 0)
for rows.Next() {
account := Account{Kind: kind}
arguments := []any{
&account.ID, &account.Login, &account.DisplayName, &account.Role,
&account.Permissions, &account.Active, &account.TOTPEnabled,
&account.LockedUntil, &account.AuthSource, &account.CreatedAt, &account.UpdatedAt,
}
if kind == KindPortal {
arguments = append(arguments, &account.DepartmentID, &account.DepartmentName)
}
if err := rows.Scan(arguments...); err != nil {
return nil, fmt.Errorf("%w: %v", ErrUnavailable, err)
}
accounts = append(accounts, account)
}
return accounts, mapRepositoryError(rows.Err())
}
func (r *Repository) CreateIdentity(ctx context.Context, account Account, actorID string) (Account, error) {
if r.pool == nil {
return Account{}, ErrUnavailable
}
id, err := platformid.NewUUID()
if err != nil {
return Account{}, err
}
eventID, err := platformid.NewUUID()
if err != nil {
return Account{}, err
}
account.ID = id
tx, err := r.pool.Begin(ctx)
if err != nil {
return Account{}, fmt.Errorf("%w: %v", ErrUnavailable, err)
}
defer func() { _ = tx.Rollback(ctx) }()
if account.Kind == KindAdmin {
err = tx.QueryRow(ctx, `
INSERT INTO gateway.admin_accounts
(id, username, display_name, role, permissions, password_hash, active)
VALUES ($1, $2, $3, $4, $5, $6, $7)
RETURNING created_at, updated_at`, account.ID, account.Login, account.DisplayName,
account.Role, account.Permissions, account.PasswordHash, account.Active,
).Scan(&account.CreatedAt, &account.UpdatedAt)
} else {
account.AuthSource = "local"
err = tx.QueryRow(ctx, `
INSERT INTO gateway.portal_users
(id, account, name, role, permissions, password_hash, auth_source, active, department_id)
VALUES ($1, $2, $3, $4, $5, $6, 'local', $7, $8)
RETURNING created_at, updated_at`, account.ID, account.Login, account.DisplayName,
account.Role, account.Permissions, account.PasswordHash, account.Active, account.DepartmentID,
).Scan(&account.CreatedAt, &account.UpdatedAt)
}
if err != nil {
return Account{}, mapManagementError(err)
}
payload, _ := json.Marshal(map[string]any{"identity_id": account.ID, "kind": account.Kind, "actor_id": actorID})
if _, err := tx.Exec(ctx, `
INSERT INTO gateway.outbox_events
(event_id, event_type, event_version, aggregate_type, aggregate_id, payload)
VALUES ($1, 'identity.created', 1, 'identity', $2, $3)`, eventID, account.ID, payload); err != nil {
return Account{}, fmt.Errorf("%w: %v", ErrUnavailable, err)
}
if err := tx.Commit(ctx); err != nil {
return Account{}, fmt.Errorf("%w: %v", ErrUnavailable, err)
}
return account, nil
}
func (r *Repository) UpdateIdentity(ctx context.Context, account Account, passwordHash *string, actorID string) (Account, error) {
if r.pool == nil {
return Account{}, ErrUnavailable
}
eventID, err := platformid.NewUUID()
if err != nil {
return Account{}, err
}
tx, err := r.pool.Begin(ctx)
if err != nil {
return Account{}, fmt.Errorf("%w: %v", ErrUnavailable, err)
}
defer func() { _ = tx.Rollback(ctx) }()
if account.Kind == KindAdmin {
err = tx.QueryRow(ctx, `
UPDATE gateway.admin_accounts
SET username = $2, display_name = $3, role = $4, permissions = $5,
active = $6, password_hash = COALESCE($7, password_hash),
updated_at = clock_timestamp()
WHERE id = $1
RETURNING totp_enabled, locked_until, created_at, updated_at`, account.ID, account.Login,
account.DisplayName, account.Role, account.Permissions, account.Active, passwordHash,
).Scan(&account.TOTPEnabled, &account.LockedUntil, &account.CreatedAt, &account.UpdatedAt)
} else {
account.AuthSource = "local"
err = tx.QueryRow(ctx, `
UPDATE gateway.portal_users
SET account = $2, name = $3, role = $4, permissions = $5,
active = $6, password_hash = COALESCE($7, password_hash), department_id = $8,
updated_at = clock_timestamp()
WHERE id = $1
RETURNING auth_source, totp_enabled, locked_until, created_at, updated_at`, account.ID, account.Login,
account.DisplayName, account.Role, account.Permissions, account.Active, passwordHash, account.DepartmentID,
).Scan(&account.AuthSource, &account.TOTPEnabled, &account.LockedUntil, &account.CreatedAt, &account.UpdatedAt)
}
if err != nil {
return Account{}, mapManagementError(err)
}
payload, _ := json.Marshal(map[string]any{"identity_id": account.ID, "kind": account.Kind, "actor_id": actorID})
if _, err := tx.Exec(ctx, `
INSERT INTO gateway.outbox_events
(event_id, event_type, event_version, aggregate_type, aggregate_id, payload)
VALUES ($1, 'identity.updated', 1, 'identity', $2, $3)`, eventID, account.ID, payload); err != nil {
return Account{}, fmt.Errorf("%w: %v", ErrUnavailable, err)
}
if err := tx.Commit(ctx); err != nil {
return Account{}, fmt.Errorf("%w: %v", ErrUnavailable, err)
}
return account, nil
}
func mapManagementError(err error) error {
if errors.Is(err, pgx.ErrNoRows) {
return ErrNotFound
}
var pgError *pgconn.PgError
if errors.As(err, &pgError) && pgError.Code == "23505" {
return ErrIdentityConflict
}
if err != nil {
return fmt.Errorf("%w: %v", ErrUnavailable, err)
}
return nil
}
+818
View File
@@ -0,0 +1,818 @@
package identity
import (
"context"
"crypto"
"crypto/rand"
"crypto/rsa"
"crypto/sha256"
"encoding/base64"
"encoding/json"
"errors"
"fmt"
"io"
"math/big"
"net"
"net/http"
"net/url"
"regexp"
"strings"
"time"
"aigateway.local/core/internal/platform/apiresponse"
platformid "aigateway.local/core/internal/platform/id"
"github.com/jackc/pgx/v5"
)
const oidcMaxResponse = 2 << 20
// oidcMaxTokenLifetimeSeconds bounds how far an ID token's exp may sit past
// its iat, preventing long-lived or replayed tokens from being accepted.
const oidcMaxTokenLifetimeSeconds = 24 * 3600
var oidcCodePattern = regexp.MustCompile(`^[a-z][a-z0-9_-]{1,63}$`)
type OIDCProvider struct {
ID string `json:"id"`
Code string `json:"code"`
DisplayName string `json:"display_name"`
IssuerURL string `json:"issuer_url"`
ClientID string `json:"client_id"`
EncryptedCredentials []byte `json:"-"`
CredentialKEKVersion int `json:"-"`
RedirectURI string `json:"redirect_uri"`
PortalReturnURL string `json:"portal_return_url"`
Scopes []string `json:"scopes"`
AutoProvision bool `json:"auto_provision"`
DefaultDepartmentID *string `json:"default_department_id,omitempty"`
Enabled bool `json:"enabled"`
Revision int64 `json:"revision"`
CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"`
}
type oidcProviderInput struct {
Code string `json:"code"`
DisplayName string `json:"display_name"`
IssuerURL string `json:"issuer_url"`
ClientID string `json:"client_id"`
ClientSecret *string `json:"client_secret"`
RedirectURI string `json:"redirect_uri"`
PortalReturnURL string `json:"portal_return_url"`
Scopes []string `json:"scopes"`
AutoProvision bool `json:"auto_provision"`
DefaultDepartmentID *string `json:"default_department_id"`
Enabled bool `json:"enabled"`
}
type oidcCredentials struct {
ClientSecret string `json:"client_secret"`
}
type oidcChallenge struct{ ProviderID, Verifier, Nonce string }
type oidcExchange struct {
Token string `json:"token"`
}
type oidcDiscovery struct{ Issuer, AuthorizationEndpoint, TokenEndpoint, JWKSURI string }
func (h *ManagementHTTPHandler) registerOIDC() {
h.mux.HandleFunc("GET /api/v1/admin/identity-providers", h.listOIDCProviders)
h.mux.HandleFunc("POST /api/v1/admin/identity-providers", h.createOIDCProvider)
h.mux.HandleFunc("PUT /api/v1/admin/identity-providers/{provider_id}", h.updateOIDCProvider)
h.registerSAML()
}
func (h *HTTPHandler) registerOIDC() {
h.mux.HandleFunc("GET /api/v1/portal/sso/providers", h.listPublicOIDCProviders)
h.mux.HandleFunc("GET /api/v1/portal/sso/{provider_code}/start", h.startSSO)
h.mux.HandleFunc("GET /api/v1/portal/sso/{provider_code}/callback", h.callbackOIDC)
h.mux.HandleFunc("POST /api/v1/portal/sso/{provider_code}/callback", h.callbackSAML)
h.mux.HandleFunc("GET /api/v1/portal/sso/{provider_code}/metadata", h.samlMetadata)
h.mux.HandleFunc("POST /api/v1/portal/sso/exchange", h.exchangeOIDC)
}
func (h *ManagementHTTPHandler) listOIDCProviders(w http.ResponseWriter, r *http.Request) {
if _, ok := h.requirePermission(w, r); !ok {
return
}
records, err := h.service.repository.ListOIDCProviders(r.Context(), false)
if err != nil {
h.writeError(w, err)
return
}
items := make([]map[string]any, 0, len(records))
for _, record := range records {
items = append(items, h.oidcView(record))
}
apiresponse.OK(w, items)
}
func (h *ManagementHTTPHandler) createOIDCProvider(w http.ResponseWriter, r *http.Request) {
actor, ok := h.requirePermission(w, r)
if !ok {
return
}
input, record, ok := h.decodeOIDC(w, r, true)
if !ok {
return
}
if err := h.setOIDCCredentials(&record, strings.TrimSpace(*input.ClientSecret)); err != nil {
h.writeError(w, err)
return
}
created, err := h.service.repository.CreateOIDCProvider(r.Context(), record, actor.ID)
if err != nil {
h.writeError(w, err)
return
}
apiresponse.OK(w, h.oidcView(created))
}
func (h *ManagementHTTPHandler) updateOIDCProvider(w http.ResponseWriter, r *http.Request) {
actor, ok := h.requirePermission(w, r)
if !ok {
return
}
input, record, ok := h.decodeOIDC(w, r, false)
if !ok {
return
}
record.ID = r.PathValue("provider_id")
replace := input.ClientSecret != nil
if replace {
if err := h.setOIDCCredentials(&record, strings.TrimSpace(*input.ClientSecret)); err != nil {
h.writeError(w, err)
return
}
}
updated, err := h.service.repository.UpdateOIDCProvider(r.Context(), record, actor.ID, replace)
if err != nil {
h.writeError(w, err)
return
}
apiresponse.OK(w, h.oidcView(updated))
}
func (h *ManagementHTTPHandler) decodeOIDC(w http.ResponseWriter, r *http.Request, creating bool) (oidcProviderInput, OIDCProvider, bool) {
var input oidcProviderInput
decoder := json.NewDecoder(http.MaxBytesReader(w, r.Body, 1<<20))
decoder.DisallowUnknownFields()
if decoder.Decode(&input) != nil {
apiresponse.Error(w, 400, "请求格式无效")
return input, OIDCProvider{}, false
}
input.Code = strings.ToLower(strings.TrimSpace(input.Code))
input.DisplayName = strings.TrimSpace(input.DisplayName)
input.ClientID = strings.TrimSpace(input.ClientID)
if !oidcCodePattern.MatchString(input.Code) || input.DisplayName == "" || input.ClientID == "" ||
(creating && input.ClientSecret == nil) || (input.ClientSecret != nil && strings.TrimSpace(*input.ClientSecret) == "") {
apiresponse.Error(w, 400, "身份源代码、名称、Client ID 或 Client Secret 无效")
return input, OIDCProvider{}, false
}
issuer, err := validateOIDCURL(r.Context(), input.IssuerURL, h.service.allowPrivateIdentityProvider)
if err != nil {
apiresponse.Error(w, 400, "Issuer URL 无效")
return input, OIDCProvider{}, false
}
redirectURI, err := validateAbsoluteURL(input.RedirectURI)
redirectURL, _ := url.Parse(redirectURI)
if err != nil || redirectURL.Fragment != "" {
apiresponse.Error(w, 400, "回调 URL 无效")
return input, OIDCProvider{}, false
}
returnURL, err := validateAbsoluteURL(input.PortalReturnURL)
if err != nil {
apiresponse.Error(w, 400, "门户返回 URL 无效")
return input, OIDCProvider{}, false
}
if len(input.Scopes) == 0 {
input.Scopes = []string{"openid", "profile", "email"}
}
input.Scopes, err = normalizeOIDCScopes(input.Scopes)
if err != nil {
apiresponse.Error(w, 400, "OIDC scopes 无效")
return input, OIDCProvider{}, false
}
if !contains(input.Scopes, "openid") {
apiresponse.Error(w, 400, "OIDC scopes 必须包含 openid")
return input, OIDCProvider{}, false
}
if input.DefaultDepartmentID != nil && *input.DefaultDepartmentID != "" {
department, err := h.service.repository.GetDepartment(r.Context(), *input.DefaultDepartmentID)
if err != nil || !department.Active {
apiresponse.Error(w, 400, "默认部门不存在或已停用")
return input, OIDCProvider{}, false
}
}
return input, OIDCProvider{Code: input.Code, DisplayName: input.DisplayName, IssuerURL: issuer, ClientID: input.ClientID,
RedirectURI: redirectURI, PortalReturnURL: returnURL, Scopes: input.Scopes, AutoProvision: input.AutoProvision,
DefaultDepartmentID: input.DefaultDepartmentID, Enabled: input.Enabled}, true
}
func (h *ManagementHTTPHandler) setOIDCCredentials(record *OIDCProvider, secret string) error {
payload, _ := json.Marshal(oidcCredentials{ClientSecret: secret})
encrypted, version, err := h.service.idpCipher.Encrypt(payload)
if err != nil {
return err
}
record.EncryptedCredentials, record.CredentialKEKVersion = encrypted, version
return nil
}
func (h *ManagementHTTPHandler) oidcView(record OIDCProvider) map[string]any {
configured := false
if plaintext, err := h.service.idpCipher.Decrypt(record.EncryptedCredentials, record.CredentialKEKVersion); err == nil {
var credentials oidcCredentials
configured = json.Unmarshal(plaintext, &credentials) == nil && credentials.ClientSecret != ""
}
return map[string]any{"id": record.ID, "code": record.Code, "display_name": record.DisplayName, "issuer_url": record.IssuerURL,
"client_id": record.ClientID, "secret_configured": configured, "redirect_uri": record.RedirectURI, "portal_return_url": record.PortalReturnURL,
"scopes": record.Scopes, "auto_provision": record.AutoProvision, "default_department_id": record.DefaultDepartmentID,
"enabled": record.Enabled, "revision": record.Revision, "credential_kek_version": record.CredentialKEKVersion}
}
func (h *HTTPHandler) listPublicOIDCProviders(w http.ResponseWriter, r *http.Request) {
records, err := h.service.repository.ListPublicIdentityProviders(r.Context())
if err != nil {
h.writeIdentityError(w, err)
return
}
apiresponse.OK(w, records)
}
func (h *HTTPHandler) startOIDC(w http.ResponseWriter, r *http.Request) {
p, err := h.service.repository.GetOIDCProviderByCode(r.Context(), r.PathValue("provider_code"))
if err != nil || !p.Enabled {
http.NotFound(w, r)
return
}
discovery, err := h.service.discoverOIDC(r.Context(), p)
if err != nil {
apiresponse.Error(w, 502, "身份源发现失败")
return
}
verifier, err := randomURLToken(48)
if err != nil {
h.writeIdentityError(w, ErrUnavailable)
return
}
nonce, err := randomURLToken(32)
if err != nil {
h.writeIdentityError(w, ErrUnavailable)
return
}
state, err := h.service.sessions.StoreOneTime(r.Context(), "oidc-state", oidcChallenge{ProviderID: p.ID, Verifier: verifier, Nonce: nonce}, 5*time.Minute)
if err != nil {
h.writeIdentityError(w, err)
return
}
challenge := sha256.Sum256([]byte(verifier))
target, _ := url.Parse(discovery.AuthorizationEndpoint)
query := target.Query()
query.Set("response_type", "code")
query.Set("client_id", p.ClientID)
query.Set("redirect_uri", p.RedirectURI)
query.Set("scope", strings.Join(p.Scopes, " "))
query.Set("state", state)
query.Set("nonce", nonce)
query.Set("code_challenge", base64.RawURLEncoding.EncodeToString(challenge[:]))
query.Set("code_challenge_method", "S256")
target.RawQuery = query.Encode()
http.Redirect(w, r, target.String(), http.StatusFound)
}
func (h *HTTPHandler) callbackOIDC(w http.ResponseWriter, r *http.Request) {
if r.URL.Query().Get("error") != "" {
apiresponse.Error(w, 401, "OIDC 登录被拒绝")
return
}
if r.URL.Query().Get("state") == "" || r.URL.Query().Get("code") == "" {
apiresponse.Error(w, 401, "OIDC 回调参数无效")
return
}
var challenge oidcChallenge
if err := h.service.sessions.ConsumeOneTime(r.Context(), "oidc-state", r.URL.Query().Get("state"), &challenge); err != nil {
apiresponse.Error(w, 401, "OIDC state 无效或已使用")
return
}
p, err := h.service.repository.GetOIDCProviderByCode(r.Context(), r.PathValue("provider_code"))
if err != nil || p.ID != challenge.ProviderID || !p.Enabled {
apiresponse.Error(w, 401, "OIDC 身份源无效")
return
}
discovery, err := h.service.discoverOIDC(r.Context(), p)
if err != nil {
apiresponse.Error(w, 502, "身份源发现失败")
return
}
credentials, err := h.service.oidcCredentials(p)
if err != nil {
apiresponse.Error(w, 503, "身份源凭据不可用")
return
}
form := url.Values{"grant_type": {"authorization_code"}, "code": {r.URL.Query().Get("code")}, "redirect_uri": {p.RedirectURI}, "code_verifier": {challenge.Verifier}}
req, err := http.NewRequestWithContext(r.Context(), http.MethodPost, discovery.TokenEndpoint, strings.NewReader(form.Encode()))
if err != nil {
apiresponse.Error(w, 502, "OIDC token 交换失败")
return
}
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
req.SetBasicAuth(p.ClientID, credentials.ClientSecret)
response, err := h.service.oidcHTTPClient().Do(req)
if err != nil {
apiresponse.Error(w, 502, "OIDC token 交换失败")
return
}
defer response.Body.Close()
payload, _ := io.ReadAll(io.LimitReader(response.Body, oidcMaxResponse+1))
if response.StatusCode/100 != 2 || len(payload) > oidcMaxResponse {
apiresponse.Error(w, 502, "OIDC token 交换失败")
return
}
var tokens struct {
IDToken string `json:"id_token"`
}
if json.Unmarshal(payload, &tokens) != nil || tokens.IDToken == "" {
apiresponse.Error(w, 502, "OIDC 响应缺少 ID Token")
return
}
claims, err := h.service.verifyIDToken(r.Context(), discovery, p.ClientID, challenge.Nonce, tokens.IDToken)
if err != nil {
apiresponse.Error(w, 401, "OIDC ID Token 校验失败")
return
}
account, err := h.service.repository.ResolveExternalAccount(r.Context(), p, claims)
if err != nil {
h.writeIdentityError(w, err)
return
}
if !account.Active {
h.writeIdentityError(w, ErrAccountDisabled)
return
}
token, err := h.service.sessions.Create(r.Context(), principalFor(account))
if err != nil {
h.writeIdentityError(w, err)
return
}
exchange, err := h.service.sessions.StoreOneTime(r.Context(), "oidc-exchange", oidcExchange{Token: token}, time.Minute)
if err != nil {
h.writeIdentityError(w, err)
return
}
returnURL, _ := url.Parse(p.PortalReturnURL)
query := returnURL.Query()
query.Set("sso_code", exchange)
returnURL.RawQuery = query.Encode()
http.Redirect(w, r, returnURL.String(), http.StatusFound)
}
func (h *HTTPHandler) exchangeOIDC(w http.ResponseWriter, r *http.Request) {
var input struct {
Code string `json:"code"`
}
if !decodeJSON(w, r, &input) {
apiresponse.Error(w, 400, "请求格式无效")
return
}
var exchange oidcExchange
if h.service.sessions.ConsumeOneTime(r.Context(), "oidc-exchange", input.Code, &exchange) != nil {
apiresponse.Error(w, 401, "SSO 交换码无效或已使用")
return
}
apiresponse.OK(w, map[string]string{"token": exchange.Token, "refreshToken": ""})
}
type oidcClaims struct {
Issuer string `json:"iss"`
Subject string `json:"sub"`
Audience json.RawMessage `json:"aud"`
AuthorizedParty string `json:"azp"`
ExpiresAt int64 `json:"exp"`
IssuedAt int64 `json:"iat"`
NotBefore int64 `json:"nbf"`
Nonce string `json:"nonce"`
Email string `json:"email"`
PreferredUsername string `json:"preferred_username"`
Name string `json:"name"`
}
type externalProvider struct {
ID string
Code string
AuthSource string
AutoProvision bool
DefaultDepartmentID *string
}
type externalClaims struct {
Subject string
Email string
PreferredUsername string
Name string
}
func (s *Service) discoverOIDC(ctx context.Context, p OIDCProvider) (oidcDiscovery, error) {
target := strings.TrimRight(p.IssuerURL, "/") + "/.well-known/openid-configuration"
req, err := http.NewRequestWithContext(ctx, http.MethodGet, target, nil)
if err != nil {
return oidcDiscovery{}, err
}
response, err := s.oidcHTTPClient().Do(req)
if err != nil {
return oidcDiscovery{}, err
}
defer response.Body.Close()
payload, _ := io.ReadAll(io.LimitReader(response.Body, oidcMaxResponse+1))
var raw struct {
Issuer string `json:"issuer"`
AuthorizationEndpoint string `json:"authorization_endpoint"`
TokenEndpoint string `json:"token_endpoint"`
JWKSURI string `json:"jwks_uri"`
}
if response.StatusCode/100 != 2 || len(payload) > oidcMaxResponse || json.Unmarshal(payload, &raw) != nil || raw.Issuer != p.IssuerURL {
return oidcDiscovery{}, errors.New("invalid discovery")
}
for _, value := range []string{raw.AuthorizationEndpoint, raw.TokenEndpoint, raw.JWKSURI} {
if _, err := validateOIDCURL(ctx, value, s.allowPrivateIdentityProvider); err != nil {
return oidcDiscovery{}, err
}
}
return oidcDiscovery{Issuer: raw.Issuer, AuthorizationEndpoint: raw.AuthorizationEndpoint, TokenEndpoint: raw.TokenEndpoint, JWKSURI: raw.JWKSURI}, nil
}
func (s *Service) oidcCredentials(p OIDCProvider) (oidcCredentials, error) {
plaintext, err := s.idpCipher.Decrypt(p.EncryptedCredentials, p.CredentialKEKVersion)
var c oidcCredentials
if err == nil {
err = json.Unmarshal(plaintext, &c)
}
return c, err
}
func (s *Service) oidcHTTPClient() *http.Client {
if s.oidcClient == nil {
s.oidcClient = newOIDCHTTPClient(s.allowPrivateIdentityProvider)
}
return s.oidcClient
}
func newOIDCHTTPClient(allowPrivate bool) *http.Client {
return &http.Client{
Timeout: 10 * time.Second,
Transport: &http.Transport{
DialContext: safeOIDCDial(allowPrivate),
TLSHandshakeTimeout: 5 * time.Second,
ResponseHeaderTimeout: 8 * time.Second,
MaxIdleConns: 32,
MaxIdleConnsPerHost: 8,
IdleConnTimeout: 90 * time.Second,
},
CheckRedirect: func(*http.Request, []*http.Request) error { return errors.New("redirect rejected") },
}
}
func (s *Service) verifyIDToken(ctx context.Context, d oidcDiscovery, clientID, nonce, token string) (oidcClaims, error) {
parts := strings.Split(token, ".")
if len(parts) != 3 {
return oidcClaims{}, errors.New("jwt format")
}
headerBytes, err := base64.RawURLEncoding.DecodeString(parts[0])
if err != nil {
return oidcClaims{}, err
}
var header struct{ Alg, Kid string }
if json.Unmarshal(headerBytes, &header) != nil || header.Alg != "RS256" || header.Kid == "" {
return oidcClaims{}, errors.New("jwt header")
}
req, err := http.NewRequestWithContext(ctx, http.MethodGet, d.JWKSURI, nil)
if err != nil {
return oidcClaims{}, err
}
response, err := s.oidcHTTPClient().Do(req)
if err != nil {
return oidcClaims{}, err
}
defer response.Body.Close()
payload, _ := io.ReadAll(io.LimitReader(response.Body, oidcMaxResponse+1))
var keys struct {
Keys []struct{ Kid, Kty, N, E string }
}
if response.StatusCode/100 != 2 || len(payload) > oidcMaxResponse || json.Unmarshal(payload, &keys) != nil {
return oidcClaims{}, errors.New("jwks")
}
var key *rsa.PublicKey
for _, j := range keys.Keys {
if j.Kid == header.Kid && j.Kty == "RSA" {
nBytes, nErr := base64.RawURLEncoding.DecodeString(j.N)
eBytes, eErr := base64.RawURLEncoding.DecodeString(j.E)
e := 0
for _, b := range eBytes {
e = e<<8 + int(b)
}
if nErr == nil && eErr == nil && len(nBytes) >= 256 && e >= 3 && e%2 == 1 {
key = &rsa.PublicKey{N: new(big.Int).SetBytes(nBytes), E: e}
}
}
}
if key == nil {
return oidcClaims{}, errors.New("key")
}
signature, err := base64.RawURLEncoding.DecodeString(parts[2])
if err != nil {
return oidcClaims{}, err
}
digest := sha256.Sum256([]byte(parts[0] + "." + parts[1]))
if rsa.VerifyPKCS1v15(key, crypto.SHA256, digest[:], signature) != nil {
return oidcClaims{}, errors.New("signature")
}
claimsBytes, err := base64.RawURLEncoding.DecodeString(parts[1])
if err != nil {
return oidcClaims{}, err
}
var claims oidcClaims
now := time.Now().Unix()
if s.now != nil {
now = s.now().Unix()
}
if json.Unmarshal(claimsBytes, &claims) != nil || claims.Issuer != d.Issuer || claims.Subject == "" || claims.Nonce != nonce ||
claims.ExpiresAt <= now-30 || claims.IssuedAt == 0 || claims.IssuedAt > now+30 || claims.NotBefore > now+30 ||
claims.ExpiresAt-claims.IssuedAt > oidcMaxTokenLifetimeSeconds {
return oidcClaims{}, errors.New("claims")
}
audiences, ok := parseAudience(claims.Audience)
if !ok || !contains(audiences, clientID) || (len(audiences) > 1 && claims.AuthorizedParty != clientID) ||
(claims.AuthorizedParty != "" && claims.AuthorizedParty != clientID) {
return oidcClaims{}, errors.New("audience")
}
return claims, nil
}
func (r *Repository) ListOIDCProviders(ctx context.Context, enabledOnly bool) ([]OIDCProvider, error) {
query := `SELECT id::text,code,display_name,issuer_url,client_id,encrypted_credentials,credential_kek_version,redirect_uri,portal_return_url,scopes,auto_provision,default_department_id::text,enabled,revision,created_at,updated_at FROM gateway.identity_providers WHERE kind='oidc'`
if enabledOnly {
query += " AND enabled"
}
query += " ORDER BY code"
rows, err := r.pool.Query(ctx, query)
if err != nil {
return nil, fmt.Errorf("%w: %v", ErrUnavailable, err)
}
defer rows.Close()
items := []OIDCProvider{}
for rows.Next() {
var p OIDCProvider
if err := rows.Scan(&p.ID, &p.Code, &p.DisplayName, &p.IssuerURL, &p.ClientID, &p.EncryptedCredentials, &p.CredentialKEKVersion, &p.RedirectURI, &p.PortalReturnURL, &p.Scopes, &p.AutoProvision, &p.DefaultDepartmentID, &p.Enabled, &p.Revision, &p.CreatedAt, &p.UpdatedAt); err != nil {
return nil, err
}
items = append(items, p)
}
return items, mapRepositoryError(rows.Err())
}
func (r *Repository) GetOIDCProviderByCode(ctx context.Context, code string) (OIDCProvider, error) {
var p OIDCProvider
err := r.pool.QueryRow(ctx, `SELECT id::text,code,display_name,issuer_url,client_id,encrypted_credentials,credential_kek_version,redirect_uri,portal_return_url,scopes,auto_provision,default_department_id::text,enabled,revision,created_at,updated_at FROM gateway.identity_providers WHERE code=$1 AND kind='oidc'`, strings.ToLower(code)).Scan(&p.ID, &p.Code, &p.DisplayName, &p.IssuerURL, &p.ClientID, &p.EncryptedCredentials, &p.CredentialKEKVersion, &p.RedirectURI, &p.PortalReturnURL, &p.Scopes, &p.AutoProvision, &p.DefaultDepartmentID, &p.Enabled, &p.Revision, &p.CreatedAt, &p.UpdatedAt)
return p, mapRepositoryError(err)
}
func (r *Repository) CreateOIDCProvider(ctx context.Context, p OIDCProvider, actor string) (OIDCProvider, error) {
id, _ := platformid.NewUUID()
p.ID = id
return r.storeOIDCProvider(ctx, p, actor, true, true)
}
func (r *Repository) UpdateOIDCProvider(ctx context.Context, p OIDCProvider, actor string, replace bool) (OIDCProvider, error) {
return r.storeOIDCProvider(ctx, p, actor, false, replace)
}
func (r *Repository) storeOIDCProvider(ctx context.Context, p OIDCProvider, actor string, creating, replace bool) (OIDCProvider, error) {
tx, err := r.pool.Begin(ctx)
if err != nil {
return p, ErrUnavailable
}
defer tx.Rollback(ctx)
if creating {
err = tx.QueryRow(ctx, `INSERT INTO gateway.identity_providers(id,code,kind,display_name,issuer_url,client_id,encrypted_credentials,credential_kek_version,redirect_uri,portal_return_url,scopes,auto_provision,default_department_id,enabled) VALUES($1,$2,'oidc',$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13) RETURNING revision,created_at,updated_at`, p.ID, p.Code, p.DisplayName, p.IssuerURL, p.ClientID, p.EncryptedCredentials, p.CredentialKEKVersion, p.RedirectURI, p.PortalReturnURL, p.Scopes, p.AutoProvision, p.DefaultDepartmentID, p.Enabled).Scan(&p.Revision, &p.CreatedAt, &p.UpdatedAt)
} else {
err = tx.QueryRow(ctx, `UPDATE gateway.identity_providers SET code=$2,display_name=$3,issuer_url=$4,client_id=$5,encrypted_credentials=CASE WHEN $14 THEN $6 ELSE encrypted_credentials END,credential_kek_version=CASE WHEN $14 THEN $7 ELSE credential_kek_version END,redirect_uri=$8,portal_return_url=$9,scopes=$10,auto_provision=$11,default_department_id=$12,enabled=$13,revision=revision+1,updated_at=clock_timestamp() WHERE id=$1 AND kind='oidc' RETURNING encrypted_credentials,credential_kek_version,revision,created_at,updated_at`, p.ID, p.Code, p.DisplayName, p.IssuerURL, p.ClientID, p.EncryptedCredentials, p.CredentialKEKVersion, p.RedirectURI, p.PortalReturnURL, p.Scopes, p.AutoProvision, p.DefaultDepartmentID, p.Enabled, replace).Scan(&p.EncryptedCredentials, &p.CredentialKEKVersion, &p.Revision, &p.CreatedAt, &p.UpdatedAt)
}
if err != nil {
return p, mapManagementError(err)
}
eventID, _ := platformid.NewUUID()
eventType := "identity_provider.updated"
if creating {
eventType = "identity_provider.created"
}
payload, _ := json.Marshal(map[string]any{"identity_provider_id": p.ID, "actor_id": actor})
if _, err = tx.Exec(ctx, `INSERT INTO gateway.outbox_events(event_id,event_type,event_version,aggregate_type,aggregate_id,payload) VALUES($1,$2,1,'identity_provider',$3,$4)`, eventID, eventType, p.ID, payload); err != nil {
return p, ErrUnavailable
}
if tx.Commit(ctx) != nil {
return p, ErrUnavailable
}
return p, nil
}
func (r *Repository) ResolveExternalAccount(ctx context.Context, p OIDCProvider, c oidcClaims) (Account, error) {
return r.resolveExternalAccount(ctx, externalProvider{
ID: p.ID, Code: p.Code, AuthSource: "oidc", AutoProvision: p.AutoProvision, DefaultDepartmentID: p.DefaultDepartmentID,
}, externalClaims{Subject: c.Subject, Email: c.Email, PreferredUsername: c.PreferredUsername, Name: c.Name})
}
func (r *Repository) resolveExternalAccount(ctx context.Context, p externalProvider, c externalClaims) (Account, error) {
var id string
err := r.pool.QueryRow(ctx, `SELECT id::text FROM gateway.portal_users WHERE identity_provider_id=$1 AND external_subject=$2`, p.ID, c.Subject).Scan(&id)
if err == nil {
return r.FindPortalByID(ctx, id)
}
if !errors.Is(err, pgx.ErrNoRows) {
return Account{}, ErrUnavailable
}
if !p.AutoProvision {
return Account{}, ErrNotFound
}
if p.DefaultDepartmentID != nil {
var active bool
if err := r.pool.QueryRow(ctx, `SELECT active FROM gateway.departments WHERE id=$1`, *p.DefaultDepartmentID).Scan(&active); err != nil || !active {
return Account{}, ErrUnavailable
}
}
login := strings.ToLower(strings.TrimSpace(c.Email))
if login == "" {
login = strings.ToLower(strings.TrimSpace(c.PreferredUsername))
}
if login == "" {
sum := sha256.Sum256([]byte(c.Subject))
login = p.Code + "_" + fmt.Sprintf("%x", sum[:6])
}
login = truncate(login, 128)
id, _ = platformid.NewUUID()
eventID, _ := platformid.NewUUID()
tx, err := r.pool.Begin(ctx)
if err != nil {
return Account{}, ErrUnavailable
}
defer tx.Rollback(ctx)
name := truncate(firstNonEmpty(c.Name, c.PreferredUsername, login), 64)
insert := func(candidate string) (bool, error) {
var insertedID string
err := tx.QueryRow(ctx, `INSERT INTO gateway.portal_users(id,account,name,role,permissions,password_hash,auth_source,external_subject,identity_provider_id,department_id,active) VALUES($1,$2,$3,'member','{}',NULL,$4,$5,$6,$7,true) ON CONFLICT DO NOTHING RETURNING id::text`, id, candidate, name, p.AuthSource, c.Subject, p.ID, p.DefaultDepartmentID).Scan(&insertedID)
if errors.Is(err, pgx.ErrNoRows) {
return false, nil
}
return err == nil, err
}
inserted, err := insert(login)
if err != nil {
return Account{}, ErrUnavailable
}
if !inserted {
var existingID string
lookupErr := tx.QueryRow(ctx, `SELECT id::text FROM gateway.portal_users WHERE identity_provider_id=$1 AND external_subject=$2`, p.ID, c.Subject).Scan(&existingID)
if lookupErr == nil {
_ = tx.Rollback(ctx)
return r.FindPortalByID(ctx, existingID)
}
if !errors.Is(lookupErr, pgx.ErrNoRows) {
return Account{}, ErrUnavailable
}
sum := sha256.Sum256([]byte(p.Code + "|" + c.Subject))
login = truncate(login, 110) + "-" + fmt.Sprintf("%x", sum[:8])
inserted, err = insert(login)
}
if err != nil || !inserted {
return Account{}, ErrUnavailable
}
payload, _ := json.Marshal(map[string]any{"identity_id": id, "identity_provider_id": p.ID})
if _, err = tx.Exec(ctx, `INSERT INTO gateway.outbox_events(event_id,event_type,event_version,aggregate_type,aggregate_id,payload) VALUES($1,'identity.external_provisioned',1,'identity',$2,$3)`, eventID, id, payload); err != nil {
return Account{}, ErrUnavailable
}
if tx.Commit(ctx) != nil {
return Account{}, ErrUnavailable
}
return r.FindPortalByID(ctx, id)
}
func validateAbsoluteURL(raw string) (string, error) {
u, err := url.Parse(strings.TrimSpace(raw))
if err != nil || (u.Scheme != "http" && u.Scheme != "https") || u.Hostname() == "" || u.User != nil {
return "", errors.New("invalid url")
}
return u.String(), nil
}
func validateOIDCURL(ctx context.Context, raw string, allowPrivate bool) (string, error) {
value, err := validateAbsoluteURL(raw)
if err != nil {
return "", err
}
u, _ := url.Parse(value)
if u.RawQuery != "" || u.Fragment != "" || (!allowPrivate && u.Scheme != "https") {
return "", errors.New("invalid issuer url")
}
if !allowPrivate {
addresses, err := net.DefaultResolver.LookupIPAddr(ctx, u.Hostname())
if err != nil || len(addresses) == 0 {
return "", errors.New("resolve")
}
for _, a := range addresses {
if !publicIP(a.IP) {
return "", errors.New("blocked address")
}
}
}
u.Path = strings.TrimRight(u.Path, "/")
return u.String(), nil
}
func safeOIDCDial(allowPrivate bool) func(context.Context, string, string) (net.Conn, error) {
d := &net.Dialer{Timeout: 5 * time.Second, KeepAlive: 30 * time.Second}
if allowPrivate {
return d.DialContext
}
return func(ctx context.Context, network, address string) (net.Conn, error) {
host, port, err := net.SplitHostPort(address)
if err != nil {
return nil, err
}
addresses, err := net.DefaultResolver.LookupIPAddr(ctx, host)
if err != nil {
return nil, err
}
for _, a := range addresses {
if !publicIP(a.IP) {
return nil, errors.New("blocked address")
}
}
if len(addresses) == 0 {
return nil, errors.New("resolve")
}
return d.DialContext(ctx, network, net.JoinHostPort(addresses[0].IP.String(), port))
}
}
func publicIP(ip net.IP) bool {
return ip != nil && !ip.IsPrivate() && !ip.IsLoopback() && !ip.IsLinkLocalUnicast() && !ip.IsLinkLocalMulticast() && !ip.IsMulticast() && !ip.IsUnspecified()
}
func randomURLToken(size int) (string, error) {
b := make([]byte, size)
_, err := rand.Read(b)
return base64.RawURLEncoding.EncodeToString(b), err
}
func contains(values []string, target string) bool {
for _, v := range values {
if v == target {
return true
}
}
return false
}
func audienceContains(raw json.RawMessage, target string) bool {
values, ok := parseAudience(raw)
return ok && contains(values, target)
}
func parseAudience(raw json.RawMessage) ([]string, bool) {
var one string
if json.Unmarshal(raw, &one) == nil {
return []string{one}, one != ""
}
var many []string
if json.Unmarshal(raw, &many) != nil || len(many) == 0 {
return nil, false
}
for _, value := range many {
if value == "" {
return nil, false
}
}
return many, true
}
func normalizeOIDCScopes(values []string) ([]string, error) {
if len(values) > 16 {
return nil, errors.New("too many scopes")
}
seen := make(map[string]struct{}, len(values))
result := make([]string, 0, len(values))
for _, value := range values {
value = strings.TrimSpace(value)
if value == "" || len(value) > 64 || strings.ContainsAny(value, " \t\r\n\"") {
return nil, errors.New("invalid scope")
}
if _, exists := seen[value]; exists {
continue
}
seen[value] = struct{}{}
result = append(result, value)
}
return result, nil
}
func firstNonEmpty(values ...string) string {
for _, v := range values {
if strings.TrimSpace(v) != "" {
return strings.TrimSpace(v)
}
}
return "用户"
}
func truncate(value string, max int) string {
runes := []rune(value)
if len(runes) > max {
return string(runes[:max])
}
return value
}
+140
View File
@@ -0,0 +1,140 @@
package identity
import (
"context"
"crypto"
"crypto/rand"
"crypto/rsa"
"crypto/sha256"
"encoding/base64"
"encoding/json"
"math/big"
"net/http"
"net/http/httptest"
"testing"
"time"
)
func TestVerifyIDTokenValidatesOIDCSecurityClaims(t *testing.T) {
key, err := rsa.GenerateKey(rand.Reader, 2048)
if err != nil {
t.Fatal(err)
}
kid := "test-key"
jwks := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
_ = json.NewEncoder(w).Encode(map[string]any{"keys": []map[string]string{{
"kid": kid,
"kty": "RSA",
"n": base64.RawURLEncoding.EncodeToString(key.N.Bytes()),
"e": base64.RawURLEncoding.EncodeToString(big.NewInt(int64(key.E)).Bytes()),
}}})
}))
defer jwks.Close()
now := time.Unix(1_800_000_000, 0)
service := &Service{
allowPrivateIdentityProvider: true,
oidcClient: newOIDCHTTPClient(true),
now: func() time.Time { return now },
}
discovery := oidcDiscovery{Issuer: "https://issuer.example", JWKSURI: jwks.URL}
baseClaims := map[string]any{
"iss": discovery.Issuer,
"sub": "subject-1",
"aud": "client-1",
"exp": now.Add(time.Minute).Unix(),
"iat": now.Unix(),
"nonce": "nonce-1",
}
if _, err := service.verifyIDToken(context.Background(), discovery, "client-1", "nonce-1", signTestIDToken(t, key, kid, baseClaims)); err != nil {
t.Fatalf("valid ID token rejected: %v", err)
}
tests := []struct {
name string
change func(map[string]any)
}{
{name: "wrong issuer", change: func(c map[string]any) { c["iss"] = "https://attacker.example" }},
{name: "wrong nonce", change: func(c map[string]any) { c["nonce"] = "other" }},
{name: "expired", change: func(c map[string]any) { c["exp"] = now.Add(-time.Minute).Unix() }},
{name: "issued in future", change: func(c map[string]any) { c["iat"] = now.Add(time.Minute).Unix() }},
{name: "missing subject", change: func(c map[string]any) { delete(c, "sub") }},
{name: "wrong audience", change: func(c map[string]any) { c["aud"] = "other-client" }},
{name: "multiple audiences without azp", change: func(c map[string]any) { c["aud"] = []string{"client-1", "other-client"} }},
{name: "wrong authorized party", change: func(c map[string]any) { c["aud"] = []string{"client-1", "other-client"}; c["azp"] = "other-client" }},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
claims := cloneClaims(baseClaims)
test.change(claims)
if _, err := service.verifyIDToken(context.Background(), discovery, "client-1", "nonce-1", signTestIDToken(t, key, kid, claims)); err == nil {
t.Fatal("expected ID token to be rejected")
}
})
}
multiple := cloneClaims(baseClaims)
multiple["aud"] = []string{"client-1", "other-client"}
multiple["azp"] = "client-1"
if _, err := service.verifyIDToken(context.Background(), discovery, "client-1", "nonce-1", signTestIDToken(t, key, kid, multiple)); err != nil {
t.Fatalf("valid multi-audience ID token rejected: %v", err)
}
}
func TestOIDCURLAndScopeValidation(t *testing.T) {
got, err := validateOIDCURL(context.Background(), "http://127.0.0.1:9090/issuer/", true)
if err != nil || got != "http://127.0.0.1:9090/issuer" {
t.Fatalf("private development issuer rejected: %q %v", got, err)
}
for _, raw := range []string{
"http://127.0.0.1:9090/issuer?tenant=1",
"http://127.0.0.1:9090/issuer#fragment",
"http://8.8.8.8/issuer",
} {
allowPrivate := raw != "http://8.8.8.8/issuer"
if _, err := validateOIDCURL(context.Background(), raw, allowPrivate); err == nil {
t.Fatalf("unsafe issuer accepted: %s", raw)
}
}
if _, err := validateAbsoluteURL("https://user:secret@example.com/callback"); err == nil {
t.Fatal("URL containing credentials was accepted")
}
scopes, err := normalizeOIDCScopes([]string{"openid", " profile ", "openid", "email"})
if err != nil || len(scopes) != 3 || scopes[0] != "openid" || scopes[1] != "profile" || scopes[2] != "email" {
t.Fatalf("unexpected normalized scopes: %#v %v", scopes, err)
}
for _, scopes := range [][]string{{"openid email"}, {"openid", ""}} {
if _, err := normalizeOIDCScopes(scopes); err == nil {
t.Fatalf("invalid scopes accepted: %#v", scopes)
}
}
}
func signTestIDToken(t *testing.T, key *rsa.PrivateKey, kid string, claims map[string]any) string {
t.Helper()
header, err := json.Marshal(map[string]string{"alg": "RS256", "kid": kid, "typ": "JWT"})
if err != nil {
t.Fatal(err)
}
payload, err := json.Marshal(claims)
if err != nil {
t.Fatal(err)
}
signingInput := base64.RawURLEncoding.EncodeToString(header) + "." + base64.RawURLEncoding.EncodeToString(payload)
digest := sha256.Sum256([]byte(signingInput))
signature, err := rsa.SignPKCS1v15(rand.Reader, key, crypto.SHA256, digest[:])
if err != nil {
t.Fatal(err)
}
return signingInput + "." + base64.RawURLEncoding.EncodeToString(signature)
}
func cloneClaims(source map[string]any) map[string]any {
result := make(map[string]any, len(source))
for key, value := range source {
result[key] = value
}
return result
}
+103
View File
@@ -0,0 +1,103 @@
package identity
import (
"crypto/hmac"
"crypto/rand"
"crypto/sha256"
"crypto/subtle"
"encoding/binary"
"encoding/hex"
"errors"
"fmt"
"strconv"
"strings"
)
const (
currentPBKDF2Iterations = 600_000
legacyPBKDF2Iterations = 120_000
maximumPBKDF2Iterations = 2_000_000
passwordDigestBytes = 32
)
var ErrInvalidPasswordHash = errors.New("invalid password hash")
type PasswordHasher struct{}
func (PasswordHasher) Hash(password string) (string, error) {
saltBytes := make([]byte, 16)
if _, err := rand.Read(saltBytes); err != nil {
return "", fmt.Errorf("generate password salt: %w", err)
}
return hashWithSalt(password, hex.EncodeToString(saltBytes), currentPBKDF2Iterations), nil
}
func (PasswordHasher) Verify(password, stored string) bool {
iterations, salt, expected, err := parsePasswordHash(stored)
if err != nil {
return false
}
actual := pbkdf2SHA256([]byte(password), []byte(salt), iterations, len(expected))
return subtle.ConstantTimeCompare(actual, expected) == 1
}
func (PasswordHasher) NeedsUpgrade(stored string) bool {
iterations, _, _, err := parsePasswordHash(stored)
return err == nil && iterations < currentPBKDF2Iterations
}
func hashWithSalt(password, salt string, iterations int) string {
digest := pbkdf2SHA256([]byte(password), []byte(salt), iterations, passwordDigestBytes)
return fmt.Sprintf("pbkdf2_sha256$%d$%s$%s", iterations, salt, hex.EncodeToString(digest))
}
func parsePasswordHash(stored string) (int, string, []byte, error) {
parts := strings.Split(stored, "$")
iterations := legacyPBKDF2Iterations
var salt, encodedDigest string
switch {
case len(parts) == 4 && parts[0] == "pbkdf2_sha256":
parsed, err := strconv.Atoi(parts[1])
if err != nil {
return 0, "", nil, ErrInvalidPasswordHash
}
iterations, salt, encodedDigest = parsed, parts[2], parts[3]
case len(parts) == 2:
salt, encodedDigest = parts[0], parts[1]
default:
return 0, "", nil, ErrInvalidPasswordHash
}
if iterations < 1 || iterations > maximumPBKDF2Iterations || salt == "" {
return 0, "", nil, ErrInvalidPasswordHash
}
digest, err := hex.DecodeString(encodedDigest)
if err != nil || len(digest) != passwordDigestBytes {
return 0, "", nil, ErrInvalidPasswordHash
}
return iterations, salt, digest, nil
}
func pbkdf2SHA256(password, salt []byte, iterations, keyLength int) []byte {
const hashLength = sha256.Size
blocks := (keyLength + hashLength - 1) / hashLength
result := make([]byte, 0, blocks*hashLength)
buffer := make([]byte, len(salt)+4)
copy(buffer, salt)
for block := 1; block <= blocks; block++ {
binary.BigEndian.PutUint32(buffer[len(salt):], uint32(block))
mac := hmac.New(sha256.New, password)
_, _ = mac.Write(buffer)
u := mac.Sum(nil)
t := append([]byte(nil), u...)
for round := 1; round < iterations; round++ {
mac.Reset()
_, _ = mac.Write(u)
u = mac.Sum(nil)
for index := range t {
t[index] ^= u[index]
}
}
result = append(result, t...)
}
return result[:keyLength]
}
+36
View File
@@ -0,0 +1,36 @@
package identity
import "testing"
func TestPasswordHasherMatchesPythonGatewayFormat(t *testing.T) {
stored := "pbkdf2_sha256$600000$00112233445566778899aabbccddeeff$afca0887b188255f525e15e30f5aa5a0b210a3e253bfaf9630411f0782bb6573"
hasher := PasswordHasher{}
if !hasher.Verify("correct horse battery staple", stored) {
t.Fatal("expected Python-compatible password to verify")
}
if hasher.Verify("wrong", stored) {
t.Fatal("wrong password must not verify")
}
}
func TestPasswordHasherAcceptsLegacyFormat(t *testing.T) {
stored := "abcd1234$a377be9840f0f48e6a0f3577f08a9e56e095561e1e313517e3d589739f8d6907"
hasher := PasswordHasher{}
if !hasher.Verify("legacy", stored) {
t.Fatal("expected legacy password to verify")
}
if !hasher.NeedsUpgrade(stored) {
t.Fatal("legacy password should require upgrade")
}
}
func TestHashRoundTrip(t *testing.T) {
hasher := PasswordHasher{}
stored, err := hasher.Hash("long-enough-password")
if err != nil {
t.Fatal(err)
}
if !hasher.Verify("long-enough-password", stored) {
t.Fatal("new hash did not verify")
}
}
+29
View File
@@ -0,0 +1,29 @@
package identity
import "testing"
func TestRoleAndDirectPermissionsAreMerged(t *testing.T) {
account := Account{Role: "auditor", Permissions: []string{PermissionProviderManage}}
if !HasPermission(account, PermissionProviderRead) || !HasPermission(account, PermissionProviderManage) {
t.Fatal("expected role and direct permissions to be merged")
}
if HasPermission(account, PermissionIdentityManage) {
t.Fatal("unexpected identity management permission")
}
}
func TestSuperadminWildcard(t *testing.T) {
if !HasPermission(Account{Role: "superadmin"}, "future_resource:future_action") {
t.Fatal("superadmin wildcard must cover future permissions")
}
}
func TestNormalizePermissions(t *testing.T) {
permissions, err := normalizePermissions([]string{"provider:read", " provider:read ", "future.feature:execute"})
if err != nil || len(permissions) != 2 {
t.Fatalf("unexpected normalized permissions: %#v, %v", permissions, err)
}
if _, err := normalizePermissions([]string{"INVALID PERMISSION"}); err == nil {
t.Fatal("expected invalid permission to be rejected")
}
}
+104
View File
@@ -0,0 +1,104 @@
package identity
import (
"context"
"crypto/rand"
"encoding/hex"
"fmt"
"net"
"net/http"
"strings"
"time"
"github.com/redis/go-redis/v9"
)
// LoginLimiter bounds login attempts per client IP with a Redis-backed
// sliding-window counter. It is defense-in-depth layered on top of the
// per-account lockout (repository.RecordFailure):
//
// - IP limiter: throttles credential-stuffing spread across many accounts
// from one source (returns HTTP 429).
// - Account lockout: stops repeated attempts on a single account.
//
// If Redis is unavailable or not configured the limiter fails OPEN (login
// proceeds, account lockout still applies) rather than locking every user out.
type LoginLimiter struct {
client *redis.Client
max int
window time.Duration
}
func NewLoginLimiter(client *redis.Client, max int, window time.Duration) *LoginLimiter {
return &LoginLimiter{client: client, max: max, window: window}
}
// Allow reports whether a login attempt from ip may proceed.
func (l *LoginLimiter) Allow(ctx context.Context, ip string) bool {
if l == nil || l.client == nil || l.max <= 0 || l.window <= 0 || ip == "" {
return true // not configured -> fail open
}
now := time.Now().UnixMilli()
// Member must be unique per attempt so ZADD appends instead of overwriting
// the score of an identical timestamp.
res, err := allowLoginScript.Run(ctx, l.client,
[]string{loginLimitKey(ip)},
now, l.window.Milliseconds(), l.max, uniqueMember(now),
int(l.window.Seconds())+60,
).Int64Slice()
if err != nil {
return true // Redis hiccup -> fail open
}
// Script returns {1, count} when limited, {0, count+1} when admitted.
return len(res) == 2 && res[0] == 0
}
func loginLimitKey(ip string) string {
return "gateway:login-limit:" + ip
}
func uniqueMember(now int64) string {
var b [8]byte
if _, err := rand.Read(b[:]); err != nil {
return fmt.Sprintf("%d:%d", now, time.Now().UnixNano())
}
return fmt.Sprintf("%d:%s", now, hex.EncodeToString(b[:]))
}
// allowLoginScript atomically trims the window, counts entries, and (when
// under the limit) records the attempt and refreshes the key TTL.
//
// KEYS[1] = key
// ARGV[1] = now (ms) ARGV[2] = window (ms) ARGV[3] = max
// ARGV[4] = unique member ARGV[5] = TTL (s)
var allowLoginScript = redis.NewScript(`
local key = KEYS[1]
local now = tonumber(ARGV[1])
local window = tonumber(ARGV[2])
local max = tonumber(ARGV[3])
redis.call('ZREMRANGEBYSCORE', key, '-inf', now - window)
local count = redis.call('ZCARD', key)
if count >= max then
return {1, count}
end
redis.call('ZADD', key, now, ARGV[4])
redis.call('EXPIRE', key, ARGV[5])
return {0, count + 1}
`)
// ClientIP extracts the caller's IP for login rate limiting. X-Forwarded-For
// is trusted here because nginx is the only ingress and overwrites the header
// on every proxy hop; the first value is the client address. Falls back to
// RemoteAddr for direct connections.
func ClientIP(r *http.Request) string {
if fwd := r.Header.Get("X-Forwarded-For"); fwd != "" {
if first := strings.TrimSpace(strings.Split(fwd, ",")[0]); first != "" {
return first
}
}
host, _, err := net.SplitHostPort(r.RemoteAddr)
if err != nil {
host = r.RemoteAddr
}
return host
}
+34
View File
@@ -0,0 +1,34 @@
package identity
import (
"net/http"
"net/http/httptest"
"testing"
)
func TestClientIP(t *testing.T) {
cases := []struct {
name string
remoteAddr string
xfwd string
want string
}{
{"xfwd first value", "10.0.0.1:52341", "203.0.113.9, 10.0.0.2", "203.0.113.9"},
{"xfwd single", "10.0.0.1:52341", "198.51.100.7", "198.51.100.7"},
{"xfwd with spaces", "10.0.0.1:52341", " 192.0.2.5 ", "192.0.2.5"},
{"no xfwd falls back to remote", "203.0.113.9:8080", "", "203.0.113.9"},
{"no xfwd and no port", "[2001:db8::1]:443", "", "2001:db8::1"},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
req := httptest.NewRequest(http.MethodPost, "/", nil)
req.RemoteAddr = tc.remoteAddr
if tc.xfwd != "" {
req.Header.Set("X-Forwarded-For", tc.xfwd)
}
if got := ClientIP(req); got != tc.want {
t.Fatalf("ClientIP() = %q, want %q", got, tc.want)
}
})
}
}
+335
View File
@@ -0,0 +1,335 @@
package identity
import (
"context"
"crypto/subtle"
"encoding/json"
"errors"
"fmt"
"strings"
"time"
platformid "aigateway.local/core/internal/platform/id"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgxpool"
)
type Repository struct {
pool *pgxpool.Pool
}
func NewRepository(pool *pgxpool.Pool) *Repository {
return &Repository{pool: pool}
}
func (r *Repository) FindAdminByLogin(ctx context.Context, login string) (Account, error) {
if r.pool == nil {
return Account{}, ErrUnavailable
}
var account Account
account.Kind = KindAdmin
err := r.pool.QueryRow(ctx, `
SELECT id::text, username, display_name, role, permissions, password_hash, active,
failed_logins, locked_until, totp_enabled, encrypted_totp_secret,
totp_kek_version, totp_last_step, totp_backup_codes
FROM gateway.admin_accounts
WHERE lower(username) = lower($1)`, strings.TrimSpace(login)).Scan(
&account.ID, &account.Login, &account.DisplayName, &account.Role, &account.Permissions,
&account.PasswordHash, &account.Active, &account.FailedLogins,
&account.LockedUntil, &account.TOTPEnabled, &account.EncryptedTOTPSecret,
&account.TOTPKekVersion, &account.TOTPLastStep, &account.TOTPBackupCodes,
)
return account, mapRepositoryError(err)
}
func (r *Repository) FindAdminByID(ctx context.Context, id string) (Account, error) {
if r.pool == nil {
return Account{}, ErrUnavailable
}
var account Account
account.Kind = KindAdmin
err := r.pool.QueryRow(ctx, `
SELECT id::text, username, display_name, role, permissions, password_hash, active,
failed_logins, locked_until, totp_enabled, encrypted_totp_secret,
totp_kek_version, totp_last_step, totp_backup_codes
FROM gateway.admin_accounts
WHERE id = $1`, id).Scan(
&account.ID, &account.Login, &account.DisplayName, &account.Role, &account.Permissions,
&account.PasswordHash, &account.Active, &account.FailedLogins,
&account.LockedUntil, &account.TOTPEnabled, &account.EncryptedTOTPSecret,
&account.TOTPKekVersion, &account.TOTPLastStep, &account.TOTPBackupCodes,
)
return account, mapRepositoryError(err)
}
func (r *Repository) FindPortalByLogin(ctx context.Context, login string) (Account, error) {
if r.pool == nil {
return Account{}, ErrUnavailable
}
var account Account
account.Kind = KindPortal
err := r.pool.QueryRow(ctx, `
SELECT id::text, account, name, role, permissions, COALESCE(password_hash, ''), auth_source, active,
failed_logins, locked_until, totp_enabled, encrypted_totp_secret,
totp_kek_version, totp_last_step, totp_backup_codes, department_id::text
FROM gateway.portal_users
WHERE lower(account) = lower($1)`, strings.TrimSpace(login)).Scan(
&account.ID, &account.Login, &account.DisplayName, &account.Role, &account.Permissions, &account.PasswordHash,
&account.AuthSource, &account.Active, &account.FailedLogins,
&account.LockedUntil, &account.TOTPEnabled, &account.EncryptedTOTPSecret,
&account.TOTPKekVersion, &account.TOTPLastStep, &account.TOTPBackupCodes, &account.DepartmentID,
)
return account, mapRepositoryError(err)
}
func (r *Repository) FindPortalByID(ctx context.Context, id string) (Account, error) {
if r.pool == nil {
return Account{}, ErrUnavailable
}
var account Account
account.Kind = KindPortal
err := r.pool.QueryRow(ctx, `
SELECT id::text, account, name, role, permissions, COALESCE(password_hash, ''), auth_source, active,
failed_logins, locked_until, totp_enabled, encrypted_totp_secret,
totp_kek_version, totp_last_step, totp_backup_codes, department_id::text
FROM gateway.portal_users
WHERE id = $1`, id).Scan(
&account.ID, &account.Login, &account.DisplayName, &account.Role, &account.Permissions, &account.PasswordHash,
&account.AuthSource, &account.Active, &account.FailedLogins,
&account.LockedUntil, &account.TOTPEnabled, &account.EncryptedTOTPSecret,
&account.TOTPKekVersion, &account.TOTPLastStep, &account.TOTPBackupCodes, &account.DepartmentID,
)
return account, mapRepositoryError(err)
}
func (r *Repository) RecordFailure(ctx context.Context, account Account, maximum int, lockDuration time.Duration) (*time.Time, error) {
if r.pool == nil {
return nil, ErrUnavailable
}
table := "gateway.admin_accounts"
if account.Kind == KindPortal {
table = "gateway.portal_users"
}
query := fmt.Sprintf(`
UPDATE %s
SET locked_until = CASE
WHEN failed_logins + 1 >= $2
THEN clock_timestamp() + make_interval(secs => $3)
ELSE locked_until
END,
failed_logins = CASE WHEN failed_logins + 1 >= $2 THEN 0 ELSE failed_logins + 1 END,
updated_at = clock_timestamp()
WHERE id = $1
RETURNING locked_until`, table)
var lockedUntil *time.Time
err := r.pool.QueryRow(ctx, query, account.ID, maximum, int64(lockDuration.Seconds())).Scan(&lockedUntil)
return lockedUntil, mapRepositoryError(err)
}
func (r *Repository) CompleteLogin(ctx context.Context, account Account, upgradedHash *string) error {
if r.pool == nil {
return ErrUnavailable
}
table := "gateway.admin_accounts"
if account.Kind == KindPortal {
table = "gateway.portal_users"
}
query := fmt.Sprintf(`
UPDATE %s
SET failed_logins = 0,
locked_until = NULL,
last_login = clock_timestamp(),
password_hash = COALESCE($2::text, password_hash),
updated_at = clock_timestamp()
WHERE id = $1`, table)
result, err := r.pool.Exec(ctx, query, account.ID, upgradedHash)
if err != nil {
return fmt.Errorf("%w: %v", ErrUnavailable, err)
}
if result.RowsAffected() != 1 {
return ErrNotFound
}
return nil
}
func (r *Repository) SetPassword(ctx context.Context, account Account, passwordHash string) error {
query := fmt.Sprintf(`UPDATE %s SET password_hash=$2, failed_logins=0, locked_until=NULL, updated_at=clock_timestamp() WHERE id=$1`, identityTable(account.Kind))
return expectOne(r.pool, ctx, query, account.ID, passwordHash)
}
func (r *Repository) SetTOTPSecret(ctx context.Context, account Account, encrypted []byte, version int) error {
query := fmt.Sprintf(`
UPDATE %s
SET encrypted_totp_secret = $2, totp_kek_version = $3,
totp_enabled = false, totp_last_step = NULL,
totp_backup_codes = '[]'::jsonb, totp_confirmed_at = NULL,
updated_at = clock_timestamp()
WHERE id = $1`, identityTable(account.Kind))
return expectOne(r.pool, ctx, query, account.ID, encrypted, version)
}
func (r *Repository) EnableTOTP(ctx context.Context, account Account, step int64, records []BackupCodeRecord) error {
payload, err := json.Marshal(records)
if err != nil {
return err
}
query := fmt.Sprintf(`
UPDATE %s
SET totp_enabled = true, totp_last_step = $2,
totp_backup_codes = $3::jsonb, totp_confirmed_at = clock_timestamp(),
updated_at = clock_timestamp()
WHERE id = $1 AND encrypted_totp_secret IS NOT NULL AND NOT totp_enabled
AND (totp_last_step IS NULL OR totp_last_step < $2)`, identityTable(account.Kind))
return expectOne(r.pool, ctx, query, account.ID, step, payload)
}
func (r *Repository) ConsumeTOTPStep(ctx context.Context, account Account, step int64) (bool, error) {
if r.pool == nil {
return false, ErrUnavailable
}
query := fmt.Sprintf(`
UPDATE %s SET totp_last_step = $2, updated_at = clock_timestamp()
WHERE id = $1 AND totp_enabled
AND (totp_last_step IS NULL OR totp_last_step < $2)`, identityTable(account.Kind))
result, err := r.pool.Exec(ctx, query, account.ID, step)
if err != nil {
return false, fmt.Errorf("%w: %v", ErrUnavailable, err)
}
return result.RowsAffected() == 1, nil
}
func (r *Repository) ConsumeBackupCode(ctx context.Context, account Account, hash string) (bool, error) {
if r.pool == nil {
return false, ErrUnavailable
}
tx, err := r.pool.BeginTx(ctx, pgx.TxOptions{})
if err != nil {
return false, fmt.Errorf("%w: %v", ErrUnavailable, err)
}
defer func() { _ = tx.Rollback(ctx) }()
query := fmt.Sprintf(`SELECT totp_backup_codes FROM %s WHERE id = $1 AND totp_enabled FOR UPDATE`, identityTable(account.Kind))
var payload []byte
if err := tx.QueryRow(ctx, query, account.ID).Scan(&payload); err != nil {
return false, mapRepositoryError(err)
}
var records []BackupCodeRecord
if err := json.Unmarshal(payload, &records); err != nil {
return false, fmt.Errorf("invalid stored TOTP backup codes: %w", err)
}
found := -1
for index := range records {
if records[index].UsedAt == nil && subtle.ConstantTimeCompare([]byte(records[index].Hash), []byte(hash)) == 1 {
found = index
}
}
if found < 0 {
return false, nil
}
now := time.Now().UTC()
records[found].UsedAt = &now
payload, err = json.Marshal(records)
if err != nil {
return false, err
}
update := fmt.Sprintf(`UPDATE %s SET totp_backup_codes = $2::jsonb, updated_at = clock_timestamp() WHERE id = $1`, identityTable(account.Kind))
if _, err := tx.Exec(ctx, update, account.ID, payload); err != nil {
return false, fmt.Errorf("%w: %v", ErrUnavailable, err)
}
if err := tx.Commit(ctx); err != nil {
return false, fmt.Errorf("%w: %v", ErrUnavailable, err)
}
return true, nil
}
func (r *Repository) DisableTOTP(ctx context.Context, account Account) error {
query := fmt.Sprintf(`
UPDATE %s SET encrypted_totp_secret = NULL, totp_kek_version = NULL,
totp_enabled = false, totp_last_step = NULL,
totp_backup_codes = '[]'::jsonb, totp_confirmed_at = NULL,
updated_at = clock_timestamp()
WHERE id = $1 AND totp_enabled`, identityTable(account.Kind))
return expectOne(r.pool, ctx, query, account.ID)
}
func (r *Repository) ReplaceBackupCodes(ctx context.Context, account Account, records []BackupCodeRecord) error {
payload, err := json.Marshal(records)
if err != nil {
return err
}
query := fmt.Sprintf(`UPDATE %s SET totp_backup_codes = $2::jsonb, updated_at = clock_timestamp() WHERE id = $1 AND totp_enabled`, identityTable(account.Kind))
return expectOne(r.pool, ctx, query, account.ID, payload)
}
func identityTable(kind Kind) string {
if kind == KindPortal {
return "gateway.portal_users"
}
return "gateway.admin_accounts"
}
func expectOne(pool *pgxpool.Pool, ctx context.Context, query string, arguments ...any) error {
if pool == nil {
return ErrUnavailable
}
result, err := pool.Exec(ctx, query, arguments...)
if err != nil {
return fmt.Errorf("%w: %v", ErrUnavailable, err)
}
if result.RowsAffected() != 1 {
return ErrNotFound
}
return nil
}
func (r *Repository) CreateAdmin(ctx context.Context, login, displayName, role, passwordHash string) (string, error) {
if r.pool == nil {
return "", ErrUnavailable
}
id, err := platformid.NewUUID()
if err != nil {
return "", err
}
result, err := r.pool.Exec(ctx, `
INSERT INTO gateway.admin_accounts
(id, username, display_name, role, password_hash)
VALUES ($1, $2, $3, $4, $5)
ON CONFLICT DO NOTHING`, id, strings.ToLower(strings.TrimSpace(login)), displayName, role, passwordHash)
if err != nil {
return "", err
}
if result.RowsAffected() != 1 {
return "", errors.New("administrator already exists")
}
return id, nil
}
func (r *Repository) CreatePortalUser(ctx context.Context, login, displayName, passwordHash string) (string, error) {
if r.pool == nil {
return "", ErrUnavailable
}
id, err := platformid.NewUUID()
if err != nil {
return "", err
}
result, err := r.pool.Exec(ctx, `
INSERT INTO gateway.portal_users (id, account, name, password_hash, auth_source)
VALUES ($1, $2, $3, $4, 'local')
ON CONFLICT DO NOTHING`, id, strings.ToLower(strings.TrimSpace(login)), displayName, passwordHash)
if err != nil {
return "", err
}
if result.RowsAffected() != 1 {
return "", errors.New("portal user already exists")
}
return id, nil
}
func mapRepositoryError(err error) error {
if errors.Is(err, pgx.ErrNoRows) {
return ErrNotFound
}
if err != nil {
return fmt.Errorf("%w: %v", ErrUnavailable, err)
}
return nil
}
+636
View File
@@ -0,0 +1,636 @@
package identity
import (
"context"
"crypto/x509"
"encoding/base64"
"encoding/json"
"encoding/xml"
"errors"
"io"
"net/http"
"net/url"
"strings"
"time"
"aigateway.local/core/internal/platform/apiresponse"
platformid "aigateway.local/core/internal/platform/id"
"github.com/beevik/etree"
"github.com/crewjam/saml"
"github.com/crewjam/saml/samlsp"
dsig "github.com/russellhaering/goxmldsig"
)
const samlMaxResponse = 4 << 20
type SAMLConfig struct {
MetadataURL string `json:"metadata_url"`
SPEntityID string `json:"sp_entity_id"`
ACSURL string `json:"acs_url"`
EmailAttribute string `json:"email_attribute"`
NameAttribute string `json:"name_attribute"`
}
type SAMLProvider struct {
ID string
Code string
DisplayName string
PortalReturnURL string
AutoProvision bool
DefaultDepartmentID *string
Enabled bool
Revision int64
Config SAMLConfig
CreatedAt time.Time
UpdatedAt time.Time
}
type samlProviderInput struct {
Code string `json:"code"`
DisplayName string `json:"display_name"`
MetadataURL string `json:"metadata_url"`
SPEntityID string `json:"sp_entity_id"`
ACSURL string `json:"acs_url"`
PortalReturnURL string `json:"portal_return_url"`
EmailAttribute string `json:"email_attribute"`
NameAttribute string `json:"name_attribute"`
AutoProvision bool `json:"auto_provision"`
DefaultDepartmentID *string `json:"default_department_id"`
Enabled bool `json:"enabled"`
}
type samlChallenge struct {
ProviderID string `json:"provider_id"`
RequestID string `json:"request_id"`
}
type samlMetadataCacheEntry struct {
Metadata *saml.EntityDescriptor
ExpiresAt time.Time
}
type PublicIdentityProvider struct {
Code string `json:"code"`
DisplayName string `json:"display_name"`
Kind string `json:"kind"`
}
func (h *ManagementHTTPHandler) registerSAML() {
h.mux.HandleFunc("GET /api/v1/admin/saml-providers", h.listSAMLProviders)
h.mux.HandleFunc("POST /api/v1/admin/saml-providers", h.createSAMLProvider)
h.mux.HandleFunc("PUT /api/v1/admin/saml-providers/{provider_id}", h.updateSAMLProvider)
}
func (h *ManagementHTTPHandler) listSAMLProviders(w http.ResponseWriter, r *http.Request) {
if _, ok := h.requirePermission(w, r); !ok {
return
}
providers, err := h.service.repository.ListSAMLProviders(r.Context())
if err != nil {
h.writeError(w, err)
return
}
items := make([]map[string]any, 0, len(providers))
for _, provider := range providers {
items = append(items, samlView(provider))
}
apiresponse.OK(w, items)
}
func (h *ManagementHTTPHandler) createSAMLProvider(w http.ResponseWriter, r *http.Request) {
actor, ok := h.requirePermission(w, r)
if !ok {
return
}
provider, ok := h.decodeSAMLProvider(w, r)
if !ok {
return
}
created, err := h.service.repository.CreateSAMLProvider(r.Context(), provider, actor.ID)
if err != nil {
h.writeError(w, err)
return
}
apiresponse.OK(w, samlView(created))
}
func (h *ManagementHTTPHandler) updateSAMLProvider(w http.ResponseWriter, r *http.Request) {
actor, ok := h.requirePermission(w, r)
if !ok {
return
}
provider, ok := h.decodeSAMLProvider(w, r)
if !ok {
return
}
provider.ID = r.PathValue("provider_id")
updated, err := h.service.repository.UpdateSAMLProvider(r.Context(), provider, actor.ID)
if err != nil {
h.writeError(w, err)
return
}
apiresponse.OK(w, samlView(updated))
}
func (h *ManagementHTTPHandler) decodeSAMLProvider(w http.ResponseWriter, r *http.Request) (SAMLProvider, bool) {
var input samlProviderInput
decoder := json.NewDecoder(http.MaxBytesReader(w, r.Body, 1<<20))
decoder.DisallowUnknownFields()
if decoder.Decode(&input) != nil {
apiresponse.Error(w, http.StatusBadRequest, "请求格式无效")
return SAMLProvider{}, false
}
input.Code = strings.ToLower(strings.TrimSpace(input.Code))
input.DisplayName = strings.TrimSpace(input.DisplayName)
input.EmailAttribute = strings.TrimSpace(input.EmailAttribute)
input.NameAttribute = strings.TrimSpace(input.NameAttribute)
if !oidcCodePattern.MatchString(input.Code) || input.DisplayName == "" || len(input.DisplayName) > 128 {
apiresponse.Error(w, http.StatusBadRequest, "身份源代码或名称无效")
return SAMLProvider{}, false
}
metadataURL, err := validateOIDCURL(r.Context(), input.MetadataURL, h.service.allowPrivateIdentityProvider)
if err != nil {
apiresponse.Error(w, http.StatusBadRequest, "SAML metadata URL 无效")
return SAMLProvider{}, false
}
entityID, err := validateSAMLEntityID(input.SPEntityID)
if err != nil {
apiresponse.Error(w, http.StatusBadRequest, "SAML SP Entity ID 无效")
return SAMLProvider{}, false
}
acsURL, err := validateAbsoluteURL(input.ACSURL)
acs, _ := url.Parse(acsURL)
expectedACSSuffix := "/api/v1/portal/sso/" + input.Code + "/callback"
if err != nil || acs.Fragment != "" || !strings.HasSuffix(strings.TrimRight(acs.Path, "/"), expectedACSSuffix) {
apiresponse.Error(w, http.StatusBadRequest, "SAML ACS URL 无效")
return SAMLProvider{}, false
}
returnURL, err := validateAbsoluteURL(input.PortalReturnURL)
if err != nil {
apiresponse.Error(w, http.StatusBadRequest, "门户返回 URL 无效")
return SAMLProvider{}, false
}
if input.EmailAttribute == "" {
input.EmailAttribute = "email"
}
if input.NameAttribute == "" {
input.NameAttribute = "displayName"
}
if len(input.EmailAttribute) > 256 || len(input.NameAttribute) > 256 {
apiresponse.Error(w, http.StatusBadRequest, "SAML 属性名过长")
return SAMLProvider{}, false
}
if input.DefaultDepartmentID != nil && *input.DefaultDepartmentID != "" {
department, err := h.service.repository.GetDepartment(r.Context(), *input.DefaultDepartmentID)
if err != nil || !department.Active {
apiresponse.Error(w, http.StatusBadRequest, "默认部门不存在或已停用")
return SAMLProvider{}, false
}
}
provider := SAMLProvider{
Code: input.Code, DisplayName: input.DisplayName, PortalReturnURL: returnURL,
AutoProvision: input.AutoProvision, DefaultDepartmentID: input.DefaultDepartmentID, Enabled: input.Enabled,
Config: SAMLConfig{MetadataURL: metadataURL, SPEntityID: entityID, ACSURL: acsURL, EmailAttribute: input.EmailAttribute, NameAttribute: input.NameAttribute},
}
if input.Enabled {
if _, err := h.service.samlServiceProvider(r.Context(), provider); err != nil {
apiresponse.Error(w, http.StatusBadRequest, "SAML metadata 无法验证或缺少有效签名证书")
return SAMLProvider{}, false
}
}
return provider, true
}
func samlView(provider SAMLProvider) map[string]any {
return map[string]any{
"id": provider.ID, "kind": "saml", "code": provider.Code, "display_name": provider.DisplayName,
"metadata_url": provider.Config.MetadataURL, "sp_entity_id": provider.Config.SPEntityID,
"acs_url": provider.Config.ACSURL, "portal_return_url": provider.PortalReturnURL,
"email_attribute": provider.Config.EmailAttribute, "name_attribute": provider.Config.NameAttribute,
"auto_provision": provider.AutoProvision, "default_department_id": provider.DefaultDepartmentID,
"enabled": provider.Enabled, "revision": provider.Revision,
}
}
func (h *HTTPHandler) startSSO(w http.ResponseWriter, r *http.Request) {
kind, err := h.service.repository.GetIdentityProviderKind(r.Context(), r.PathValue("provider_code"))
if err != nil {
http.NotFound(w, r)
return
}
if kind == "saml" {
h.startSAML(w, r)
return
}
h.startOIDC(w, r)
}
func (h *HTTPHandler) startSAML(w http.ResponseWriter, r *http.Request) {
provider, err := h.service.repository.GetSAMLProviderByCode(r.Context(), r.PathValue("provider_code"))
if err != nil || !provider.Enabled {
http.NotFound(w, r)
return
}
sp, err := h.service.samlServiceProvider(r.Context(), provider)
if err != nil {
apiresponse.Error(w, http.StatusBadGateway, "SAML metadata 获取失败")
return
}
idpURL := sp.GetSSOBindingLocation(saml.HTTPRedirectBinding)
if idpURL == "" {
apiresponse.Error(w, http.StatusBadGateway, "SAML 身份源不支持 Redirect 登录")
return
}
authnRequest, err := sp.MakeAuthenticationRequest(idpURL, saml.HTTPRedirectBinding, saml.HTTPPostBinding)
if err != nil {
apiresponse.Error(w, http.StatusBadGateway, "SAML 登录请求生成失败")
return
}
relayState, err := h.service.sessions.StoreOneTime(r.Context(), "saml-state", samlChallenge{ProviderID: provider.ID, RequestID: authnRequest.ID}, 5*time.Minute)
if err != nil {
h.writeIdentityError(w, err)
return
}
target, err := authnRequest.Redirect(relayState, sp)
if err != nil {
apiresponse.Error(w, http.StatusBadGateway, "SAML 登录请求生成失败")
return
}
http.Redirect(w, r, target.String(), http.StatusFound)
}
func (h *HTTPHandler) callbackSAML(w http.ResponseWriter, r *http.Request) {
r.Body = http.MaxBytesReader(w, r.Body, samlMaxResponse)
if err := r.ParseForm(); err != nil || r.PostForm.Get("SAMLResponse") == "" || r.PostForm.Get("RelayState") == "" || r.PostForm.Get("SAMLart") != "" {
apiresponse.Error(w, http.StatusBadRequest, "SAML 回调格式无效")
return
}
var challenge samlChallenge
if err := h.service.sessions.ConsumeOneTime(r.Context(), "saml-state", r.PostForm.Get("RelayState"), &challenge); err != nil {
apiresponse.Error(w, http.StatusUnauthorized, "SAML RelayState 无效或已使用")
return
}
provider, err := h.service.repository.GetSAMLProviderByCode(r.Context(), r.PathValue("provider_code"))
if err != nil || !provider.Enabled || provider.ID != challenge.ProviderID {
apiresponse.Error(w, http.StatusUnauthorized, "SAML 身份源无效")
return
}
sp, err := h.service.samlServiceProvider(r.Context(), provider)
if err != nil {
apiresponse.Error(w, http.StatusBadGateway, "SAML metadata 获取失败")
return
}
assertion, err := sp.ParseResponse(r, []string{challenge.RequestID})
if err != nil || assertion == nil || assertion.Subject == nil || assertion.Subject.NameID == nil || strings.TrimSpace(assertion.Subject.NameID.Value) == "" {
apiresponse.Error(w, http.StatusUnauthorized, "SAML 断言校验失败")
return
}
claimed, err := h.service.sessions.ClaimIdentifier(r.Context(), "saml-assertion", assertion.ID, 10*time.Minute)
if err != nil {
h.writeIdentityError(w, err)
return
}
if !claimed {
apiresponse.Error(w, http.StatusUnauthorized, "SAML 断言已使用")
return
}
subject := strings.TrimSpace(assertion.Subject.NameID.Value)
email := samlAttribute(assertion, provider.Config.EmailAttribute)
name := samlAttribute(assertion, provider.Config.NameAttribute)
account, err := h.service.repository.resolveExternalAccount(r.Context(), externalProvider{
ID: provider.ID, Code: provider.Code, AuthSource: "saml", AutoProvision: provider.AutoProvision, DefaultDepartmentID: provider.DefaultDepartmentID,
}, externalClaims{Subject: subject, Email: email, PreferredUsername: subject, Name: name})
if err != nil {
h.writeIdentityError(w, err)
return
}
if !account.Active {
h.writeIdentityError(w, ErrAccountDisabled)
return
}
token, err := h.service.sessions.Create(r.Context(), principalFor(account))
if err != nil {
h.writeIdentityError(w, err)
return
}
exchange, err := h.service.sessions.StoreOneTime(r.Context(), "oidc-exchange", oidcExchange{Token: token}, time.Minute)
if err != nil {
h.writeIdentityError(w, err)
return
}
returnURL, _ := url.Parse(provider.PortalReturnURL)
query := returnURL.Query()
query.Set("sso_code", exchange)
returnURL.RawQuery = query.Encode()
http.Redirect(w, r, returnURL.String(), http.StatusFound)
}
func (h *HTTPHandler) samlMetadata(w http.ResponseWriter, r *http.Request) {
provider, err := h.service.repository.GetSAMLProviderByCode(r.Context(), r.PathValue("provider_code"))
if err != nil {
http.NotFound(w, r)
return
}
sp, err := newSAMLServiceProvider(provider, nil, h.service.oidcHTTPClient())
if err != nil {
apiresponse.Error(w, http.StatusBadGateway, "SAML metadata 获取失败")
return
}
payload, err := xml.MarshalIndent(sp.Metadata(), "", " ")
if err != nil {
apiresponse.Error(w, http.StatusInternalServerError, "SAML SP metadata 生成失败")
return
}
w.Header().Set("Content-Type", "application/samlmetadata+xml; charset=utf-8")
w.Header().Set("X-Content-Type-Options", "nosniff")
_, _ = w.Write(append([]byte(xml.Header), payload...))
}
func (s *Service) samlServiceProvider(ctx context.Context, provider SAMLProvider) (*saml.ServiceProvider, error) {
metadata, err := s.fetchSAMLMetadata(ctx, provider.Config.MetadataURL)
if err != nil {
return nil, err
}
return newSAMLServiceProvider(provider, metadata, s.oidcHTTPClient())
}
func newSAMLServiceProvider(provider SAMLProvider, metadata *saml.EntityDescriptor, client *http.Client) (*saml.ServiceProvider, error) {
acsURL, err := url.Parse(provider.Config.ACSURL)
if err != nil {
return nil, err
}
metadataURL := *acsURL
metadataURL.Path = strings.TrimSuffix(metadataURL.Path, "/callback") + "/metadata"
metadataURL.RawQuery = ""
metadataURL.Fragment = ""
return &saml.ServiceProvider{
EntityID: provider.Config.SPEntityID, MetadataURL: metadataURL, AcsURL: *acsURL,
IDPMetadata: metadata, HTTPClient: client, AllowIDPInitiated: false,
SignatureVerifier: modernSAMLSignatureVerifier{},
}, nil
}
type modernSAMLSignatureVerifier struct{}
func (modernSAMLSignatureVerifier) VerifySignature(validationContext *dsig.ValidationContext, element *etree.Element) error {
if err := validateSAMLSignatureAlgorithms(element); err != nil {
return err
}
_, err := validationContext.Validate(element)
return err
}
func validateSAMLSignatureAlgorithms(element *etree.Element) error {
method := element.FindElement("./Signature/SignedInfo/SignatureMethod")
if method == nil {
return errors.New("SAML signature method missing")
}
switch method.SelectAttrValue("Algorithm", "") {
case dsig.RSASHA256SignatureMethod, dsig.RSASHA384SignatureMethod, dsig.RSASHA512SignatureMethod,
dsig.ECDSASHA256SignatureMethod, dsig.ECDSASHA384SignatureMethod, dsig.ECDSASHA512SignatureMethod:
default:
return errors.New("legacy or unsupported SAML signature algorithm")
}
allowedDigests := map[string]bool{
"http://www.w3.org/2001/04/xmlenc#sha256": true,
"http://www.w3.org/2001/04/xmldsig-more#sha384": true,
"http://www.w3.org/2001/04/xmlenc#sha512": true,
}
digests := element.FindElements("./Signature/SignedInfo/Reference/DigestMethod")
if len(digests) == 0 {
return errors.New("SAML digest method missing")
}
for _, digest := range digests {
if !allowedDigests[digest.SelectAttrValue("Algorithm", "")] {
return errors.New("legacy or unsupported SAML digest algorithm")
}
}
return nil
}
func (s *Service) fetchSAMLMetadata(ctx context.Context, rawURL string) (*saml.EntityDescriptor, error) {
target, err := validateOIDCURL(ctx, rawURL, s.allowPrivateIdentityProvider)
if err != nil {
return nil, err
}
now := time.Now()
s.samlMetadataMu.RLock()
cached, found := s.samlMetadata[target]
s.samlMetadataMu.RUnlock()
if found && now.Before(cached.ExpiresAt) {
return cached.Metadata, nil
}
req, err := http.NewRequestWithContext(ctx, http.MethodGet, target, nil)
if err != nil {
return nil, err
}
response, err := s.oidcHTTPClient().Do(req)
if err != nil {
return nil, err
}
defer response.Body.Close()
payload, err := io.ReadAll(io.LimitReader(response.Body, oidcMaxResponse+1))
if err != nil || response.StatusCode/100 != 2 || len(payload) > oidcMaxResponse {
return nil, errors.New("invalid SAML metadata response")
}
metadata, err := samlsp.ParseMetadata(payload)
if err != nil {
return nil, err
}
if err := validateSAMLMetadata(ctx, metadata, s.allowPrivateIdentityProvider, now); err != nil {
return nil, err
}
expiresAt := now.Add(time.Minute)
if !metadata.ValidUntil.IsZero() && metadata.ValidUntil.Before(expiresAt) {
expiresAt = metadata.ValidUntil
}
s.samlMetadataMu.Lock()
if s.samlMetadata == nil {
s.samlMetadata = make(map[string]samlMetadataCacheEntry)
}
s.samlMetadata[target] = samlMetadataCacheEntry{Metadata: metadata, ExpiresAt: expiresAt}
s.samlMetadataMu.Unlock()
return metadata, nil
}
func validateSAMLMetadata(ctx context.Context, metadata *saml.EntityDescriptor, allowPrivate bool, now time.Time) error {
if metadata == nil || strings.TrimSpace(metadata.EntityID) == "" || len(metadata.IDPSSODescriptors) == 0 {
return errors.New("metadata has no IDP descriptor")
}
if !metadata.ValidUntil.IsZero() && !metadata.ValidUntil.After(now) {
return errors.New("metadata expired")
}
descriptor := metadata.IDPSSODescriptors[0]
redirectFound := false
for _, endpoint := range descriptor.SingleSignOnServices {
if endpoint.Binding == saml.HTTPRedirectBinding {
if _, err := validateOIDCURL(ctx, endpoint.Location, allowPrivate); err != nil {
return err
}
redirectFound = true
}
}
if !redirectFound {
return errors.New("metadata has no redirect SSO endpoint")
}
validSigningCertificate := false
for _, key := range descriptor.KeyDescriptors {
if key.Use != "" && key.Use != "signing" {
continue
}
for _, encoded := range key.KeyInfo.X509Data.X509Certificates {
der, err := base64.StdEncoding.DecodeString(strings.Join(strings.Fields(encoded.Data), ""))
if err != nil {
continue
}
certificate, err := x509.ParseCertificate(der)
if err == nil && !now.Before(certificate.NotBefore) && now.Before(certificate.NotAfter) {
validSigningCertificate = true
}
}
}
if !validSigningCertificate {
return errors.New("metadata has no currently valid signing certificate")
}
return nil
}
func validateSAMLEntityID(raw string) (string, error) {
value := strings.TrimSpace(raw)
if value == "" || len(value) > 512 {
return "", errors.New("invalid entity ID")
}
parsed, err := url.Parse(value)
if err != nil || parsed.Scheme == "" || parsed.Fragment != "" {
return "", errors.New("invalid entity ID")
}
if (parsed.Scheme == "http" || parsed.Scheme == "https") && (parsed.Hostname() == "" || parsed.User != nil) {
return "", errors.New("invalid entity ID")
}
return value, nil
}
func samlAttribute(assertion *saml.Assertion, name string) string {
for _, statement := range assertion.AttributeStatements {
for _, attribute := range statement.Attributes {
if attribute.Name != name && attribute.FriendlyName != name {
continue
}
for _, value := range attribute.Values {
if strings.TrimSpace(value.Value) != "" {
return strings.TrimSpace(value.Value)
}
}
}
}
return ""
}
func (r *Repository) GetIdentityProviderKind(ctx context.Context, code string) (string, error) {
var kind string
err := r.pool.QueryRow(ctx, `SELECT kind FROM gateway.identity_providers WHERE code=$1 AND enabled`, strings.ToLower(strings.TrimSpace(code))).Scan(&kind)
return kind, mapRepositoryError(err)
}
func (r *Repository) ListPublicIdentityProviders(ctx context.Context) ([]PublicIdentityProvider, error) {
rows, err := r.pool.Query(ctx, `SELECT code,display_name,kind FROM gateway.identity_providers WHERE enabled ORDER BY code`)
if err != nil {
return nil, ErrUnavailable
}
defer rows.Close()
providers := []PublicIdentityProvider{}
for rows.Next() {
var provider PublicIdentityProvider
if err := rows.Scan(&provider.Code, &provider.DisplayName, &provider.Kind); err != nil {
return nil, ErrUnavailable
}
providers = append(providers, provider)
}
return providers, mapRepositoryError(rows.Err())
}
func (r *Repository) ListSAMLProviders(ctx context.Context) ([]SAMLProvider, error) {
rows, err := r.pool.Query(ctx, `SELECT id::text,code,display_name,portal_return_url,auto_provision,default_department_id::text,enabled,revision,config,created_at,updated_at FROM gateway.identity_providers WHERE kind='saml' ORDER BY code`)
if err != nil {
return nil, ErrUnavailable
}
defer rows.Close()
providers := []SAMLProvider{}
for rows.Next() {
provider, err := scanSAMLProvider(rows)
if err != nil {
return nil, err
}
providers = append(providers, provider)
}
return providers, mapRepositoryError(rows.Err())
}
func (r *Repository) GetSAMLProviderByCode(ctx context.Context, code string) (SAMLProvider, error) {
return scanSAMLProvider(r.pool.QueryRow(ctx, `SELECT id::text,code,display_name,portal_return_url,auto_provision,default_department_id::text,enabled,revision,config,created_at,updated_at FROM gateway.identity_providers WHERE kind='saml' AND code=$1`, strings.ToLower(strings.TrimSpace(code))))
}
type rowScanner interface {
Scan(dest ...any) error
}
func scanSAMLProvider(row rowScanner) (SAMLProvider, error) {
var provider SAMLProvider
var configJSON []byte
if err := row.Scan(&provider.ID, &provider.Code, &provider.DisplayName, &provider.PortalReturnURL, &provider.AutoProvision, &provider.DefaultDepartmentID, &provider.Enabled, &provider.Revision, &configJSON, &provider.CreatedAt, &provider.UpdatedAt); err != nil {
return provider, mapRepositoryError(err)
}
if json.Unmarshal(configJSON, &provider.Config) != nil {
return provider, ErrUnavailable
}
return provider, nil
}
func (r *Repository) CreateSAMLProvider(ctx context.Context, provider SAMLProvider, actor string) (SAMLProvider, error) {
id, err := platformid.NewUUID()
if err != nil {
return provider, ErrUnavailable
}
provider.ID = id
return r.storeSAMLProvider(ctx, provider, actor, true)
}
func (r *Repository) UpdateSAMLProvider(ctx context.Context, provider SAMLProvider, actor string) (SAMLProvider, error) {
return r.storeSAMLProvider(ctx, provider, actor, false)
}
func (r *Repository) storeSAMLProvider(ctx context.Context, provider SAMLProvider, actor string, creating bool) (SAMLProvider, error) {
configJSON, err := json.Marshal(provider.Config)
if err != nil {
return provider, ErrUnavailable
}
tx, err := r.pool.Begin(ctx)
if err != nil {
return provider, ErrUnavailable
}
defer tx.Rollback(ctx)
if creating {
err = tx.QueryRow(ctx, `INSERT INTO gateway.identity_providers(id,code,kind,display_name,portal_return_url,auto_provision,default_department_id,enabled,config) VALUES($1,$2,'saml',$3,$4,$5,$6,$7,$8) RETURNING revision,created_at,updated_at`, provider.ID, provider.Code, provider.DisplayName, provider.PortalReturnURL, provider.AutoProvision, provider.DefaultDepartmentID, provider.Enabled, configJSON).Scan(&provider.Revision, &provider.CreatedAt, &provider.UpdatedAt)
} else {
err = tx.QueryRow(ctx, `UPDATE gateway.identity_providers SET code=$2,display_name=$3,portal_return_url=$4,auto_provision=$5,default_department_id=$6,enabled=$7,config=$8,revision=revision+1,updated_at=clock_timestamp() WHERE id=$1 AND kind='saml' RETURNING revision,created_at,updated_at`, provider.ID, provider.Code, provider.DisplayName, provider.PortalReturnURL, provider.AutoProvision, provider.DefaultDepartmentID, provider.Enabled, configJSON).Scan(&provider.Revision, &provider.CreatedAt, &provider.UpdatedAt)
}
if err != nil {
return provider, mapManagementError(err)
}
eventID, err := platformid.NewUUID()
if err != nil {
return provider, ErrUnavailable
}
eventType := "identity_provider.updated"
if creating {
eventType = "identity_provider.created"
}
payload, _ := json.Marshal(map[string]any{"identity_provider_id": provider.ID, "kind": "saml", "actor_id": actor})
if _, err = tx.Exec(ctx, `INSERT INTO gateway.outbox_events(event_id,event_type,event_version,aggregate_type,aggregate_id,payload) VALUES($1,$2,1,'identity_provider',$3,$4)`, eventID, eventType, provider.ID, payload); err != nil {
return provider, ErrUnavailable
}
if err := tx.Commit(ctx); err != nil {
return provider, ErrUnavailable
}
return provider, nil
}
+115
View File
@@ -0,0 +1,115 @@
package identity
import (
"context"
"crypto/rand"
"crypto/rsa"
"crypto/x509"
"crypto/x509/pkix"
"encoding/base64"
"math/big"
"testing"
"time"
"github.com/beevik/etree"
"github.com/crewjam/saml"
dsig "github.com/russellhaering/goxmldsig"
)
func TestValidateSAMLMetadataRequiresLiveSigningCertificate(t *testing.T) {
now := time.Unix(1_800_000_000, 0)
certificate := makeTestCertificate(t, now.Add(-time.Hour), now.Add(time.Hour))
metadata := testIDPMetadata(certificate)
if err := validateSAMLMetadata(context.Background(), metadata, true, now); err != nil {
t.Fatalf("valid metadata rejected: %v", err)
}
metadata.IDPSSODescriptors[0].KeyDescriptors = nil
if err := validateSAMLMetadata(context.Background(), metadata, true, now); err == nil {
t.Fatal("metadata without signing certificate was accepted")
}
metadata = testIDPMetadata(makeTestCertificate(t, now.Add(-2*time.Hour), now.Add(-time.Hour)))
if err := validateSAMLMetadata(context.Background(), metadata, true, now); err == nil {
t.Fatal("expired signing certificate was accepted")
}
metadata = testIDPMetadata(certificate)
metadata.IDPSSODescriptors[0].SingleSignOnServices = nil
if err := validateSAMLMetadata(context.Background(), metadata, true, now); err == nil {
t.Fatal("metadata without redirect SSO endpoint was accepted")
}
}
func TestSAMLEntityIDAndAttributes(t *testing.T) {
if got, err := validateSAMLEntityID(" urn:example:gateway "); err != nil || got != "urn:example:gateway" {
t.Fatalf("valid URN entity ID rejected: %q %v", got, err)
}
for _, raw := range []string{"", "relative", "https://user:secret@example.com/sp", "urn:example:sp#fragment"} {
if _, err := validateSAMLEntityID(raw); err == nil {
t.Fatalf("invalid entity ID accepted: %q", raw)
}
}
assertion := &saml.Assertion{AttributeStatements: []saml.AttributeStatement{{Attributes: []saml.Attribute{{
FriendlyName: "mail", Name: "urn:oid:0.9.2342.19200300.100.1.3",
Values: []saml.AttributeValue{{Value: " user@example.com "}},
}}}}}
if got := samlAttribute(assertion, "mail"); got != "user@example.com" {
t.Fatalf("unexpected friendly-name attribute %q", got)
}
if got := samlAttribute(assertion, "urn:oid:0.9.2342.19200300.100.1.3"); got != "user@example.com" {
t.Fatalf("unexpected named attribute %q", got)
}
}
func TestSAMLSignatureAlgorithmPolicyRejectsSHA1(t *testing.T) {
document := etree.NewDocument()
if err := document.ReadFromString(`<Response><Signature><SignedInfo><SignatureMethod Algorithm="` + dsig.RSASHA256SignatureMethod + `"/><Reference><DigestMethod Algorithm="http://www.w3.org/2001/04/xmlenc#sha256"/></Reference></SignedInfo></Signature></Response>`); err != nil {
t.Fatal(err)
}
if err := validateSAMLSignatureAlgorithms(document.Root()); err != nil {
t.Fatalf("SHA-256 signature policy rejected: %v", err)
}
document.Root().FindElement("./Signature/SignedInfo/SignatureMethod").SelectAttr("Algorithm").Value = dsig.RSASHA1SignatureMethod
if err := validateSAMLSignatureAlgorithms(document.Root()); err == nil {
t.Fatal("SHA-1 signature algorithm was accepted")
}
document.Root().FindElement("./Signature/SignedInfo/SignatureMethod").SelectAttr("Algorithm").Value = dsig.RSASHA256SignatureMethod
document.Root().FindElement("./Signature/SignedInfo/Reference/DigestMethod").SelectAttr("Algorithm").Value = "http://www.w3.org/2000/09/xmldsig#sha1"
if err := validateSAMLSignatureAlgorithms(document.Root()); err == nil {
t.Fatal("SHA-1 digest algorithm was accepted")
}
}
func testIDPMetadata(certificate *x509.Certificate) *saml.EntityDescriptor {
return &saml.EntityDescriptor{
EntityID: "https://idp.example/metadata",
IDPSSODescriptors: []saml.IDPSSODescriptor{{
SSODescriptor: saml.SSODescriptor{RoleDescriptor: saml.RoleDescriptor{KeyDescriptors: []saml.KeyDescriptor{{
Use: "signing", KeyInfo: saml.KeyInfo{X509Data: saml.X509Data{X509Certificates: []saml.X509Certificate{{Data: base64.StdEncoding.EncodeToString(certificate.Raw)}}}},
}}}},
SingleSignOnServices: []saml.Endpoint{{Binding: saml.HTTPRedirectBinding, Location: "http://127.0.0.1:9091/sso"}},
}},
}
}
func makeTestCertificate(t *testing.T, notBefore, notAfter time.Time) *x509.Certificate {
t.Helper()
key, err := rsa.GenerateKey(rand.Reader, 2048)
if err != nil {
t.Fatal(err)
}
template := &x509.Certificate{
SerialNumber: big.NewInt(1), Subject: pkix.Name{CommonName: "SAML test"},
NotBefore: notBefore, NotAfter: notAfter, KeyUsage: x509.KeyUsageDigitalSignature,
}
der, err := x509.CreateCertificate(rand.Reader, template, template, &key.PublicKey, key)
if err != nil {
t.Fatal(err)
}
certificate, err := x509.ParseCertificate(der)
if err != nil {
t.Fatal(err)
}
return certificate
}
+362
View File
@@ -0,0 +1,362 @@
package identity
import (
"context"
"errors"
"fmt"
"net/http"
"strings"
"sync"
"time"
"aigateway.local/core/internal/platform/config"
"aigateway.local/core/internal/platform/cryptox"
)
var (
ErrInvalidCredentials = errors.New("invalid credentials")
ErrAccountDisabled = errors.New("account disabled")
ErrInvalidTOTP = errors.New("invalid or reused TOTP code")
ErrTOTPAlreadyEnabled = errors.New("TOTP is already enabled")
ErrTOTPNotEnabled = errors.New("TOTP is not enabled")
ErrTOTPSetupRequired = errors.New("TOTP setup is required")
)
const dummyPasswordHash = "pbkdf2_sha256$600000$00112233445566778899aabbccddeeff$afca0887b188255f525e15e30f5aa5a0b210a3e253bfaf9630411f0782bb6573"
type LockedError struct {
Until time.Time
}
func (e LockedError) Error() string {
return fmt.Sprintf("account locked until %s", e.Until.Format(time.RFC3339))
}
type LoginResult struct {
Token string
TempToken string
RequireTOTP bool
Account Account
}
type TOTPSetupResult struct {
Secret string
ProvisioningURI string
}
type Service struct {
repository *Repository
sessions *SessionStore
limiter *LoginLimiter
hasher PasswordHasher
config config.Auth
totpCipher cryptox.Cipher
idpCipher cryptox.Cipher
allowPrivateIdentityProvider bool
oidcClient *http.Client
samlMetadataMu sync.RWMutex
samlMetadata map[string]samlMetadataCacheEntry
now func() time.Time
}
func (s *Service) SetIdentityProviderCipher(cipher cryptox.Cipher, allowPrivate bool) {
s.idpCipher = cipher
s.allowPrivateIdentityProvider = allowPrivate
s.oidcClient = newOIDCHTTPClient(allowPrivate)
}
func NewService(repository *Repository, sessions *SessionStore, limiter *LoginLimiter, cfg config.Auth, totpCipher cryptox.Cipher) *Service {
return &Service{repository: repository, sessions: sessions, limiter: limiter, hasher: PasswordHasher{}, config: cfg, totpCipher: totpCipher, oidcClient: newOIDCHTTPClient(false), samlMetadata: make(map[string]samlMetadataCacheEntry), now: time.Now}
}
// AllowLogin reports whether a login attempt from ip may proceed. When the
// per-IP sliding-window limit is exceeded it returns false (caller responds 429).
func (s *Service) AllowLogin(ctx context.Context, ip string) bool {
return s.limiter == nil || s.limiter.Allow(ctx, ip)
}
func (s *Service) Login(ctx context.Context, kind Kind, login, password string) (LoginResult, error) {
account, err := s.findByLogin(ctx, kind, login)
if errors.Is(err, ErrNotFound) {
_ = s.hasher.Verify(password, dummyPasswordHash)
return LoginResult{}, ErrInvalidCredentials
}
if err != nil {
return LoginResult{}, err
}
if !account.Active {
return LoginResult{}, ErrAccountDisabled
}
if account.Locked(s.now()) {
return LoginResult{}, LockedError{Until: *account.LockedUntil}
}
if account.PasswordHash == "" || !s.hasher.Verify(password, account.PasswordHash) {
lockedUntil, recordErr := s.repository.RecordFailure(ctx, account, s.config.MaxFailures, s.config.LockDuration)
if recordErr != nil {
return LoginResult{}, recordErr
}
if lockedUntil != nil && lockedUntil.After(s.now()) {
return LoginResult{}, LockedError{Until: *lockedUntil}
}
return LoginResult{}, ErrInvalidCredentials
}
if account.TOTPEnabled {
principal := principalFor(account)
token, tokenErr := s.sessions.CreatePending(ctx, principal, s.config.TOTPChallengeTTL)
if tokenErr != nil {
return LoginResult{}, tokenErr
}
return LoginResult{TempToken: token, RequireTOTP: true, Account: account}, nil
}
var upgradedHash *string
if s.hasher.NeedsUpgrade(account.PasswordHash) {
hash, hashErr := s.hasher.Hash(password)
if hashErr != nil {
return LoginResult{}, hashErr
}
upgradedHash = &hash
}
principal := principalFor(account)
token, err := s.sessions.Create(ctx, principal)
if err != nil {
return LoginResult{}, err
}
if err := s.repository.CompleteLogin(ctx, account, upgradedHash); err != nil {
_ = s.sessions.Delete(ctx, "Bearer "+token)
return LoginResult{}, err
}
return LoginResult{Token: token, Account: account}, nil
}
func (s *Service) CompleteTOTPLogin(ctx context.Context, kind Kind, tempToken, code, backupCode string) (LoginResult, error) {
principal, err := s.sessions.AuthenticatePending(ctx, tempToken, kind)
if err != nil {
return LoginResult{}, err
}
account, err := s.findByID(ctx, kind, principal.SubjectID)
if err != nil {
return LoginResult{}, err
}
if !account.Active {
return LoginResult{}, ErrAccountDisabled
}
if account.Locked(s.now()) {
return LoginResult{}, LockedError{Until: *account.LockedUntil}
}
if !account.TOTPEnabled {
return LoginResult{}, ErrTOTPNotEnabled
}
valid, err := s.verifyAndConsumeFactor(ctx, account, code, backupCode)
if err != nil {
return LoginResult{}, err
}
if !valid {
lockedUntil, recordErr := s.repository.RecordFailure(ctx, account, s.config.MaxFailures, s.config.LockDuration)
if recordErr != nil {
return LoginResult{}, recordErr
}
if lockedUntil != nil && lockedUntil.After(s.now()) {
return LoginResult{}, LockedError{Until: *lockedUntil}
}
return LoginResult{}, ErrInvalidTOTP
}
token, err := s.sessions.Create(ctx, principalFor(account))
if err != nil {
return LoginResult{}, err
}
if err := s.repository.CompleteLogin(ctx, account, nil); err != nil {
_ = s.sessions.Delete(ctx, "Bearer "+token)
return LoginResult{}, err
}
_ = s.sessions.DeleteToken(ctx, tempToken)
return LoginResult{Token: token, Account: account}, nil
}
func (s *Service) SetupTOTP(ctx context.Context, account Account, password string) (TOTPSetupResult, error) {
if account.TOTPEnabled {
return TOTPSetupResult{}, ErrTOTPAlreadyEnabled
}
if !s.hasher.Verify(password, account.PasswordHash) {
return TOTPSetupResult{}, ErrInvalidCredentials
}
secret, err := GenerateTOTPSecret()
if err != nil {
return TOTPSetupResult{}, err
}
encrypted, version, err := s.totpCipher.Encrypt([]byte(secret))
if err != nil {
return TOTPSetupResult{}, err
}
if err := s.repository.SetTOTPSecret(ctx, account, encrypted, version); err != nil {
return TOTPSetupResult{}, err
}
return TOTPSetupResult{Secret: secret, ProvisioningURI: TOTPProvisioningURI(secret, account.Kind, account.Login)}, nil
}
func (s *Service) ConfirmTOTP(ctx context.Context, account Account, code string) ([]string, error) {
account, err := s.findByID(ctx, account.Kind, account.ID)
if err != nil {
return nil, err
}
if account.TOTPEnabled {
return nil, ErrTOTPAlreadyEnabled
}
secret, err := s.decryptTOTPSecret(account)
if err != nil {
return nil, err
}
step, valid := VerifyTOTP(secret, code, s.now())
if !valid {
return nil, ErrInvalidTOTP
}
codes, records, err := GenerateBackupCodes()
if err != nil {
return nil, err
}
if err := s.repository.EnableTOTP(ctx, account, step, records); err != nil {
return nil, err
}
return codes, nil
}
func (s *Service) DisableTOTP(ctx context.Context, account Account, password, code, backupCode string) error {
account, err := s.findByID(ctx, account.Kind, account.ID)
if err != nil {
return err
}
if !account.TOTPEnabled {
return ErrTOTPNotEnabled
}
if !s.hasher.Verify(password, account.PasswordHash) {
return ErrInvalidCredentials
}
valid, err := s.verifyAndConsumeFactor(ctx, account, code, backupCode)
if err != nil {
return err
}
if !valid {
return ErrInvalidTOTP
}
return s.repository.DisableTOTP(ctx, account)
}
func (s *Service) RegenerateBackupCodes(ctx context.Context, account Account, password, code, backupCode string) ([]string, error) {
account, err := s.findByID(ctx, account.Kind, account.ID)
if err != nil {
return nil, err
}
if !account.TOTPEnabled {
return nil, ErrTOTPNotEnabled
}
if !s.hasher.Verify(password, account.PasswordHash) {
return nil, ErrInvalidCredentials
}
valid, err := s.verifyAndConsumeFactor(ctx, account, code, backupCode)
if err != nil {
return nil, err
}
if !valid {
return nil, ErrInvalidTOTP
}
codes, records, err := GenerateBackupCodes()
if err != nil {
return nil, err
}
if err := s.repository.ReplaceBackupCodes(ctx, account, records); err != nil {
return nil, err
}
return codes, nil
}
func (s *Service) verifyAndConsumeFactor(ctx context.Context, account Account, code, backupCode string) (bool, error) {
if strings.TrimSpace(backupCode) != "" {
return s.repository.ConsumeBackupCode(ctx, account, HashBackupCode(backupCode))
}
secret, err := s.decryptTOTPSecret(account)
if err != nil {
return false, err
}
step, valid := VerifyTOTP(secret, code, s.now())
if !valid {
return false, nil
}
return s.repository.ConsumeTOTPStep(ctx, account, step)
}
func (s *Service) decryptTOTPSecret(account Account) (string, error) {
if len(account.EncryptedTOTPSecret) == 0 || account.TOTPKekVersion == nil {
return "", ErrTOTPSetupRequired
}
plaintext, err := s.totpCipher.Decrypt(account.EncryptedTOTPSecret, *account.TOTPKekVersion)
if err != nil {
return "", err
}
return string(plaintext), nil
}
func (s *Service) Authenticate(ctx context.Context, kind Kind, authorization string) (Account, error) {
principal, err := s.sessions.Authenticate(ctx, authorization, kind)
if err != nil {
return Account{}, err
}
var account Account
if kind == KindAdmin {
account, err = s.repository.FindAdminByID(ctx, principal.SubjectID)
} else {
account, err = s.repository.FindPortalByID(ctx, principal.SubjectID)
}
if errors.Is(err, ErrNotFound) {
return Account{}, ErrInvalidSession
}
if err != nil {
return Account{}, err
}
if !account.Active {
return Account{}, ErrAccountDisabled
}
return account, nil
}
func (s *Service) Logout(ctx context.Context, authorization string) error {
return s.sessions.Delete(ctx, authorization)
}
// ChangePassword updates the authenticated account password after verifying the
// current password. Externally provisioned portal accounts without a local
// password may set their first password without an old-password check.
func (s *Service) ChangePassword(ctx context.Context, account Account, oldPassword, newPassword string) error {
account, err := s.findByID(ctx, account.Kind, account.ID)
if err != nil {
return err
}
if account.PasswordHash != "" && !s.hasher.Verify(oldPassword, account.PasswordHash) {
return ErrInvalidCredentials
}
if len(newPassword) < 12 || len(newPassword) > 1024 {
return errors.New("new password must contain 12 to 1024 characters")
}
hash, err := s.hasher.Hash(newPassword)
if err != nil {
return err
}
return s.repository.SetPassword(ctx, account, hash)
}
func (s *Service) findByLogin(ctx context.Context, kind Kind, login string) (Account, error) {
if kind == KindAdmin {
return s.repository.FindAdminByLogin(ctx, login)
}
return s.repository.FindPortalByLogin(ctx, login)
}
func (s *Service) findByID(ctx context.Context, kind Kind, id string) (Account, error) {
if kind == KindAdmin {
return s.repository.FindAdminByID(ctx, id)
}
return s.repository.FindPortalByID(ctx, id)
}
func principalFor(account Account) Principal {
return Principal{Kind: account.Kind, SubjectID: account.ID, Login: account.Login, DisplayName: account.DisplayName, Role: account.Role}
}
+211
View File
@@ -0,0 +1,211 @@
package identity
import (
"context"
"crypto/rand"
"crypto/sha256"
"encoding/base64"
"encoding/hex"
"encoding/json"
"errors"
"fmt"
"strings"
"time"
"github.com/redis/go-redis/v9"
)
var ErrInvalidSession = errors.New("invalid or expired session")
type Principal struct {
Kind Kind `json:"kind"`
SubjectID string `json:"subject_id"`
Login string `json:"login"`
DisplayName string `json:"display_name"`
Role string `json:"role,omitempty"`
Purpose string `json:"purpose"`
IssuedAt int64 `json:"issued_at"`
}
type SessionStore struct {
client *redis.Client
ttl time.Duration
}
func NewSessionStore(client *redis.Client, ttl time.Duration) *SessionStore {
return &SessionStore{client: client, ttl: ttl}
}
func (s *SessionStore) Create(ctx context.Context, principal Principal) (string, error) {
principal.Purpose = "session"
return s.create(ctx, principal, s.ttl)
}
func (s *SessionStore) CreatePending(ctx context.Context, principal Principal, ttl time.Duration) (string, error) {
principal.Purpose = "totp_pending"
return s.create(ctx, principal, ttl)
}
func (s *SessionStore) create(ctx context.Context, principal Principal, ttl time.Duration) (string, error) {
if s.client == nil {
return "", ErrUnavailable
}
random := make([]byte, 32)
if _, err := rand.Read(random); err != nil {
return "", err
}
token := base64.RawURLEncoding.EncodeToString(random)
principal.IssuedAt = time.Now().Unix()
payload, err := json.Marshal(principal)
if err != nil {
return "", err
}
if err := s.client.Set(ctx, sessionKey(token), payload, ttl).Err(); err != nil {
return "", fmt.Errorf("%w: %v", ErrUnavailable, err)
}
return token, nil
}
func (s *SessionStore) Authenticate(ctx context.Context, authorization string, expected Kind) (Principal, error) {
if s.client == nil {
return Principal{}, ErrUnavailable
}
token, ok := bearerToken(authorization)
if !ok {
return Principal{}, ErrInvalidSession
}
payload, err := s.client.Get(ctx, sessionKey(token)).Bytes()
if errors.Is(err, redis.Nil) {
return Principal{}, ErrInvalidSession
}
if err != nil {
return Principal{}, fmt.Errorf("%w: %v", ErrUnavailable, err)
}
var principal Principal
if err := json.Unmarshal(payload, &principal); err != nil || principal.Kind != expected || principal.SubjectID == "" || principal.Purpose != "session" {
return Principal{}, ErrInvalidSession
}
return principal, nil
}
func (s *SessionStore) AuthenticatePending(ctx context.Context, token string, expected Kind) (Principal, error) {
principal, err := s.authenticateToken(ctx, token)
if err != nil || principal.Kind != expected || principal.SubjectID == "" || principal.Purpose != "totp_pending" {
return Principal{}, ErrInvalidSession
}
return principal, nil
}
func (s *SessionStore) authenticateToken(ctx context.Context, token string) (Principal, error) {
if s.client == nil {
return Principal{}, ErrUnavailable
}
payload, err := s.client.Get(ctx, sessionKey(strings.TrimSpace(token))).Bytes()
if errors.Is(err, redis.Nil) {
return Principal{}, ErrInvalidSession
}
if err != nil {
return Principal{}, fmt.Errorf("%w: %v", ErrUnavailable, err)
}
var principal Principal
if err := json.Unmarshal(payload, &principal); err != nil {
return Principal{}, ErrInvalidSession
}
return principal, nil
}
func (s *SessionStore) DeleteToken(ctx context.Context, token string) error {
if s.client == nil {
return ErrUnavailable
}
if err := s.client.Del(ctx, sessionKey(strings.TrimSpace(token))).Err(); err != nil {
return fmt.Errorf("%w: %v", ErrUnavailable, err)
}
return nil
}
func (s *SessionStore) Delete(ctx context.Context, authorization string) error {
if s.client == nil {
return ErrUnavailable
}
token, ok := bearerToken(authorization)
if !ok {
return nil
}
if err := s.client.Del(ctx, sessionKey(token)).Err(); err != nil {
return fmt.Errorf("%w: %v", ErrUnavailable, err)
}
return nil
}
func (s *SessionStore) StoreOneTime(ctx context.Context, namespace string, value any, ttl time.Duration) (string, error) {
if s.client == nil {
return "", ErrUnavailable
}
random := make([]byte, 32)
if _, err := rand.Read(random); err != nil {
return "", err
}
token := base64.RawURLEncoding.EncodeToString(random)
payload, err := json.Marshal(value)
if err != nil {
return "", err
}
if err := s.client.Set(ctx, oneTimeKey(namespace, token), payload, ttl).Err(); err != nil {
return "", fmt.Errorf("%w: %v", ErrUnavailable, err)
}
return token, nil
}
func (s *SessionStore) ConsumeOneTime(ctx context.Context, namespace, token string, target any) error {
if s.client == nil {
return ErrUnavailable
}
payload, err := s.client.GetDel(ctx, oneTimeKey(namespace, strings.TrimSpace(token))).Bytes()
if errors.Is(err, redis.Nil) {
return ErrInvalidSession
}
if err != nil {
return fmt.Errorf("%w: %v", ErrUnavailable, err)
}
if err := json.Unmarshal(payload, target); err != nil {
return ErrInvalidSession
}
return nil
}
// ClaimIdentifier atomically claims a caller-provided replay identifier for
// the TTL. It is used for signed protocol message IDs, not bearer secrets.
func (s *SessionStore) ClaimIdentifier(ctx context.Context, namespace, identifier string, ttl time.Duration) (bool, error) {
if s.client == nil {
return false, ErrUnavailable
}
identifier = strings.TrimSpace(identifier)
if identifier == "" || len(identifier) > 512 {
return false, ErrInvalidSession
}
claimed, err := s.client.SetNX(ctx, oneTimeKey(namespace, identifier), "1", ttl).Result()
if err != nil {
return false, fmt.Errorf("%w: %v", ErrUnavailable, err)
}
return claimed, nil
}
func bearerToken(authorization string) (string, bool) {
authorization = strings.TrimSpace(authorization)
if len(authorization) <= len("Bearer ") || !strings.EqualFold(authorization[:len("Bearer ")], "Bearer ") {
return "", false
}
token := strings.TrimSpace(authorization[len("Bearer "):])
return token, token != ""
}
func sessionKey(token string) string {
digest := sha256.Sum256([]byte(token))
return "gateway:session:v1:" + hex.EncodeToString(digest[:])
}
func oneTimeKey(namespace, token string) string {
digest := sha256.Sum256([]byte(token))
return "gateway:one-time:" + namespace + ":" + hex.EncodeToString(digest[:])
}
+118
View File
@@ -0,0 +1,118 @@
package identity
import (
"crypto/hmac"
"crypto/rand"
"crypto/sha1"
"crypto/sha256"
"crypto/subtle"
"encoding/base32"
"encoding/binary"
"encoding/hex"
"errors"
"fmt"
"net/url"
"strings"
"time"
)
const (
totpPeriod = int64(30)
totpDigits = 6
totpWindow = int64(1)
backupCodeCount = 10
backupCodeLength = 8
)
var backupAlphabet = []byte("23456789ABCDEFGHJKLMNPQRSTUVWXYZ")
type BackupCodeRecord struct {
Hash string `json:"hash"`
UsedAt *time.Time `json:"used_at,omitempty"`
}
func GenerateTOTPSecret() (string, error) {
secret := make([]byte, 20)
if _, err := rand.Read(secret); err != nil {
return "", err
}
return base32.StdEncoding.WithPadding(base32.NoPadding).EncodeToString(secret), nil
}
func TOTPProvisioningURI(secret string, kind Kind, login string) string {
issuer := "AI Gateway"
label := issuer + ":" + string(kind) + ":" + login
query := url.Values{
"secret": []string{secret},
"issuer": []string{issuer},
"algorithm": []string{"SHA1"},
"digits": []string{fmt.Sprint(totpDigits)},
"period": []string{fmt.Sprint(totpPeriod)},
}
return "otpauth://totp/" + url.PathEscape(label) + "?" + query.Encode()
}
func VerifyTOTP(secret, code string, now time.Time) (int64, bool) {
code = strings.TrimSpace(code)
if len(code) != totpDigits {
return 0, false
}
step := now.Unix() / totpPeriod
for offset := -totpWindow; offset <= totpWindow; offset++ {
candidate, err := hotp(secret, step+offset, totpDigits)
if err == nil && subtle.ConstantTimeCompare([]byte(candidate), []byte(code)) == 1 {
return step + offset, true
}
}
return 0, false
}
func hotp(secret string, counter int64, digits int) (string, error) {
if counter < 0 || digits < 6 || digits > 8 {
return "", errors.New("invalid HOTP parameters")
}
decoded, err := base32.StdEncoding.WithPadding(base32.NoPadding).DecodeString(strings.ToUpper(strings.TrimSpace(secret)))
if err != nil {
return "", err
}
message := make([]byte, 8)
binary.BigEndian.PutUint64(message, uint64(counter))
mac := hmac.New(sha1.New, decoded)
_, _ = mac.Write(message)
digest := mac.Sum(nil)
offset := digest[len(digest)-1] & 0x0f
value := (uint32(digest[offset])&0x7f)<<24 |
uint32(digest[offset+1])<<16 |
uint32(digest[offset+2])<<8 |
uint32(digest[offset+3])
modulus := uint32(1)
for i := 0; i < digits; i++ {
modulus *= 10
}
return fmt.Sprintf("%0*d", digits, value%modulus), nil
}
func GenerateBackupCodes() ([]string, []BackupCodeRecord, error) {
codes := make([]string, 0, backupCodeCount)
records := make([]BackupCodeRecord, 0, backupCodeCount)
for range backupCodeCount {
random := make([]byte, backupCodeLength)
if _, err := rand.Read(random); err != nil {
return nil, nil, err
}
for index := range random {
random[index] = backupAlphabet[int(random[index])%len(backupAlphabet)]
}
raw := string(random)
code := raw[:4] + "-" + raw[4:]
codes = append(codes, code)
records = append(records, BackupCodeRecord{Hash: HashBackupCode(code)})
}
return codes, records, nil
}
func HashBackupCode(code string) string {
normalized := strings.ToUpper(strings.ReplaceAll(strings.TrimSpace(code), "-", ""))
digest := sha256.Sum256([]byte(normalized))
return hex.EncodeToString(digest[:])
}
+43
View File
@@ -0,0 +1,43 @@
package identity
import (
"testing"
"time"
)
func TestHOTPUsesRFC6238Vector(t *testing.T) {
// RFC 6238 SHA-1 shared secret, time 59 seconds, 8-digit expected value.
code, err := hotp("GEZDGNBVGY3TQOJQGEZDGNBVGY3TQOJQ", 59/30, 8)
if err != nil {
t.Fatal(err)
}
if code != "94287082" {
t.Fatalf("got %s", code)
}
}
func TestVerifyTOTPAcceptsWindow(t *testing.T) {
now := time.Unix(1_700_000_000, 0)
secret := "GEZDGNBVGY3TQOJQGEZDGNBVGY3TQOJQ"
code, err := hotp(secret, now.Unix()/30-1, 6)
if err != nil {
t.Fatal(err)
}
step, ok := VerifyTOTP(secret, code, now)
if !ok || step != now.Unix()/30-1 {
t.Fatalf("step=%d ok=%v", step, ok)
}
}
func TestBackupCodeNormalization(t *testing.T) {
if HashBackupCode("abcd-2345") != HashBackupCode(" ABCD2345 ") {
t.Fatal("backup code normalization differs")
}
codes, records, err := GenerateBackupCodes()
if err != nil {
t.Fatal(err)
}
if len(codes) != 10 || len(records) != 10 || records[0].Hash != HashBackupCode(codes[0]) {
t.Fatal("invalid backup code generation")
}
}
+84
View File
@@ -0,0 +1,84 @@
package operations
import (
"context"
"net/http"
"runtime"
"time"
"aigateway.local/core/internal/identity"
"aigateway.local/core/internal/platform/apiresponse"
"github.com/jackc/pgx/v5/pgxpool"
)
type AdminHTTPHandler struct {
pool *pgxpool.Pool
identity *identity.Service
version string
startedAt time.Time
reload func(context.Context) error
mux *http.ServeMux
}
func NewAdminHTTPHandler(pool *pgxpool.Pool, identityService *identity.Service, version string, startedAt time.Time, reload func(context.Context) error) *AdminHTTPHandler {
h := &AdminHTTPHandler{pool: pool, identity: identityService, version: version, startedAt: startedAt, reload: reload, mux: http.NewServeMux()}
h.mux.HandleFunc("GET /api/v1/admin/system-info", h.systemInfo)
h.mux.HandleFunc("GET /api/v1/admin/monitoring/overview", h.overview)
h.mux.HandleFunc("POST /api/v1/admin/reload", h.reloadSnapshots)
return h
}
func (h *AdminHTTPHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { h.mux.ServeHTTP(w, r) }
func (h *AdminHTTPHandler) account(w http.ResponseWriter, r *http.Request) (identity.Account, bool) {
account, err := h.identity.Authenticate(r.Context(), identity.KindAdmin, r.Header.Get("Authorization"))
if err != nil {
apiresponse.Error(w, http.StatusUnauthorized, "登录状态无效")
return identity.Account{}, false
}
return account, true
}
func (h *AdminHTTPHandler) systemInfo(w http.ResponseWriter, r *http.Request) {
if _, ok := h.account(w, r); !ok {
return
}
var providers, models, keys int64
if err := h.pool.QueryRow(r.Context(), `SELECT (SELECT count(*) FROM gateway.providers WHERE enabled),(SELECT count(*) FROM gateway.provider_models WHERE enabled),(SELECT count(*) FROM gateway.api_keys WHERE enabled)`).Scan(&providers, &models, &keys); err != nil {
apiresponse.Error(w, http.StatusServiceUnavailable, "系统信息查询失败")
return
}
apiresponse.OK(w, map[string]any{"version": h.version, "go_version": runtime.Version(), "uptime_seconds": int64(time.Since(h.startedAt).Seconds()), "database": "postgresql", "object_storage": false, "clickhouse": false, "enabled_providers": providers, "enabled_models": models, "enabled_api_keys": keys})
}
func (h *AdminHTTPHandler) overview(w http.ResponseWriter, r *http.Request) {
if _, ok := h.account(w, r); !ok {
return
}
var requests, failures, promptTokens, completionTokens, cost int64
var latency *float64
err := h.pool.QueryRow(r.Context(), `SELECT count(*),count(*) FILTER(WHERE coalesce(status_code,500)>=400),coalesce(sum(prompt_tokens),0),coalesce(sum(completion_tokens),0),coalesce(sum(cost_microunits),0),avg(latency_ms) FROM gateway.audit_events WHERE recorded_at>=clock_timestamp()-interval '24 hours'`).Scan(&requests, &failures, &promptTokens, &completionTokens, &cost, &latency)
if err != nil {
apiresponse.Error(w, http.StatusServiceUnavailable, "监控汇总查询失败")
return
}
apiresponse.OK(w, map[string]any{"window": "24h", "requests": requests, "failed_requests": failures, "prompt_tokens": promptTokens, "completion_tokens": completionTokens, "cost_microunits": cost, "avg_latency_ms": latency})
}
func (h *AdminHTTPHandler) reloadSnapshots(w http.ResponseWriter, r *http.Request) {
account, ok := h.account(w, r)
if !ok {
return
}
if !identity.HasPermission(account, identity.PermissionProviderManage) {
apiresponse.Error(w, http.StatusForbidden, "缺少运行时配置管理权限")
return
}
if h.reload != nil {
if err := h.reload(r.Context()); err != nil {
apiresponse.Error(w, http.StatusServiceUnavailable, "运行时快照刷新不完整: "+err.Error())
return
}
}
apiresponse.OK(w, map[string]bool{"reloaded": true})
}
+95
View File
@@ -0,0 +1,95 @@
package outbox
import (
"errors"
"net/http"
"strconv"
"strings"
"aigateway.local/core/internal/identity"
"aigateway.local/core/internal/platform/apiresponse"
"github.com/jackc/pgx/v5/pgtype"
)
type AdminHTTPHandler struct {
store *Store
identity *identity.Service
mux *http.ServeMux
}
func NewAdminHTTPHandler(store *Store, identityService *identity.Service) *AdminHTTPHandler {
h := &AdminHTTPHandler{store: store, identity: identityService, mux: http.NewServeMux()}
h.mux.HandleFunc("GET /api/v1/admin/outbox-events", h.list)
h.mux.HandleFunc("POST /api/v1/admin/outbox-events/{event_id}/retry", h.retry)
return h
}
func (h *AdminHTTPHandler) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
h.mux.ServeHTTP(writer, request)
}
func (h *AdminHTTPHandler) list(writer http.ResponseWriter, request *http.Request) {
if !h.requirePermission(writer, request, identity.PermissionOutboxRead) {
return
}
status := strings.TrimSpace(request.URL.Query().Get("status"))
if status != "" && status != "pending" && status != "dead" && status != "processed" {
apiresponse.Error(writer, http.StatusBadRequest, "status 必须是 pending、dead 或 processed")
return
}
limit := 100
if value := request.URL.Query().Get("limit"); value != "" {
parsed, err := strconv.Atoi(value)
if err != nil || parsed < 1 || parsed > 500 {
apiresponse.Error(writer, http.StatusBadRequest, "limit 必须在 1 到 500 之间")
return
}
limit = parsed
}
items, err := h.store.List(request.Context(), status, strings.TrimSpace(request.URL.Query().Get("type")), limit)
if err != nil {
apiresponse.Error(writer, http.StatusServiceUnavailable, "事件投递查询服务暂不可用")
return
}
apiresponse.OK(writer, items)
}
func (h *AdminHTTPHandler) retry(writer http.ResponseWriter, request *http.Request) {
if !h.requirePermission(writer, request, identity.PermissionOutboxManage) {
return
}
eventID := strings.TrimSpace(request.PathValue("event_id"))
var parsedID pgtype.UUID
if err := parsedID.Scan(eventID); err != nil || !parsedID.Valid {
apiresponse.Error(writer, http.StatusBadRequest, "event_id 必须是有效 UUID")
return
}
if err := h.store.Retry(request.Context(), eventID); err != nil {
if errors.Is(err, ErrEventNotFound) {
apiresponse.Error(writer, http.StatusNotFound, "待处理或死信事件不存在")
return
}
apiresponse.Error(writer, http.StatusServiceUnavailable, "事件重试服务暂不可用")
return
}
apiresponse.OK(writer, map[string]bool{"retried": true})
}
func (h *AdminHTTPHandler) requirePermission(writer http.ResponseWriter, request *http.Request, permission string) bool {
account, err := h.identity.Authenticate(request.Context(), identity.KindAdmin, request.Header.Get("Authorization"))
if err != nil {
if errors.Is(err, identity.ErrInvalidSession) || errors.Is(err, identity.ErrNotFound) {
apiresponse.Error(writer, http.StatusUnauthorized, "登录状态无效或已过期")
} else if errors.Is(err, identity.ErrAccountDisabled) {
apiresponse.Error(writer, http.StatusForbidden, "管理员账号已被停用")
} else {
apiresponse.Error(writer, http.StatusServiceUnavailable, "身份服务暂不可用")
}
return false
}
if !identity.HasPermission(account, permission) {
apiresponse.Error(writer, http.StatusForbidden, "缺少事件投递操作权限")
return false
}
return true
}
@@ -0,0 +1,55 @@
package outbox
import (
"context"
"errors"
"os"
"testing"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgxpool"
)
func TestConsumeTransactionRollbackAndIdempotency(t *testing.T) {
databaseURL := os.Getenv("OUTBOX_TEST_DATABASE_URL")
if databaseURL == "" {
t.Skip("OUTBOX_TEST_DATABASE_URL is not configured")
}
ctx := context.Background()
pool, err := pgxpool.New(ctx, databaseURL)
if err != nil {
t.Fatal(err)
}
defer pool.Close()
store := NewStore(pool)
const subscriber = "outbox-integration-test"
const eventID = "33333333-3333-4333-8333-333333333333"
_, _ = pool.Exec(ctx, `DELETE FROM gateway.event_consumptions WHERE subscriber=$1 AND event_id=$2`, subscriber, eventID)
t.Cleanup(func() {
_, _ = pool.Exec(context.Background(), `DELETE FROM gateway.event_consumptions WHERE subscriber=$1 AND event_id=$2`, subscriber, eventID)
})
calls := 0
wantErr := errors.New("rollback handler")
consumed, err := store.Consume(ctx, subscriber, eventID, func(context.Context, pgx.Tx) error {
calls++
return wantErr
})
if consumed || !errors.Is(err, wantErr) {
t.Fatalf("failed handler must roll back: consumed=%v err=%v", consumed, err)
}
consumed, err = store.Consume(ctx, subscriber, eventID, func(context.Context, pgx.Tx) error {
calls++
return nil
})
if !consumed || err != nil {
t.Fatalf("second attempt must consume: consumed=%v err=%v", consumed, err)
}
consumed, err = store.Consume(ctx, subscriber, eventID, func(context.Context, pgx.Tx) error {
calls++
return nil
})
if consumed || err != nil || calls != 2 {
t.Fatalf("duplicate must skip handler: consumed=%v calls=%d err=%v", consumed, calls, err)
}
}
+98
View File
@@ -0,0 +1,98 @@
package outbox
import (
"context"
"errors"
"fmt"
"strconv"
"time"
"github.com/redis/go-redis/v9"
)
var ErrPublisherUnavailable = errors.New("outbox publisher unavailable")
type PublishResult struct {
Published bool
StreamID string
}
type RedisPublisher struct {
client *redis.Client
stream string
maxLength int64
markerTTL time.Duration
script *redis.Script
}
func NewRedisPublisher(client *redis.Client, stream string, maxLength int64, markerTTL time.Duration) *RedisPublisher {
return &RedisPublisher{client: client, stream: stream, maxLength: maxLength, markerTTL: markerTTL, script: redis.NewScript(publishScript)}
}
// Ping reports whether the downstream Redis stream is reachable. The outbox
// worker calls it before claiming events so a Redis outage never consumes the
// events' delivery budget (claiming increments attempts; dead-lettering then
// burns the whole queue for a fault that was never the events').
func (p *RedisPublisher) Ping(ctx context.Context) error {
if p == nil || p.client == nil {
return ErrPublisherUnavailable
}
if err := p.client.Ping(ctx).Err(); err != nil {
return fmt.Errorf("%w: %v", ErrPublisherUnavailable, err)
}
return nil
}
func (p *RedisPublisher) Publish(ctx context.Context, event Event) (PublishResult, error) {
if p == nil || p.client == nil {
return PublishResult{}, ErrPublisherUnavailable
}
marker := "gateway:{outbox}:published:" + event.EventID
result, err := p.script.Run(ctx, p.client, []string{marker, p.stream},
int64(p.markerTTL.Seconds()), p.maxLength, event.EventID, event.EventType, event.EventVersion,
valueOrEmpty(event.TenantID), event.AggregateType, event.AggregateID, string(event.Payload), string(event.TraceContext), event.OccurredAt.UTC().Format(time.RFC3339Nano),
).Slice()
if err != nil || len(result) != 2 {
return PublishResult{}, fmt.Errorf("%w: %v", ErrPublisherUnavailable, err)
}
published, err := redisInt(result[0])
if err != nil {
return PublishResult{}, fmt.Errorf("%w: %v", ErrPublisherUnavailable, err)
}
streamID := "duplicate"
if result[1] != nil && fmt.Sprint(result[1]) != "" {
streamID = fmt.Sprint(result[1])
}
return PublishResult{Published: published == 1, StreamID: streamID}, nil
}
func valueOrEmpty(value *string) string {
if value == nil {
return ""
}
return *value
}
func redisInt(value any) (int64, error) {
switch typed := value.(type) {
case int64:
return typed, nil
case string:
return strconv.ParseInt(typed, 10, 64)
case []byte:
return strconv.ParseInt(string(typed), 10, 64)
default:
return 0, fmt.Errorf("unexpected redis integer %T", value)
}
}
const publishScript = `
if redis.call('SET', KEYS[1], '1', 'NX', 'EX', tonumber(ARGV[1])) then
local id = redis.call('XADD', KEYS[2], 'MAXLEN', '~', tonumber(ARGV[2]), '*',
'event_id', ARGV[3], 'event_type', ARGV[4], 'event_version', ARGV[5],
'tenant_id', ARGV[6], 'aggregate_type', ARGV[7], 'aggregate_id', ARGV[8],
'payload', ARGV[9], 'trace_context', ARGV[10], 'occurred_at', ARGV[11])
return {1, id}
end
return {0, ''}
`
+201
View File
@@ -0,0 +1,201 @@
package outbox
import (
"context"
"encoding/json"
"errors"
"fmt"
"strings"
"time"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgxpool"
)
var (
ErrStoreUnavailable = errors.New("outbox store unavailable")
ErrEventNotFound = errors.New("outbox event not found")
)
type Event struct {
EventID string `json:"event_id"`
EventType string `json:"event_type"`
EventVersion int `json:"event_version"`
TenantID *string `json:"tenant_id"`
AggregateType string `json:"aggregate_type"`
AggregateID string `json:"aggregate_id"`
Payload json.RawMessage `json:"payload"`
TraceContext json.RawMessage `json:"trace_context"`
OccurredAt time.Time `json:"occurred_at"`
Attempts int `json:"attempts"`
}
type Store struct{ pool *pgxpool.Pool }
func NewStore(pool *pgxpool.Pool) *Store { return &Store{pool: pool} }
func (s *Store) Claim(ctx context.Context, workerID string, limit int, lease time.Duration) ([]Event, error) {
if s == nil || s.pool == nil {
return nil, ErrStoreUnavailable
}
rows, err := s.pool.Query(ctx, `
WITH candidates AS (
SELECT event_id FROM gateway.outbox_events
WHERE processed_at IS NULL AND dead_lettered_at IS NULL AND available_at <= clock_timestamp()
AND (locked_at IS NULL OR locked_at < clock_timestamp()-($1 * interval '1 second'))
ORDER BY available_at,occurred_at
FOR UPDATE SKIP LOCKED LIMIT $2
)
UPDATE gateway.outbox_events e
SET locked_at=clock_timestamp(),locked_by=$3,attempts=e.attempts+1
FROM candidates c WHERE e.event_id=c.event_id
RETURNING e.event_id::text,e.event_type,e.event_version,e.tenant_id::text,e.aggregate_type,
e.aggregate_id,e.payload,e.trace_context,e.occurred_at,e.attempts`, lease.Seconds(), limit, workerID)
if err != nil {
return nil, fmt.Errorf("%w: claim: %v", ErrStoreUnavailable, err)
}
defer rows.Close()
events := make([]Event, 0, limit)
for rows.Next() {
var event Event
if err := rows.Scan(&event.EventID, &event.EventType, &event.EventVersion, &event.TenantID,
&event.AggregateType, &event.AggregateID, &event.Payload, &event.TraceContext, &event.OccurredAt, &event.Attempts); err != nil {
return nil, fmt.Errorf("%w: scan claim: %v", ErrStoreUnavailable, err)
}
events = append(events, event)
}
if err := rows.Err(); err != nil {
return nil, fmt.Errorf("%w: claim rows: %v", ErrStoreUnavailable, err)
}
return events, nil
}
func (s *Store) MarkProcessed(ctx context.Context, eventID, workerID, streamID string) error {
if s == nil || s.pool == nil {
return ErrStoreUnavailable
}
result, err := s.pool.Exec(ctx, `UPDATE gateway.outbox_events SET processed_at=clock_timestamp(),locked_at=NULL,locked_by=NULL,last_error=NULL,published_stream_id=$3 WHERE event_id=$1 AND locked_by=$2 AND processed_at IS NULL`, eventID, workerID, streamID)
if err != nil {
return fmt.Errorf("%w: mark processed: %v", ErrStoreUnavailable, err)
}
if result.RowsAffected() != 1 {
return ErrEventNotFound
}
return nil
}
func (s *Store) MarkFailed(ctx context.Context, event Event, workerID string, deliveryErr error, maxAttempts int, delay time.Duration) error {
if s == nil || s.pool == nil {
return ErrStoreUnavailable
}
message := deliveryErr.Error()
if len(message) > 2048 {
message = message[:2048]
}
dead := event.Attempts >= maxAttempts
result, err := s.pool.Exec(ctx, `
UPDATE gateway.outbox_events SET locked_at=NULL,locked_by=NULL,last_error=$3,
available_at=CASE WHEN $4 THEN available_at ELSE clock_timestamp()+($5 * interval '1 second') END,
dead_lettered_at=CASE WHEN $4 THEN clock_timestamp() ELSE NULL END
WHERE event_id=$1 AND locked_by=$2 AND processed_at IS NULL`, event.EventID, workerID, message, dead, delay.Seconds())
if err != nil {
return fmt.Errorf("%w: mark failed: %v", ErrStoreUnavailable, err)
}
if result.RowsAffected() != 1 {
return ErrEventNotFound
}
return nil
}
type EventView struct {
Event
AvailableAt time.Time `json:"available_at"`
LockedAt *time.Time `json:"locked_at"`
LockedBy *string `json:"locked_by"`
ProcessedAt *time.Time `json:"processed_at"`
DeadLetteredAt *time.Time `json:"dead_lettered_at"`
LastError *string `json:"last_error"`
PublishedStream *string `json:"published_stream_id"`
}
func (s *Store) List(ctx context.Context, status, eventType string, limit int) ([]EventView, error) {
if s == nil || s.pool == nil {
return nil, ErrStoreUnavailable
}
where := "TRUE"
switch status {
case "pending":
where = "processed_at IS NULL AND dead_lettered_at IS NULL"
case "dead":
where = "dead_lettered_at IS NOT NULL"
case "processed":
where = "processed_at IS NOT NULL"
}
args := []any{limit}
if eventType != "" {
args = append(args, eventType)
where += fmt.Sprintf(" AND event_type=$%d", len(args))
}
rows, err := s.pool.Query(ctx, `SELECT event_id::text,event_type,event_version,tenant_id::text,aggregate_type,aggregate_id,payload,trace_context,occurred_at,attempts,available_at,locked_at,locked_by,processed_at,dead_lettered_at,last_error,published_stream_id FROM gateway.outbox_events WHERE `+where+` ORDER BY occurred_at DESC LIMIT $1`, args...)
if err != nil {
return nil, fmt.Errorf("%w: list: %v", ErrStoreUnavailable, err)
}
defer rows.Close()
items := make([]EventView, 0)
for rows.Next() {
var item EventView
if err := rows.Scan(&item.EventID, &item.EventType, &item.EventVersion, &item.TenantID, &item.AggregateType,
&item.AggregateID, &item.Payload, &item.TraceContext, &item.OccurredAt, &item.Attempts, &item.AvailableAt,
&item.LockedAt, &item.LockedBy, &item.ProcessedAt, &item.DeadLetteredAt, &item.LastError, &item.PublishedStream); err != nil {
return nil, fmt.Errorf("%w: scan list: %v", ErrStoreUnavailable, err)
}
items = append(items, item)
}
return items, rows.Err()
}
func (s *Store) Retry(ctx context.Context, eventID string) error {
if s == nil || s.pool == nil {
return ErrStoreUnavailable
}
result, err := s.pool.Exec(ctx, `UPDATE gateway.outbox_events SET attempts=0,available_at=clock_timestamp(),locked_at=NULL,locked_by=NULL,dead_lettered_at=NULL,last_error=NULL WHERE event_id=$1 AND processed_at IS NULL`, eventID)
if err != nil {
return fmt.Errorf("%w: retry: %v", ErrStoreUnavailable, err)
}
if result.RowsAffected() != 1 {
return ErrEventNotFound
}
return nil
}
type ConsumerHandler func(context.Context, pgx.Tx) error
func (s *Store) Consume(ctx context.Context, subscriber, eventID string, handler ConsumerHandler) (bool, error) {
if s == nil || s.pool == nil {
return false, ErrStoreUnavailable
}
subscriber = strings.TrimSpace(subscriber)
if subscriber == "" || eventID == "" || handler == nil {
return false, errors.New("subscriber, event ID, and handler are required")
}
tx, err := s.pool.Begin(ctx)
if err != nil {
return false, fmt.Errorf("%w: begin consumption: %v", ErrStoreUnavailable, err)
}
defer func() { _ = tx.Rollback(ctx) }()
var inserted bool
err = tx.QueryRow(ctx, `WITH inserted AS (INSERT INTO gateway.event_consumptions(subscriber,event_id) VALUES($1,$2) ON CONFLICT DO NOTHING RETURNING 1) SELECT EXISTS(SELECT 1 FROM inserted)`, subscriber, eventID).Scan(&inserted)
if err != nil {
return false, fmt.Errorf("%w: reserve consumption: %v", ErrStoreUnavailable, err)
}
if !inserted {
return false, nil
}
if err := handler(ctx, tx); err != nil {
return false, err
}
if err := tx.Commit(ctx); err != nil {
return false, fmt.Errorf("%w: commit consumption: %v", ErrStoreUnavailable, err)
}
return true, nil
}
+91
View File
@@ -0,0 +1,91 @@
package outbox
import (
"context"
"errors"
"log/slog"
"math"
"time"
)
type WorkerConfig struct {
WorkerID string
BatchSize int
PollInterval time.Duration
Lease time.Duration
MaxAttempts int
MaxBackoff time.Duration
}
type Worker struct {
store *Store
publisher *RedisPublisher
config WorkerConfig
logger *slog.Logger
}
func NewWorker(store *Store, publisher *RedisPublisher, config WorkerConfig, logger *slog.Logger) *Worker {
return &Worker{store: store, publisher: publisher, config: config, logger: logger}
}
func (w *Worker) Run(ctx context.Context) error {
for {
processed, err := w.runBatch(ctx)
if err != nil && ctx.Err() == nil && w.logger != nil {
w.logger.Error("outbox batch failed", "error", err)
}
if ctx.Err() != nil {
return nil
}
if err == nil && processed > 0 {
continue
}
timer := time.NewTimer(w.config.PollInterval)
select {
case <-ctx.Done():
timer.Stop()
return nil
case <-timer.C:
}
}
}
func (w *Worker) runBatch(ctx context.Context) (int, error) {
// Gate on the downstream publisher before claiming anything. Claiming
// increments each event's attempt counter, so claiming during a Redis
// outage would burn every queued event's delivery budget and dead-letter
// the whole queue the moment the budget ran out — even though the events
// themselves were never at fault. When Redis is unreachable we back off
// instead, leaving events untouched in PostgreSQL until it recovers.
if err := w.publisher.Ping(ctx); err != nil {
return 0, err
}
events, err := w.store.Claim(ctx, w.config.WorkerID, w.config.BatchSize, w.config.Lease)
if err != nil {
return 0, err
}
var batchErr error
for _, event := range events {
result, publishErr := w.publisher.Publish(ctx, event)
if publishErr == nil {
if err := w.store.MarkProcessed(ctx, event.EventID, w.config.WorkerID, result.StreamID); err != nil {
batchErr = errors.Join(batchErr, err)
}
continue
}
delay := retryDelay(event.Attempts, w.config.MaxBackoff)
if err := w.store.MarkFailed(ctx, event, w.config.WorkerID, publishErr, w.config.MaxAttempts, delay); err != nil {
batchErr = errors.Join(batchErr, err)
}
}
return len(events), batchErr
}
func retryDelay(attempt int, maximum time.Duration) time.Duration {
seconds := math.Pow(2, float64(max(attempt-1, 0)))
delay := time.Duration(seconds * float64(time.Second))
if delay > maximum {
return maximum
}
return delay
}
+18
View File
@@ -0,0 +1,18 @@
package outbox
import (
"testing"
"time"
)
func TestRetryDelayIsExponentialAndCapped(t *testing.T) {
tests := []struct {
attempt int
want time.Duration
}{{1, time.Second}, {2, 2 * time.Second}, {3, 4 * time.Second}, {20, 5 * time.Second}}
for _, test := range tests {
if got := retryDelay(test.attempt, 5*time.Second); got != test.want {
t.Fatalf("attempt %d: got %s, want %s", test.attempt, got, test.want)
}
}
}
+26
View File
@@ -0,0 +1,26 @@
package apiresponse
import (
"encoding/json"
"net/http"
)
type Envelope struct {
Code int `json:"code"`
Message string `json:"msg"`
Data any `json:"data,omitempty"`
}
func OK(writer http.ResponseWriter, data any) {
Write(writer, http.StatusOK, Envelope{Code: http.StatusOK, Message: "success", Data: data})
}
func Error(writer http.ResponseWriter, status int, message string) {
Write(writer, status, Envelope{Code: status, Message: message})
}
func Write(writer http.ResponseWriter, status int, value any) {
writer.Header().Set("Content-Type", "application/json")
writer.WriteHeader(status)
_ = json.NewEncoder(writer).Encode(value)
}
+23
View File
@@ -0,0 +1,23 @@
package cache
import (
"fmt"
"time"
"github.com/redis/go-redis/v9"
)
func Open(rawURL string) (*redis.Client, error) {
if rawURL == "" {
return nil, nil
}
options, err := redis.ParseURL(rawURL)
if err != nil {
return nil, fmt.Errorf("parse redis URL: %w", err)
}
options.DialTimeout = 2 * time.Second
options.ReadTimeout = time.Second
options.WriteTimeout = time.Second
options.PoolTimeout = 2 * time.Second
return redis.NewClient(options), nil
}
+348
View File
@@ -0,0 +1,348 @@
package config
import (
"errors"
"fmt"
"net/url"
"os"
"strconv"
"strings"
"time"
)
type Config struct {
Environment string
Server Server
Database Database
Redis Redis
Security Security
Auth Auth
Credentials Credentials
Upstream Upstream
Audit Audit
Outbox Outbox
RuntimeData RuntimeData
Shadow Shadow
}
type Server struct {
Address string
ReadHeaderTimeout time.Duration
IdleTimeout time.Duration
ShutdownTimeout time.Duration
MaxBodyBytes int64
}
type Database struct {
URL string
MaxConns int32
MinConns int32
}
type Redis struct {
CriticalURL string
CacheURL string
}
type Security struct {
BootstrapAPIKey string
BootstrapAPIKeyEnabled bool
}
type Auth struct {
SessionTTL time.Duration
TOTPChallengeTTL time.Duration
MaxFailures int
LockDuration time.Duration
LoginRateLimitMax int // 单 IP 滑动窗口内的最大登录尝试次数
LoginRateLimitWindow time.Duration // 登录限流滑动窗口
}
type Credentials struct {
MasterKey string
KEKVersion int
KEKKeyring string
AllowPrivateProviderURL bool
AllowPrivateToolURL bool
AllowPrivateWebhookURL bool
ProviderRefreshInterval time.Duration
}
type Upstream struct {
BaseURL string
APIKey string
FallbackEnabled bool
ResponseHeaderTimeout time.Duration
MaxRetries int
RetryBackoff time.Duration
CircuitThreshold int
CircuitOpenDuration time.Duration
}
type Audit struct {
QueueSize int
BatchSize int
FlushInterval time.Duration
Retention time.Duration
UsageRetention time.Duration
PartitionMonthsAhead int
MaintenanceInterval time.Duration
}
type Outbox struct {
Stream string
BatchSize int
PollInterval time.Duration
Lease time.Duration
MaxAttempts int
MaxBackoff time.Duration
StreamMaxLength int64
MarkerTTL time.Duration
}
type RuntimeData struct {
ContentPolicyRefreshInterval time.Duration
PricingRefreshInterval time.Duration
}
type Shadow struct {
BaseURL string
APIKey string
SampleRate float64
Timeout time.Duration
MaxBodyBytes int64
MaxConcurrent int
}
func Load() (Config, error) {
cfg := Config{
Environment: env("APP_ENV", "local"),
Server: Server{
Address: env("HTTP_ADDR", "127.0.0.1:8080"),
ReadHeaderTimeout: duration("HTTP_READ_HEADER_TIMEOUT", 5*time.Second),
IdleTimeout: duration("HTTP_IDLE_TIMEOUT", 120*time.Second),
ShutdownTimeout: duration("HTTP_SHUTDOWN_TIMEOUT", 20*time.Second),
MaxBodyBytes: int64Value("HTTP_MAX_BODY_BYTES", 32<<20),
},
Database: Database{
URL: strings.TrimSpace(os.Getenv("DATABASE_URL")),
MaxConns: int32(intValue("DATABASE_MAX_CONNS", 40)),
MinConns: int32(intValue("DATABASE_MIN_CONNS", 4)),
},
Redis: Redis{
CriticalURL: strings.TrimSpace(os.Getenv("REDIS_CRITICAL_URL")),
CacheURL: strings.TrimSpace(os.Getenv("REDIS_CACHE_URL")),
},
Security: Security{
BootstrapAPIKey: strings.TrimSpace(os.Getenv("GATEWAY_BOOTSTRAP_API_KEY")),
BootstrapAPIKeyEnabled: boolValue("GATEWAY_BOOTSTRAP_API_KEY_ENABLED", false),
},
Auth: Auth{
SessionTTL: duration("AUTH_SESSION_TTL", 12*time.Hour),
TOTPChallengeTTL: duration("AUTH_TOTP_CHALLENGE_TTL", 5*time.Minute),
MaxFailures: intValue("LOGIN_MAX_FAILURES", 5),
LockDuration: duration("LOGIN_LOCK_DURATION", 15*time.Minute),
LoginRateLimitMax: intValue("LOGIN_RATE_LIMIT_MAX", 30),
LoginRateLimitWindow: duration("LOGIN_RATE_LIMIT_WINDOW", 5*time.Minute),
},
Credentials: Credentials{
MasterKey: strings.TrimSpace(os.Getenv("CREDENTIAL_MASTER_KEY")),
KEKVersion: intValue("CREDENTIAL_KEK_VERSION", 1),
KEKKeyring: strings.TrimSpace(os.Getenv("CREDENTIAL_KEK_KEYRING")),
AllowPrivateProviderURL: boolValue("ALLOW_PRIVATE_PROVIDER_URLS", false),
AllowPrivateToolURL: boolValue("ALLOW_PRIVATE_TOOL_URLS", false),
AllowPrivateWebhookURL: boolValue("ALLOW_PRIVATE_WEBHOOK_URLS", false),
ProviderRefreshInterval: duration("PROVIDER_REFRESH_INTERVAL", 5*time.Second),
},
Upstream: Upstream{
BaseURL: strings.TrimRight(env("UPSTREAM_BASE_URL", "https://api.openai.com"), "/"), APIKey: strings.TrimSpace(os.Getenv("UPSTREAM_API_KEY")),
FallbackEnabled: boolValue("UPSTREAM_FALLBACK_ENABLED", true), ResponseHeaderTimeout: duration("UPSTREAM_RESPONSE_HEADER_TIMEOUT", 60*time.Second),
MaxRetries: intValue("UPSTREAM_MAX_RETRIES", 2), RetryBackoff: duration("UPSTREAM_RETRY_BACKOFF", 50*time.Millisecond),
CircuitThreshold: intValue("UPSTREAM_CIRCUIT_THRESHOLD", 5), CircuitOpenDuration: duration("UPSTREAM_CIRCUIT_OPEN_DURATION", 30*time.Second),
},
Audit: Audit{
QueueSize: intValue("AUDIT_QUEUE_SIZE", 4096), BatchSize: intValue("AUDIT_BATCH_SIZE", 200),
FlushInterval: duration("AUDIT_FLUSH_INTERVAL", time.Second), Retention: duration("AUDIT_RETENTION", 90*24*time.Hour),
UsageRetention: duration("USAGE_RETENTION", 730*24*time.Hour), PartitionMonthsAhead: intValue("AUDIT_PARTITION_MONTHS_AHEAD", 3),
MaintenanceInterval: duration("AUDIT_MAINTENANCE_INTERVAL", 6*time.Hour),
},
Outbox: Outbox{
Stream: env("OUTBOX_STREAM", "gateway:{outbox}:events"), BatchSize: intValue("OUTBOX_BATCH_SIZE", 100),
PollInterval: duration("OUTBOX_POLL_INTERVAL", 500*time.Millisecond), Lease: duration("OUTBOX_LEASE", 30*time.Second),
MaxAttempts: intValue("OUTBOX_MAX_ATTEMPTS", 10), MaxBackoff: duration("OUTBOX_MAX_BACKOFF", 5*time.Minute),
StreamMaxLength: int64Value("OUTBOX_STREAM_MAX_LENGTH", 100_000), MarkerTTL: duration("OUTBOX_MARKER_TTL", 30*24*time.Hour),
},
RuntimeData: RuntimeData{
ContentPolicyRefreshInterval: duration("CONTENT_POLICY_REFRESH_INTERVAL", 30*time.Second),
PricingRefreshInterval: duration("PRICING_REFRESH_INTERVAL", 30*time.Second),
},
Shadow: Shadow{
BaseURL: strings.TrimRight(strings.TrimSpace(os.Getenv("SHADOW_BASE_URL")), "/"),
APIKey: strings.TrimSpace(os.Getenv("SHADOW_API_KEY")), SampleRate: floatValue("SHADOW_SAMPLE_RATE", 0),
Timeout: duration("SHADOW_TIMEOUT", 20*time.Second), MaxBodyBytes: int64Value("SHADOW_MAX_BODY_BYTES", 2<<20),
MaxConcurrent: intValue("SHADOW_MAX_CONCURRENT", 16),
},
}
return cfg, cfg.Validate()
}
func (c Config) Validate() error {
var errs []error
if c.Server.MaxBodyBytes <= 0 {
errs = append(errs, errors.New("HTTP_MAX_BODY_BYTES must be positive"))
}
if c.Database.MinConns < 0 || c.Database.MaxConns < 1 || c.Database.MinConns > c.Database.MaxConns {
errs = append(errs, errors.New("database pool sizes are invalid"))
}
if c.Auth.SessionTTL < 5*time.Minute || c.Auth.TOTPChallengeTTL < time.Minute || c.Auth.TOTPChallengeTTL > 15*time.Minute || c.Auth.MaxFailures < 1 || c.Auth.LockDuration < time.Minute {
errs = append(errs, errors.New("authentication limits are invalid"))
}
if c.Auth.LoginRateLimitMax < 1 || c.Auth.LoginRateLimitWindow < time.Second {
errs = append(errs, errors.New("login rate limit is invalid"))
}
if c.Credentials.KEKVersion < 1 {
errs = append(errs, errors.New("CREDENTIAL_KEK_VERSION must be positive"))
}
if c.Credentials.ProviderRefreshInterval < time.Second || c.Credentials.ProviderRefreshInterval > time.Minute {
errs = append(errs, errors.New("PROVIDER_REFRESH_INTERVAL must be between 1s and 1m"))
}
if err := validateHTTPURL(c.Upstream.BaseURL); err != nil {
errs = append(errs, fmt.Errorf("UPSTREAM_BASE_URL: %w", err))
}
if c.Upstream.ResponseHeaderTimeout < time.Second || c.Upstream.ResponseHeaderTimeout > 10*time.Minute || c.Upstream.MaxRetries < 0 || c.Upstream.MaxRetries > 5 || c.Upstream.RetryBackoff < 0 || c.Upstream.RetryBackoff > 5*time.Second || c.Upstream.CircuitThreshold < 1 || c.Upstream.CircuitThreshold > 100 || c.Upstream.CircuitOpenDuration < time.Second || c.Upstream.CircuitOpenDuration > 10*time.Minute {
errs = append(errs, errors.New("upstream resilience settings are invalid"))
}
if c.Audit.QueueSize < 100 || c.Audit.QueueSize > 1_000_000 || c.Audit.BatchSize < 1 || c.Audit.BatchSize > c.Audit.QueueSize || c.Audit.FlushInterval < 100*time.Millisecond || c.Audit.FlushInterval > time.Minute {
errs = append(errs, errors.New("audit buffering settings are invalid"))
}
if c.Audit.Retention < 24*time.Hour || c.Audit.Retention > 10*365*24*time.Hour || c.Audit.UsageRetention < c.Audit.Retention || c.Audit.UsageRetention > 10*365*24*time.Hour || c.Audit.PartitionMonthsAhead < 1 || c.Audit.PartitionMonthsAhead > 24 || c.Audit.MaintenanceInterval < time.Hour || c.Audit.MaintenanceInterval > 7*24*time.Hour {
errs = append(errs, errors.New("audit retention settings are invalid"))
}
if !strings.Contains(c.Outbox.Stream, "{outbox}") || c.Outbox.BatchSize < 1 || c.Outbox.BatchSize > 1000 || c.Outbox.PollInterval < 50*time.Millisecond || c.Outbox.PollInterval > time.Minute || c.Outbox.Lease < 5*time.Second || c.Outbox.Lease > 10*time.Minute || c.Outbox.MaxAttempts < 1 || c.Outbox.MaxAttempts > 100 || c.Outbox.MaxBackoff < time.Second || c.Outbox.MaxBackoff > time.Hour || c.Outbox.StreamMaxLength < 1000 || c.Outbox.StreamMaxLength > 100_000_000 || c.Outbox.MarkerTTL < 24*time.Hour || c.Outbox.MarkerTTL > 365*24*time.Hour {
errs = append(errs, errors.New("outbox delivery settings are invalid"))
}
if c.RuntimeData.ContentPolicyRefreshInterval < time.Second || c.RuntimeData.ContentPolicyRefreshInterval > 10*time.Minute || c.RuntimeData.PricingRefreshInterval < time.Second || c.RuntimeData.PricingRefreshInterval > 10*time.Minute {
errs = append(errs, errors.New("runtime data refresh settings must be between 1s and 10m"))
}
if c.Shadow.SampleRate < 0 || c.Shadow.SampleRate > 1 || c.Shadow.Timeout < time.Second || c.Shadow.Timeout > 2*time.Minute || c.Shadow.MaxBodyBytes < 1024 || c.Shadow.MaxBodyBytes > 32<<20 || c.Shadow.MaxConcurrent < 1 || c.Shadow.MaxConcurrent > 1000 {
errs = append(errs, errors.New("shadow traffic settings are invalid"))
}
if c.Shadow.BaseURL != "" {
if err := validateHTTPURL(c.Shadow.BaseURL); err != nil {
errs = append(errs, fmt.Errorf("SHADOW_BASE_URL: %w", err))
}
if c.Shadow.APIKey == "" {
errs = append(errs, errors.New("SHADOW_API_KEY is required when SHADOW_BASE_URL is set"))
}
}
return errors.Join(errs...)
}
func (c Config) ValidateRuntime() error {
var errs []error
if strings.EqualFold(c.Environment, "production") {
if c.Database.URL == "" {
errs = append(errs, errors.New("DATABASE_URL is required in production"))
}
if c.Redis.CriticalURL == "" {
errs = append(errs, errors.New("REDIS_CRITICAL_URL is required in production"))
}
if c.Security.BootstrapAPIKeyEnabled {
if len(c.Security.BootstrapAPIKey) < 32 {
errs = append(errs, errors.New("GATEWAY_BOOTSTRAP_API_KEY must contain at least 32 characters in production"))
}
// Reject the documented placeholder and other obviously weak values
// so an operator cannot accidentally deploy with the example secret.
switch strings.ToLower(c.Security.BootstrapAPIKey) {
case "change-me-in-production", "local-development-key", "changeme", "change-me", "password", "secret":
errs = append(errs, errors.New("GATEWAY_BOOTSTRAP_API_KEY is set to a known weak default; generate a strong random value"))
}
}
if c.Upstream.FallbackEnabled && c.Upstream.APIKey == "" {
errs = append(errs, errors.New("UPSTREAM_API_KEY is required in production"))
}
if c.Credentials.MasterKey == "" {
errs = append(errs, errors.New("CREDENTIAL_MASTER_KEY is required in production"))
}
}
return errors.Join(errs...)
}
func validateHTTPURL(raw string) error {
u, err := url.Parse(raw)
if err != nil {
return err
}
if (u.Scheme != "http" && u.Scheme != "https") || u.Host == "" || u.User != nil {
return errors.New("must be an absolute http(s) URL without user info")
}
return nil
}
func env(key, fallback string) string {
if value := strings.TrimSpace(os.Getenv(key)); value != "" {
return value
}
return fallback
}
func duration(key string, fallback time.Duration) time.Duration {
value := strings.TrimSpace(os.Getenv(key))
if value == "" {
return fallback
}
parsed, err := time.ParseDuration(value)
if err != nil {
return fallback
}
return parsed
}
func intValue(key string, fallback int) int {
value := strings.TrimSpace(os.Getenv(key))
if value == "" {
return fallback
}
parsed, err := strconv.Atoi(value)
if err != nil {
return fallback
}
return parsed
}
func int64Value(key string, fallback int64) int64 {
value := strings.TrimSpace(os.Getenv(key))
if value == "" {
return fallback
}
parsed, err := strconv.ParseInt(value, 10, 64)
if err != nil {
return fallback
}
return parsed
}
func floatValue(key string, fallback float64) float64 {
value := strings.TrimSpace(os.Getenv(key))
if value == "" {
return fallback
}
parsed, err := strconv.ParseFloat(value, 64)
if err != nil {
return fallback
}
return parsed
}
func boolValue(key string, fallback bool) bool {
value := strings.TrimSpace(os.Getenv(key))
if value == "" {
return fallback
}
parsed, err := strconv.ParseBool(value)
if err != nil {
return fallback
}
return parsed
}
+45
View File
@@ -0,0 +1,45 @@
package config
import "testing"
func TestProductionRequiresCoreDependencies(t *testing.T) {
t.Setenv("APP_ENV", "production")
t.Setenv("DATABASE_URL", "")
t.Setenv("REDIS_CRITICAL_URL", "")
t.Setenv("GATEWAY_BOOTSTRAP_API_KEY", "short")
t.Setenv("UPSTREAM_API_KEY", "")
cfg, err := Load()
if err != nil {
t.Fatalf("unexpected structural configuration error: %v", err)
}
if err := cfg.ValidateRuntime(); err == nil {
t.Fatal("expected production validation error")
}
}
func TestLocalDefaultsAreValid(t *testing.T) {
t.Setenv("APP_ENV", "local")
t.Setenv("UPSTREAM_BASE_URL", "https://example.com")
if _, err := Load(); err != nil {
t.Fatalf("unexpected error: %v", err)
}
}
func TestProductionCanDisableBootstrapCompatibility(t *testing.T) {
t.Setenv("APP_ENV", "production")
t.Setenv("DATABASE_URL", "postgres://gateway:secret@db.example/gateway?sslmode=require")
t.Setenv("REDIS_CRITICAL_URL", "rediss://redis.example/0")
t.Setenv("GATEWAY_BOOTSTRAP_API_KEY", "")
t.Setenv("GATEWAY_BOOTSTRAP_API_KEY_ENABLED", "false")
t.Setenv("UPSTREAM_FALLBACK_ENABLED", "false")
t.Setenv("CREDENTIAL_MASTER_KEY", "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=")
t.Setenv("UPSTREAM_BASE_URL", "https://example.com")
cfg, err := Load()
if err != nil {
t.Fatal(err)
}
if err := cfg.ValidateRuntime(); err != nil {
t.Fatalf("disabled bootstrap compatibility should not require a key: %v", err)
}
}
+77
View File
@@ -0,0 +1,77 @@
package cryptox
import (
"crypto/aes"
"crypto/cipher"
"crypto/rand"
"encoding/base64"
"errors"
"fmt"
)
var ErrKeyUnavailable = errors.New("encryption key unavailable")
type AESGCM struct {
aead cipher.AEAD
version int
purpose string
}
func NewAESGCM(encodedKey string, version int, purpose string) (*AESGCM, error) {
if encodedKey == "" {
return nil, nil
}
key, err := base64.StdEncoding.DecodeString(encodedKey)
if err != nil {
key, err = base64.RawStdEncoding.DecodeString(encodedKey)
}
if err != nil || len(key) != 32 {
return nil, errors.New("encryption key must be base64-encoded 32 bytes")
}
if version < 1 || purpose == "" {
return nil, errors.New("encryption version and purpose are required")
}
block, err := aes.NewCipher(key)
if err != nil {
return nil, err
}
aead, err := cipher.NewGCM(block)
if err != nil {
return nil, err
}
return &AESGCM{aead: aead, version: version, purpose: purpose}, nil
}
func (c *AESGCM) Encrypt(plaintext []byte) ([]byte, int, error) {
if c == nil {
return nil, 0, ErrKeyUnavailable
}
nonce := make([]byte, c.aead.NonceSize())
if _, err := rand.Read(nonce); err != nil {
return nil, 0, err
}
ciphertext := c.aead.Seal(nil, nonce, plaintext, c.additionalData())
return append(nonce, ciphertext...), c.version, nil
}
func (c *AESGCM) Decrypt(encrypted []byte, version int) ([]byte, error) {
if c == nil {
return nil, ErrKeyUnavailable
}
if version != c.version {
return nil, fmt.Errorf("encryption key version %d is not loaded", version)
}
nonceSize := c.aead.NonceSize()
if len(encrypted) <= nonceSize {
return nil, errors.New("encrypted value is truncated")
}
plaintext, err := c.aead.Open(nil, encrypted[:nonceSize], encrypted[nonceSize:], c.additionalData())
if err != nil {
return nil, errors.New("encrypted value authentication failed")
}
return plaintext, nil
}
func (c *AESGCM) additionalData() []byte {
return []byte(fmt.Sprintf("ai-gateway/%s/v%d", c.purpose, c.version))
}
+100
View File
@@ -0,0 +1,100 @@
package cryptox
import (
"encoding/json"
"errors"
"fmt"
"sort"
"strconv"
"strings"
)
type Cipher interface {
Encrypt([]byte) ([]byte, int, error)
Decrypt([]byte, int) ([]byte, error)
}
type Keyring struct {
activeVersion int
ciphers map[int]*AESGCM
}
func NewKeyring(activeKey string, activeVersion int, encodedKeyring, purpose string) (*Keyring, error) {
if activeVersion < 1 || purpose == "" {
return nil, errors.New("encryption version and purpose are required")
}
encoded := make(map[int]string)
if strings.TrimSpace(encodedKeyring) != "" {
var values map[string]string
if err := json.Unmarshal([]byte(encodedKeyring), &values); err != nil {
return nil, fmt.Errorf("encryption keyring must be a JSON object: %w", err)
}
for rawVersion, key := range values {
version, err := strconv.Atoi(rawVersion)
if err != nil || version < 1 || strings.TrimSpace(key) == "" {
return nil, fmt.Errorf("invalid encryption keyring version %q", rawVersion)
}
encoded[version] = strings.TrimSpace(key)
}
}
if strings.TrimSpace(activeKey) != "" {
if existing, ok := encoded[activeVersion]; ok && existing != strings.TrimSpace(activeKey) {
return nil, fmt.Errorf("active encryption key version %d is configured twice with different values", activeVersion)
}
encoded[activeVersion] = strings.TrimSpace(activeKey)
}
if len(encoded) == 0 {
return nil, nil
}
if _, ok := encoded[activeVersion]; !ok {
return nil, fmt.Errorf("active encryption key version %d is not loaded", activeVersion)
}
keyring := &Keyring{activeVersion: activeVersion, ciphers: make(map[int]*AESGCM, len(encoded))}
for version, key := range encoded {
cipher, err := NewAESGCM(key, version, purpose)
if err != nil {
return nil, fmt.Errorf("encryption key version %d: %w", version, err)
}
keyring.ciphers[version] = cipher
}
return keyring, nil
}
func (k *Keyring) Encrypt(plaintext []byte) ([]byte, int, error) {
if k == nil {
return nil, 0, ErrKeyUnavailable
}
return k.ciphers[k.activeVersion].Encrypt(plaintext)
}
func (k *Keyring) Decrypt(encrypted []byte, version int) ([]byte, error) {
if k == nil {
return nil, ErrKeyUnavailable
}
cipher, ok := k.ciphers[version]
if !ok {
return nil, fmt.Errorf("encryption key version %d is not loaded", version)
}
return cipher.Decrypt(encrypted, version)
}
func (k *Keyring) ActiveVersion() int {
if k == nil {
return 0
}
return k.activeVersion
}
func (k *Keyring) Versions() []int {
if k == nil {
return nil
}
versions := make([]int, 0, len(k.ciphers))
for version := range k.ciphers {
versions = append(versions, version)
}
sort.Ints(versions)
return versions
}
var _ Cipher = (*Keyring)(nil)
+37
View File
@@ -0,0 +1,37 @@
package cryptox
import (
"encoding/base64"
"fmt"
"testing"
)
func TestKeyringDecryptsOldVersionAndEncryptsActiveVersion(t *testing.T) {
oldKey := base64.StdEncoding.EncodeToString(make([]byte, 32))
newBytes := make([]byte, 32)
newBytes[0] = 1
newKey := base64.StdEncoding.EncodeToString(newBytes)
oldCipher, err := NewAESGCM(oldKey, 1, "test")
if err != nil {
t.Fatal(err)
}
oldEncrypted, _, err := oldCipher.Encrypt([]byte("secret"))
if err != nil {
t.Fatal(err)
}
keyring, err := NewKeyring(newKey, 2, fmt.Sprintf(`{"1":%q}`, oldKey), "test")
if err != nil {
t.Fatal(err)
}
plaintext, err := keyring.Decrypt(oldEncrypted, 1)
if err != nil || string(plaintext) != "secret" {
t.Fatalf("old key was not usable: %q, %v", plaintext, err)
}
newEncrypted, version, err := keyring.Encrypt([]byte("new-secret"))
if err != nil || version != 2 {
t.Fatalf("active version was not used: version=%d error=%v", version, err)
}
if _, err := oldCipher.Decrypt(newEncrypted, version); err == nil {
t.Fatal("old cipher must not decrypt the new version")
}
}

Some files were not shown because too many files have changed in this diff Show More