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:
@@ -0,0 +1,9 @@
|
||||
.git
|
||||
.env
|
||||
deploy/production.env
|
||||
bin
|
||||
coverage.out
|
||||
web/node_modules
|
||||
web/**/node_modules
|
||||
web/**/dist
|
||||
web/.pnpm-store
|
||||
@@ -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=
|
||||
@@ -0,0 +1,9 @@
|
||||
.env
|
||||
deploy/production.env
|
||||
.idea/
|
||||
.vscode/
|
||||
.DS_Store
|
||||
bin/
|
||||
coverage.out
|
||||
web/node_modules/
|
||||
web/**/dist/
|
||||
@@ -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
|
||||
@@ -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 API,Go 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
@@ -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
@@ -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")
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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) }
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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:
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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))
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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"]
|
||||
@@ -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
|
||||
@@ -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.
|
||||
@@ -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:
|
||||
@@ -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:
|
||||
@@ -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;
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
@@ -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` 接口接入对象存储或分析库。
|
||||
@@ -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 备份、指标和日志,再分析问题。
|
||||
@@ -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` | 否 | `—` | UNIQUE;FK → `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` | — |
|
||||
@@ -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 数据文件,因此已完成的是可验证暂存管道;生产领域转换必须在拿到脱敏快照及旧密钥托管授权后执行,不能凭空假设旧密文格式。
|
||||
@@ -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。
|
||||
|
||||
工程实现与本地生产等价验证已关闭。正式上线仍是外部发布门禁:拿到实际旧数据库脱敏快照和旧密钥迁移授权后执行领域转换,并在真实新旧双系统与生产等价流量下完成影子观察、容量验收及切换/回退演练;这些动作不会在缺少生产数据和授权时伪造为已执行。
|
||||
@@ -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`),靠目录快照;建议尽快接入版本控制 |
|
||||
|
||||
---
|
||||
|
||||
## 二、已完成里程碑(M0–M7)
|
||||
|
||||
### 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`。
|
||||
- **上线门禁未过**:旧库脱敏快照迁移、影子观察、容量验收、切换/回退演练需真实生产数据与授权后方可执行。
|
||||
@@ -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
|
||||
)
|
||||
@@ -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=
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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 }
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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(¤tTable); 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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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 }
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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 输出 uncertain,claims 输出 []。`
|
||||
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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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}
|
||||
`
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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 }
|
||||
}
|
||||
@@ -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) }
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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": "模型权限"}},
|
||||
}},
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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]
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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}
|
||||
}
|
||||
@@ -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[:])
|
||||
}
|
||||
@@ -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[:])
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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})
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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, ''}
|
||||
`
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
Vendored
+23
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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))
|
||||
}
|
||||
@@ -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)
|
||||
@@ -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
Reference in New Issue
Block a user