From 4c632d94f03077b382f2eeb127874f6caec12e53 Mon Sep 17 00:00:00 2001 From: Futureppo Date: Mon, 13 Jul 2026 01:17:47 +0800 Subject: [PATCH 1/2] Add one-command Docker deployment --- README.md | 26 +++++- README_EN.md | 26 +++++- scripts/deploy.sh | 199 ++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 249 insertions(+), 2 deletions(-) create mode 100755 scripts/deploy.sh diff --git a/README.md b/README.md index 387f35d..d3569ea 100644 --- a/README.md +++ b/README.md @@ -11,7 +11,7 @@ [![License: AGPL v3](https://img.shields.io/badge/License-AGPL_v3-663399.svg)](LICENSE) [![Container](https://img.shields.io/badge/GHCR-grokcli2api--go-2496ED?logo=docker&logoColor=white)](https://github.com/Futureppo/grokcli2api-go/pkgs/container/grokcli2api-go) -[快速开始](#快速开始) · [API 兼容性](#api-兼容性) · [配置说明](#配置说明) · [接口一览](#接口一览) · [参与贡献](#开发与贡献) +[一键部署](#一键部署linux) · [快速开始](#快速开始) · [API 兼容性](#api-兼容性) · [配置说明](#配置说明) · [接口一览](#接口一览) · [参与贡献](#开发与贡献) **简体中文** · [English](README_EN.md) @@ -70,6 +70,30 @@ flowchart LR ## 快速开始 +### 一键部署（Linux） + +服务器已安装 Docker 与 Docker Compose v2 时，可直接运行： + +```bash +bash <(curl -fsSL https://raw.githubusercontent.com/Futureppo/grokcli2api-go/main/scripts/deploy.sh) +``` + +脚本会检查 Docker 环境、下载 Compose 配置、创建受保护的 `.env` 和 `auths/` 目录、生成随机本地 API Key、导入你指定的 OAuth JSON 凭证，并启动及验证服务。交互执行时按提示输入凭证文件路径即可；已有安装会保留 `.env` 与凭证，可用同一条命令完成镜像更新。 + +无人值守部署可预先传入参数： + +```bash +AUTH_FILE=/root/account.json \ +GROK_API_KEYS='sk-change-this-to-a-strong-random-key' \ +INSTALL_DIR=/opt/grokcli2api-go \ +bash <(curl -fsSL https://raw.githubusercontent.com/Futureppo/grokcli2api-go/main/scripts/deploy.sh) +``` + +可选变量包括 `GROK2API_PORT`（默认 `8088`）、`INSTALL_DIR`（默认 `~/grokcli2api-go`）、`AUTH_FILE` 与 `GROK_API_KEYS`。未提供凭证时，脚本只初始化安全配置而不会启动一个无法工作的服务。 + +> [!TIP] +> 一键脚本解决的是服务部署，不会替你获取上游凭证。OAuth JSON 属于敏感信息，请只从可信来源导出，并在服务器上以最小权限保存。正式对公网开放前还应配置 HTTPS、反向代理、访问控制和限流。 + ### 1. 准备项目 运行前需要： diff --git a/README_EN.md b/README_EN.md index 4a11dd9..0f07035 100644 --- a/README_EN.md +++ b/README_EN.md @@ -11,7 +11,7 @@ A lightweight, deployable Go compatibility layer with streaming and multi-accoun [![License: AGPL v3](https://img.shields.io/badge/License-AGPL_v3-663399.svg)](LICENSE) [![Container](https://img.shields.io/badge/GHCR-grokcli2api--go-2496ED?logo=docker&logoColor=white)](https://github.com/Futureppo/grokcli2api-go/pkgs/container/grokcli2api-go) -[Quick Start](#quick-start) · [API Compatibility](#api-compatibility) · [Configuration](#configuration) · [Endpoints](#endpoints) · [Contributing](#development-and-contributing) +[One-command deployment](#one-command-deployment-linux) · [Quick Start](#quick-start) · [API Compatibility](#api-compatibility) · [Configuration](#configuration) · [Endpoints](#endpoints) · [Contributing](#development-and-contributing) [简体中文](README.md) · **English** @@ -70,6 +70,30 @@ The compatibility layer preserves commonly used request fields, response structu ## Quick Start +### One-command deployment (Linux) + +If Docker and Docker Compose v2 are already installed on the server, run: + +```bash +bash <(curl -fsSL https://raw.githubusercontent.com/Futureppo/grokcli2api-go/main/scripts/deploy.sh) +``` + +The script checks Docker, downloads the Compose configuration, creates a protected `.env` file and `auths/` directory, generates a random local API key, imports the OAuth JSON credential you select, and starts and verifies the service. Interactive runs prompt for the credential path. Existing installations retain their `.env` and credentials, so the same command also updates the container image. + +For unattended deployment, provide the inputs as environment variables: + +```bash +AUTH_FILE=/root/account.json \ +GROK_API_KEYS='sk-change-this-to-a-strong-random-key' \ +INSTALL_DIR=/opt/grokcli2api-go \ +bash <(curl -fsSL https://raw.githubusercontent.com/Futureppo/grokcli2api-go/main/scripts/deploy.sh) +``` + +Optional variables include `GROK2API_PORT` (default `8088`), `INSTALL_DIR` (default `~/grokcli2api-go`), `AUTH_FILE`, and `GROK_API_KEYS`. Without a credential, the script initializes the protected configuration but does not start a service that cannot handle requests. + +> [!TIP] +> The deployment script does not obtain upstream credentials for you. OAuth JSON is sensitive: export it only from a trusted source and store it with least-privilege permissions. Configure HTTPS, a reverse proxy, access controls, and rate limiting before exposing the service publicly. + ### 1. Prepare the Project You will need: diff --git a/scripts/deploy.sh b/scripts/deploy.sh new file mode 100755 index 0000000..ee7d665 --- /dev/null +++ b/scripts/deploy.sh @@ -0,0 +1,199 @@ +#!/usr/bin/env bash + +set -Eeuo pipefail + +REPOSITORY="${REPOSITORY:-Futureppo/grokcli2api-go}" +DEPLOY_REF="${DEPLOY_REF:-main}" +INSTALL_DIR="${INSTALL_DIR:-${HOME}/grokcli2api-go}" +PORT="${GROK2API_PORT:-8088}" +RAW_BASE="https://raw.githubusercontent.com/${REPOSITORY}/${DEPLOY_REF}" + +info() { + printf '\033[1;34m[INFO]\033[0m %s\n' "$*" +} + +success() { + printf '\033[1;32m[OK]\033[0m %s\n' "$*" +} + +warn() { + printf '\033[1;33m[WARN]\033[0m %s\n' "$*" >&2 +} + +fail() { + printf '\033[1;31m[ERROR]\033[0m %s\n' "$*" >&2 + exit 1 +} + +require_command() { + command -v "$1" >/dev/null 2>&1 || fail "缺少命令：$1" +} + +fetch() { + local url="$1" + local output="$2" + curl --fail --silent --show-error --location --retry 3 "$url" --output "$output" +} + +read_env_value() { + local key="$1" + local file="$2" + awk -v key="$key" ' + index($0, key "=") == 1 { + sub(/^[^=]*=/, "") + print + exit + } + ' "$file" +} + +upsert_env() { + local key="$1" + local value="$2" + local file="$3" + local temporary + temporary="$(mktemp "${file}.XXXXXX")" + awk -v key="$key" -v value="$value" ' + BEGIN { replaced = 0 } + index($0, key "=") == 1 { + if (!replaced) { + print key "=" value + replaced = 1 + } + next + } + { print } + END { + if (!replaced) { + print key "=" value + } + } + ' "$file" >"$temporary" + mv "$temporary" "$file" +} + +generate_api_key() { + if command -v openssl >/dev/null 2>&1; then + printf 'sk-%s' "$(openssl rand -hex 24)" + return + fi + printf 'sk-%s' "$(od -An -N24 -tx1 /dev/urandom | tr -d ' \n')" +} + +find_credential() { + find "$1" -maxdepth 1 -type f -name '*.json' -print -quit 2>/dev/null || true +} + +case "$PORT" in + ''|*[!0-9]*) fail "GROK2API_PORT 必须是 1 到 65535 之间的整数" ;; +esac +if ((PORT < 1 || PORT > 65535)); then + fail "GROK2API_PORT 必须是 1 到 65535 之间的整数" +fi + +require_command curl +require_command docker +require_command awk +require_command find + +docker compose version >/dev/null 2>&1 || fail "未检测到 Docker Compose v2，请先安装或升级 Docker" +docker info >/dev/null 2>&1 || fail "Docker 服务未运行，或当前用户没有访问 Docker 的权限" + +mkdir -p "$INSTALL_DIR/auths" +chmod 700 "$INSTALL_DIR/auths" + +compose_tmp="$(mktemp "${INSTALL_DIR}/compose.yaml.XXXXXX")" +env_tmp="" +trap 'rm -f "$compose_tmp" "$env_tmp"' EXIT + +info "下载 Docker Compose 配置" +fetch "${RAW_BASE}/compose.yaml" "$compose_tmp" +mv "$compose_tmp" "$INSTALL_DIR/compose.yaml" + +if [[ ! -f "$INSTALL_DIR/.env" ]]; then + env_tmp="$(mktemp "${INSTALL_DIR}/.env.XXXXXX")" + info "创建环境变量配置" + fetch "${RAW_BASE}/.env.example" "$env_tmp" + mv "$env_tmp" "$INSTALL_DIR/.env" +else + info "保留现有 .env 配置" +fi + +api_keys="${GROK_API_KEYS:-}" +if [[ -z "$api_keys" ]]; then + api_keys="$(read_env_value GROK_API_KEYS "$INSTALL_DIR/.env")" +fi +if [[ -z "$api_keys" || "$api_keys" == "key1,key2,key3" ]]; then + api_keys="$(generate_api_key)" + generated_api_key=1 +else + generated_api_key=0 +fi +if [[ "$api_keys" == *$'\n'* || "$api_keys" == *$'\r'* ]]; then + fail "GROK_API_KEYS 不能包含换行符" +fi + +upsert_env GROK2API_PORT "$PORT" "$INSTALL_DIR/.env" +upsert_env GROK_API_KEYS "$api_keys" "$INSTALL_DIR/.env" +chmod 600 "$INSTALL_DIR/.env" + +credential="$(find_credential "$INSTALL_DIR/auths")" +credential_source="${AUTH_FILE:-}" + +if [[ -z "$credential" && -z "$credential_source" && -t 1 && -r /dev/tty ]]; then + printf '请输入 Grok CLI OAuth JSON 凭证路径（直接回车仅初始化配置）：' >/dev/tty + IFS= read -r credential_source /dev/null 2>&1; then + ready=1 + break + fi + sleep 2 +done + +if ((ready == 0)); then + warn "服务未在预期时间内就绪，最近的容器日志如下：" + (cd "$INSTALL_DIR" && docker compose logs --tail 80) || true + fail "部署未通过健康检查" +fi + +success "grokcli2api-go 已启动：http://127.0.0.1:${PORT}" +printf '安装目录：%s\n' "$INSTALL_DIR" +if ((generated_api_key == 1)); then + printf '自动生成的本地 API Key：%s\n' "$api_keys" + warn "请立即保存此 Key；后续可在 ${INSTALL_DIR}/.env 中查看或更换。" +else + printf '本地 API Key 已按现有配置启用。\n' +fi +printf '\n常用命令：\n' +printf ' 查看状态：cd %q && docker compose ps\n' "$INSTALL_DIR" +printf ' 查看日志：cd %q && docker compose logs -f\n' "$INSTALL_DIR" +printf ' 更新服务：重新执行本部署脚本\n' From 5ec540e2619c5f00f96b017d27c65a4a27856c92 Mon Sep 17 00:00:00 2001 From: Futureppo Date: Mon, 13 Jul 2026 01:22:15 +0800 Subject: [PATCH 2/2] =?UTF-8?q?=E6=B7=BB=E5=8A=A0=E8=BF=9C=E7=A8=8B?= =?UTF-8?q?=E7=AE=A1=E7=90=86=E6=8E=A5=E5=8F=A3?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .env.example | 4 + README.md | 31 +++++ README_EN.md | 31 +++++ internal/auth/auth.go | 24 +++- internal/auth/auth_test.go | 156 +++++++++++++++++++++++++ internal/auth/pool.go | 203 ++++++++++++++++++++++++++++++++- internal/config/config.go | 2 + internal/config/config_test.go | 11 ++ internal/grok/client.go | 21 +++- internal/server/server.go | 173 ++++++++++++++++++++++++++-- internal/server/server_test.go | 182 ++++++++++++++++++++++++++++- 11 files changed, 823 insertions(+), 15 deletions(-) diff --git a/.env.example b/.env.example index fff46c5..02eee24 100644 --- a/.env.example +++ b/.env.example @@ -24,6 +24,10 @@ GROK_AFFINITY_MAX_ENTRIES=100000 GROK_API_KEYS=key1,key2,key3 +# 可选的管理员密钥。设置后启用远程凭证管理接口，并允许在凭证目录为空时启动服务。 +# 管理接口仅接受此密钥，不接受上面的普通 API Key。请只通过 HTTPS 暴露管理接口。 +GROK_ADMIN_KEY=admin + # 每个账号对应一个 OAuth JSON 凭证文件，并直接放入此目录，支持直接导入CPA格式的凭证 # 文件支持热加载，刷新后的令牌会写回。 diff --git a/README.md b/README.md index d3569ea..d1cc72a 100644 --- a/README.md +++ b/README.md @@ -265,6 +265,7 @@ X-Grok-Session-ID: conversation-123 | `GROK2API_LOG_LEVEL` | `INFO` | 日志等级：`DEBUG`、`INFO`、`WARN` 或 `ERROR` | | `GROK_API_KEYS` | 空 | 逗号分隔的本地访问密钥，可为不同客户端分配独立 Key | | `GROK_API_KEY` | 空 | 单个本地访问密钥的兼容别名 | +| `GROK_ADMIN_KEY` | 空 | 独立的管理员密钥；设置后启用远程凭证管理并允许空凭证池启动 | 启用本地访问保护后，受保护接口接受以下任一种请求头： @@ -272,6 +273,25 @@ X-Grok-Session-ID: conversation-123 - `x-api-key: ` - `api-key: ` +管理员密钥与普通 API Key 相互独立。管理接口接受 `Authorization: Bearer ` 或 `X-Admin-Key: `，且只应通过 HTTPS 和受限网络暴露。上传凭证可以直接发送 JSON： + +```bash +curl http://localhost:8088/v1/admin/credentials \ + -H "Authorization: Bearer $GROK_ADMIN_KEY" \ + -H "Content-Type: application/json" \ + --data-binary @auth.json +``` + +也可以使用文件表单： + +```bash +curl http://localhost:8088/v1/admin/credentials \ + -H "X-Admin-Key: $GROK_ADMIN_KEY" \ + -F "file=@auth.json;type=application/json" +``` + +服务端根据凭证中的稳定账号标识生成脱敏 ID；重复上传同一账号会原子覆盖原凭证。上传后会立即尝试发现模型，临时探测失败不会删除已经保存的凭证。 + ### 凭证池与调度 | 环境变量 | 未设置时的默认值 | 说明 | @@ -326,6 +346,16 @@ go run ./cmd/grok2api -version “可选”表示仅在配置了本地 API Key 时需要鉴权。 +### 管理员凭证接口 + +仅在设置 `GROK_ADMIN_KEY` 后启用，普通 API Key 无法访问。列表响应不会包含账号标识、文件路径、客户端 ID 或 Token。 + +| 方法 | 路径 | 说明 | +| --- | --- | --- | +| `GET` | `/v1/admin/credentials` | 列出脱敏的凭证状态和模型目录 | +| `POST` | `/v1/admin/credentials` | 上传或覆盖 JSON 凭证，支持 JSON 请求体和 multipart `file` 字段 | +| `DELETE` | `/v1/admin/credentials/{id}` | 删除凭证并立即从调度池移除 | + ### Grok 只读透传接口 | 方法 | 路径 | @@ -351,6 +381,7 @@ go run ./cmd/grok2api -version - 切勿提交或公开 OAuth Token、API Key、认证文件及未脱敏日志。 - 对外提供服务前，务必配置 `GROK_API_KEYS`，并在反向代理层启用 HTTPS、访问控制和限流。 +- 启用 `GROK_ADMIN_KEY` 时应使用独立的强随机密钥，并额外限制管理接口的来源地址和请求频率。 - 为凭证目录设置最小必要文件权限，并限制可访问该目录的系统用户。 - 除非处于受控调试环境，否则不要启用 `GROK_TLS_INSECURE_SKIP_VERIFY`。 - 不要把会话 ID、用户邮箱或其他敏感数据直接用作亲和标识。 diff --git a/README_EN.md b/README_EN.md index 0f07035..08e1c5e 100644 --- a/README_EN.md +++ b/README_EN.md @@ -265,6 +265,7 @@ The service loads environment variables that are not already set from a `.env` f | `GROK2API_LOG_LEVEL` | `INFO` | `DEBUG`, `INFO`, `WARN`, or `ERROR` | | `GROK_API_KEYS` | empty | Comma-separated local access keys; separate keys may be assigned to different clients | | `GROK_API_KEY` | empty | Backward-compatible alias for one local access key | +| `GROK_ADMIN_KEY` | empty | Independent administrator key; enables remote credential management and empty-pool startup | When local access protection is enabled, protected endpoints accept any of these headers: @@ -272,6 +273,25 @@ When local access protection is enabled, protected endpoints accept any of these - `x-api-key: ` - `api-key: ` +The administrator key is independent from normal API keys. Management endpoints accept `Authorization: Bearer ` or `X-Admin-Key: ` and should only be exposed through HTTPS on a restricted network. Upload a credential as a JSON request body: + +```bash +curl http://localhost:8088/v1/admin/credentials \ + -H "Authorization: Bearer $GROK_ADMIN_KEY" \ + -H "Content-Type: application/json" \ + --data-binary @auth.json +``` + +Or upload it as a file form: + +```bash +curl http://localhost:8088/v1/admin/credentials \ + -H "X-Admin-Key: $GROK_ADMIN_KEY" \ + -F "file=@auth.json;type=application/json" +``` + +The server derives a redacted ID from the credential's stable account identity. Uploading the same account again atomically replaces its existing credential. Model discovery runs immediately after upload; a temporary discovery failure does not remove the saved credential. + ### Credential Pool and Scheduling | Environment variable | Default when unset | Description | @@ -326,6 +346,16 @@ go run ./cmd/grok2api -version “Optional” means authentication is required only when a local API key has been configured. +### Administrator Credential APIs + +These endpoints are registered only when `GROK_ADMIN_KEY` is set, and normal API keys cannot access them. List responses never expose account subjects, file paths, client IDs, or tokens. + +| Method | Path | Description | +| --- | --- | --- | +| `GET` | `/v1/admin/credentials` | List redacted credential status and model catalogs | +| `POST` | `/v1/admin/credentials` | Upload or replace a JSON credential using a JSON body or multipart `file` field | +| `DELETE` | `/v1/admin/credentials/{id}` | Delete a credential and immediately remove it from scheduling | + ### Read-only Grok Passthrough APIs | Method | Path | @@ -351,6 +381,7 @@ Local `.env` and `auths/` data remain external configuration and credential stor - Never commit or disclose OAuth tokens, API keys, authentication files, or unsanitized logs. - Configure `GROK_API_KEYS` before exposing the service to a network, and enable HTTPS, access control, and rate limiting at the reverse proxy. +- When enabling `GROK_ADMIN_KEY`, use a separate strong random key and apply additional source-address and rate restrictions to the management endpoints. - Apply least-privilege file permissions to the credential directory and restrict which system users can access it. - Do not enable `GROK_TLS_INSECURE_SKIP_VERIFY` outside a controlled debugging environment. - Do not use session IDs, email addresses, or other sensitive data directly as affinity identifiers. diff --git a/internal/auth/auth.go b/internal/auth/auth.go index b094a62..9fb8e37 100644 --- a/internal/auth/auth.go +++ b/internal/auth/auth.go @@ -23,6 +23,11 @@ import ( var ErrNoAuth = errors.New("no usable credentials in auths directory") +var ( + ErrInvalidCredentialJSON = errors.New("invalid credential JSON") + ErrInvalidCredential = errors.New("invalid credential") +) + type Session struct { Token string `json:"-"` Surface string `json:"surface"` @@ -68,15 +73,22 @@ func loadCredential(path, surface string) (*credential, error) { if err != nil { return nil, err } + return parseCredential(b, path, surface) +} + +func parseCredential(b []byte, path, surface string) (*credential, error) { var raw map[string]any if err := json.Unmarshal(b, &raw); err != nil { - return nil, fmt.Errorf("invalid credential JSON: %w", err) + return nil, fmt.Errorf("%w: %v", ErrInvalidCredentialJSON, err) + } + if raw == nil { + return nil, fmt.Errorf("%w: credential must be a JSON object", ErrInvalidCredential) } node := credentialNode(raw) access := firstString(node, "access_token", "AccessToken", "key", "session_token", "SessionToken") refresh := firstString(node, "refresh_token", "RefreshToken") if access == "" && refresh == "" { - return nil, errors.New("credential has neither access_token nor refresh_token") + return nil, fmt.Errorf("%w: credential has neither access_token nor refresh_token", ErrInvalidCredential) } accessClaims := jwtClaims(access) idClaims := jwtClaims(firstString(node, "id_token", "IDToken")) @@ -88,7 +100,7 @@ func loadCredential(path, surface string) (*credential, error) { subject = claimString(idClaims, "sub") } if subject == "" { - return nil, errors.New("credential has no stable subject") + return nil, fmt.Errorf("%w: credential has no stable subject", ErrInvalidCredential) } clientID := firstString(node, "client_id", "clientId") if clientID == "" { @@ -236,6 +248,10 @@ func writeCredentialAtomic(path string, raw map[string]any) error { if err != nil { return err } + return writeCredentialAtomicMode(path, raw, info.Mode().Perm()) +} + +func writeCredentialAtomicMode(path string, raw map[string]any, mode os.FileMode) error { b, err := json.MarshalIndent(raw, "", " ") if err != nil { return err @@ -247,7 +263,7 @@ func writeCredentialAtomic(path string, raw map[string]any) error { } tmpName := tmp.Name() defer os.Remove(tmpName) - if err := tmp.Chmod(info.Mode().Perm()); err != nil { + if err := tmp.Chmod(mode); err != nil { tmp.Close() return err } diff --git a/internal/auth/auth_test.go b/internal/auth/auth_test.go index 103427e..0a8a338 100644 --- a/internal/auth/auth_test.go +++ b/internal/auth/auth_test.go @@ -5,6 +5,7 @@ import ( "encoding/json" "errors" "fmt" + "io" "net/http" "net/http/httptest" "os" @@ -712,6 +713,161 @@ func TestHotReloadAddsAndRemovesCredentials(t *testing.T) { } } +func TestPoolAllowEmptyImportReplaceAndDelete(t *testing.T) { + dir := t.TempDir() + if _, err := NewPool(context.Background(), PoolConfig{Dir: dir}, nil); !errors.Is(err, ErrNoAuth) { + t.Fatalf("NewPool without AllowEmpty error = %v, want ErrNoAuth", err) + } + pool, err := NewPool(context.Background(), PoolConfig{ + Dir: dir, Surface: "tui", AllowEmpty: true, ReloadInterval: time.Hour, + RefreshConcurrency: 1, AffinityTTL: time.Hour, AffinityMaxEntries: 128, + }, &http.Client{Timeout: time.Second}) + if err != nil { + t.Fatal(err) + } + defer pool.Close() + + first := remoteCredentialJSON(t, "remote-subject", "token-old", []string{"grok-4"}) + info, created, err := pool.ImportCredential(context.Background(), first) + if err != nil { + t.Fatal(err) + } + if !created || info.Status != "ready" || !info.Usable || info.ID == "" { + t.Fatalf("first import = %#v, created=%v", info, created) + } + if len(pool.Credentials()) != 1 { + t.Fatalf("credential count = %d", len(pool.Credentials())) + } + path := filepath.Join(dir, info.ID+".json") + if stat, err := os.Stat(path); err != nil { + t.Fatal(err) + } else if os.PathSeparator != '\\' && stat.Mode().Perm() != 0o600 { + t.Fatalf("credential mode = %o", stat.Mode().Perm()) + } + + replacement := remoteCredentialJSON(t, "remote-subject", "token-new", []string{"grok-4", "grok-new"}) + updated, created, err := pool.ImportCredential(context.Background(), replacement) + if err != nil { + t.Fatal(err) + } + if created || updated.ID != info.ID || !slices.Equal(updated.Models, []string{"grok-4", "grok-new"}) { + t.Fatalf("replacement = %#v, created=%v", updated, created) + } + stored, err := os.ReadFile(path) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(string(stored), "token-new") || strings.Contains(string(stored), "token-old") { + t.Fatal("replacement was not persisted") + } + + if err := pool.DeleteCredential(context.Background(), info.ID); err != nil { + t.Fatal(err) + } + if len(pool.Credentials()) != 0 { + t.Fatalf("credential count after delete = %d", len(pool.Credentials())) + } + if _, err := os.Stat(path); !os.IsNotExist(err) { + t.Fatalf("credential file still exists: %v", err) + } + if _, err := pool.Acquire(context.Background(), Affinity{}, "grok-4", nil); err == nil { + t.Fatal("empty pool Acquire unexpectedly succeeded") + } else { + var unavailable *UnavailableError + if !errors.As(err, &unavailable) { + t.Fatalf("empty pool Acquire error = %T %v", err, err) + } + } +} + +func TestImportCredentialValidation(t *testing.T) { + pool, err := NewPool(context.Background(), PoolConfig{Dir: t.TempDir(), AllowEmpty: true}, nil) + if err != nil { + t.Fatal(err) + } + defer pool.Close() + if _, _, err := pool.ImportCredential(context.Background(), []byte(`{"access_token":`)); !errors.Is(err, ErrInvalidCredentialJSON) { + t.Fatalf("invalid JSON error = %v", err) + } + if _, _, err := pool.ImportCredential(context.Background(), []byte(`{"access_token":"token"}`)); !errors.Is(err, ErrInvalidCredential) { + t.Fatalf("missing subject error = %v", err) + } +} + +func TestImportWaitsForRefreshAndWins(t *testing.T) { + refreshStarted := make(chan struct{}) + releaseRefresh := make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + close(refreshStarted) + <-releaseRefresh + _, _ = io.WriteString(w, `{"access_token":"token-refreshed","refresh_token":"refresh-new","expires_in":3600}`) + })) + defer server.Close() + pool, err := NewPool(context.Background(), PoolConfig{ + Dir: t.TempDir(), Surface: "tui", AllowEmpty: true, ReloadInterval: time.Hour, + RefreshConcurrency: 1, AffinityTTL: time.Hour, AffinityMaxEntries: 128, + }, server.Client()) + if err != nil { + t.Fatal(err) + } + defer pool.Close() + initial, err := json.Marshal(map[string]any{ + "access_token": "token-old", "refresh_token": "refresh-old", "client_id": "client", + "sub": "refresh-race", "expired": time.Now().Add(time.Hour).UTC().Format(time.RFC3339Nano), + "token_endpoint": server.URL, "models": []string{"grok-4"}, + }) + if err != nil { + t.Fatal(err) + } + info, _, err := pool.ImportCredential(context.Background(), initial) + if err != nil { + t.Fatal(err) + } + refreshDone := make(chan error, 1) + go func() { refreshDone <- pool.Refresh(context.Background(), info.ID) }() + <-refreshStarted + replacement := remoteCredentialJSON(t, "refresh-race", "token-uploaded", []string{"grok-4"}) + importDone := make(chan error, 1) + go func() { + _, _, err := pool.ImportCredential(context.Background(), replacement) + importDone <- err + }() + select { + case err := <-importDone: + t.Fatalf("import completed before refresh lock was released: %v", err) + case <-time.After(20 * time.Millisecond): + } + close(releaseRefresh) + if err := <-refreshDone; err != nil { + t.Fatal(err) + } + if err := <-importDone; err != nil { + t.Fatal(err) + } + path := filepath.Join(pool.cfg.Dir, info.ID+".json") + stored, err := os.ReadFile(path) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(string(stored), "token-uploaded") || strings.Contains(string(stored), "token-refreshed") { + t.Fatalf("upload did not win refresh race: %s", stored) + } +} + +func remoteCredentialJSON(t *testing.T, subject, token string, models []string) []byte { + t.Helper() + raw := map[string]any{ + "access_token": token, "refresh_token": "refresh", "client_id": "client", + "sub": subject, "expired": time.Now().Add(time.Hour).UTC().Format(time.RFC3339Nano), + "models": models, "models_updated_at": time.Now().UTC().Format(time.RFC3339Nano), + } + b, err := json.Marshal(raw) + if err != nil { + t.Fatal(err) + } + return b +} + func BenchmarkPoolAcquireTenThousandAccounts(b *testing.B) { p := &Pool{ accounts: map[string]*account{}, states: map[string]accountState{}, diff --git a/internal/auth/pool.go b/internal/auth/pool.go index f7628d3..437a5dd 100644 --- a/internal/auth/pool.go +++ b/internal/auth/pool.go @@ -22,6 +22,8 @@ const stateFileName = ".grokcli2api-state.json" var errAccountBusy = errors.New("credential account is at its in-flight limit") +var ErrCredentialNotFound = errors.New("credential not found") + type AffinityMode uint8 const ( @@ -43,6 +45,20 @@ type PoolConfig struct { AccountMaxInflight int AffinityTTL time.Duration AffinityMaxEntries int + AllowEmpty bool +} + +// CredentialInfo is a redacted view of a credential account. It deliberately +// excludes subjects, file paths, client IDs, and token values. +type CredentialInfo struct { + ID string `json:"id"` + Status string `json:"status"` + Usable bool `json:"usable"` + Disabled bool `json:"disabled"` + ExpiresAt *time.Time `json:"expires_at,omitempty"` + CooldownUntil *time.Time `json:"cooldown_until,omitempty"` + Models []string `json:"models"` + HasRefreshToken bool `json:"has_refresh_token"` } type UnavailableError struct { @@ -182,6 +198,7 @@ type Pool struct { closeOnce sync.Once wg sync.WaitGroup stateMu sync.Mutex + mutationMu sync.Mutex } type Lease struct { @@ -242,6 +259,11 @@ func NewPool(ctx context.Context, cfg PoolConfig, client *http.Client) (*Pool, e if client == nil { client = &http.Client{Timeout: 20 * time.Second} } + if cfg.AllowEmpty { + if err := os.MkdirAll(cfg.Dir, 0o700); err != nil { + return nil, fmt.Errorf("create auths directory: %w", err) + } + } p := &Pool{ cfg: cfg, http: client, accounts: map[string]*account{}, files: map[string]fileEntry{}, states: map[string]accountState{}, affinity: newAffinityCache(cfg.AffinityTTL, cfg.AffinityMaxEntries), @@ -254,10 +276,10 @@ func NewPool(ctx context.Context, cfg PoolConfig, client *http.Client) (*Pool, e if err := p.scan(); err != nil { return nil, err } - if len(p.accounts) == 0 { + if len(p.accounts) == 0 && !cfg.AllowEmpty { return nil, ErrNoAuth } - if !p.hasReadyAccount() { + if len(p.accounts) > 0 && !p.hasReadyAccount() { warmup, cancel := context.WithTimeout(ctx, 30*time.Second) err := p.warmup(warmup) cancel() @@ -312,6 +334,9 @@ func (p *Pool) Acquire(ctx context.Context, affinity Affinity, model string, exc p.rebuildActive() active = p.schedulingSnapshot(model) if len(active) == 0 { + if !p.hasAccounts() { + return nil, p.unavailable(model) + } if model != "" && !p.hasKnownModel(model) { return nil, &ModelUnavailableError{Model: model} } @@ -368,6 +393,9 @@ func (p *Pool) Acquire(ctx context.Context, affinity Affinity, model string, exc } continue } + if !p.hasAccounts() { + return nil, p.unavailable(model) + } if model != "" && !p.hasKnownModel(model) { return nil, &ModelUnavailableError{Model: model} } @@ -375,6 +403,13 @@ func (p *Pool) Acquire(ctx context.Context, affinity Affinity, model string, exc } } +func (p *Pool) hasAccounts() bool { + p.mu.RLock() + hasAccounts := len(p.accounts) > 0 + p.mu.RUnlock() + return hasAccounts +} + func (p *Pool) accountRequestUsable(id string) bool { p.mu.RLock() a := p.accounts[id] @@ -514,6 +549,161 @@ func (p *Pool) AccountIDs() []string { return ids } +// Credentials returns redacted credential metadata sorted by account ID. +func (p *Pool) Credentials() []CredentialInfo { + p.mu.RLock() + items := make([]CredentialInfo, 0, len(p.accounts)) + for id, a := range p.accounts { + items = append(items, credentialInfo(id, a, time.Now())) + } + p.mu.RUnlock() + sort.Slice(items, func(i, j int) bool { return items[i].ID < items[j].ID }) + return items +} + +// Credential returns redacted metadata for one account. +func (p *Pool) Credential(id string) (CredentialInfo, bool) { + p.mu.RLock() + a := p.accounts[id] + if a == nil { + p.mu.RUnlock() + return CredentialInfo{}, false + } + info := credentialInfo(id, a, time.Now()) + p.mu.RUnlock() + return info, true +} + +func credentialInfo(id string, a *account, now time.Time) CredentialInfo { + a.mu.RLock() + defer a.mu.RUnlock() + info := CredentialInfo{ID: id, Disabled: a.disabled, Models: []string{}} + if a.credential == nil { + info.Status = "unavailable" + return info + } + info.Models = append(info.Models, a.credential.Models...) + info.HasRefreshToken = a.credential.RefreshToken != "" + if !a.credential.ExpiresAt.IsZero() { + expires := a.credential.ExpiresAt.UTC() + info.ExpiresAt = &expires + } + if now.Before(a.cooldownUntil) { + cooldown := a.cooldownUntil.UTC() + info.CooldownUntil = &cooldown + } + info.Usable = !a.disabled && !now.Before(a.cooldownUntil) && a.credential.usable(now) + switch { + case a.disabled: + info.Status = "disabled" + case now.Before(a.cooldownUntil): + info.Status = "cooling_down" + case !a.credential.usable(now): + info.Status = "needs_refresh" + case len(a.credential.Models) == 0: + info.Status = "pending_models" + default: + info.Status = "ready" + } + return info +} + +// ImportCredential validates and atomically creates or replaces a credential. +// Account identity, rather than a caller-supplied filename, determines the +// destination so remote uploads cannot escape the configured directory. +func (p *Pool) ImportCredential(ctx context.Context, raw []byte) (CredentialInfo, bool, error) { + parsed, err := parseCredential(raw, "", p.cfg.Surface) + if err != nil { + return CredentialInfo{}, false, err + } + id := accountID(parsed.Subject) + p.mutationMu.Lock() + defer p.mutationMu.Unlock() + + p.mu.RLock() + existing := p.accounts[id] + p.mu.RUnlock() + created := existing == nil + path := filepath.Join(p.cfg.Dir, id+".json") + if existing != nil { + if err := existing.acquireRefresh(ctx); err != nil { + return CredentialInfo{}, false, err + } + defer existing.releaseRefresh() + existing.mu.RLock() + if existing.credential != nil && existing.credential.Path != "" { + path = existing.credential.Path + } + existing.mu.RUnlock() + } + + if _, statErr := os.Stat(path); statErr != nil && !os.IsNotExist(statErr) { + return CredentialInfo{}, false, statErr + } + if err := writeCredentialAtomicMode(path, parsed.Raw, 0o600); err != nil { + return CredentialInfo{}, false, err + } + // Do not let a same-size, same-timestamp replacement reuse cached content. + p.mu.Lock() + delete(p.files, path) + p.mu.Unlock() + if err := p.scanUnlocked(); err != nil { + return CredentialInfo{}, false, err + } + info, ok := p.Credential(id) + if !ok { + return CredentialInfo{}, false, errors.New("credential was saved but not loaded") + } + return info, created, nil +} + +// DeleteCredential removes every valid file for an account and immediately +// evicts that account from scheduling. Existing leases retain their snapshots. +func (p *Pool) DeleteCredential(ctx context.Context, id string) error { + p.mutationMu.Lock() + defer p.mutationMu.Unlock() + p.mu.RLock() + a := p.accounts[id] + p.mu.RUnlock() + if a == nil { + return ErrCredentialNotFound + } + if err := a.acquireRefresh(ctx); err != nil { + return err + } + defer a.releaseRefresh() + + entries, err := os.ReadDir(p.cfg.Dir) + if err != nil { + return err + } + removed := false + for _, entry := range entries { + if entry.IsDir() || strings.HasPrefix(entry.Name(), ".") || !strings.EqualFold(filepath.Ext(entry.Name()), ".json") { + continue + } + path := filepath.Join(p.cfg.Dir, entry.Name()) + cred, loadErr := loadCredential(path, p.cfg.Surface) + if loadErr != nil || accountID(cred.Subject) != id { + continue + } + if err := os.Remove(path); err != nil && !os.IsNotExist(err) { + return err + } + removed = true + } + if !removed { + return ErrCredentialNotFound + } + p.mu.Lock() + delete(p.states, id) + p.mu.Unlock() + if err := p.scanUnlocked(); err != nil { + return err + } + return p.persistState() +} + func (p *Pool) Models() []string { seen := map[string]struct{}{} p.mu.RLock() @@ -1037,6 +1227,12 @@ func (p *Pool) refreshAllDue() { } func (p *Pool) scan() error { + p.mutationMu.Lock() + defer p.mutationMu.Unlock() + return p.scanUnlocked() +} + +func (p *Pool) scanUnlocked() error { entries, err := os.ReadDir(p.cfg.Dir) if err != nil { return fmt.Errorf("read auths directory: %w", err) @@ -1093,6 +1289,9 @@ func (p *Pool) scan() error { slog.Warn("credential pool is empty") return nil } + if p.cfg.AllowEmpty { + return nil + } return ErrNoAuth } p.mu.Lock() diff --git a/internal/config/config.go b/internal/config/config.go index 10713c3..9b08d01 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -39,6 +39,7 @@ type Config struct { ProxyURL string NoProxy []string APIKeys []string + AdminKey string } func Load() (Config, error) { @@ -119,6 +120,7 @@ func Load() (Config, error) { StreamCompression: streamCompression, ProxyURL: strings.TrimSpace(os.Getenv("GROK_PROXY_URL")), NoProxy: splitCSV(os.Getenv("GROK_NO_PROXY")), + AdminKey: strings.TrimSpace(os.Getenv("GROK_ADMIN_KEY")), } cfg.APIKeys = unique(append(splitCSV(os.Getenv("GROK_API_KEYS")), splitCSV(os.Getenv("GROK_API_KEY"))...)) return cfg, nil diff --git a/internal/config/config_test.go b/internal/config/config_test.go index b2877b7..3142e6b 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -12,3 +12,14 @@ func TestLoadValidatesStreamCompression(t *testing.T) { t.Fatalf("Load() error = %v, want stream compression validation error", err) } } + +func TestLoadAdminKey(t *testing.T) { + t.Setenv("GROK_ADMIN_KEY", " admin-secret ") + cfg, err := Load() + if err != nil { + t.Fatal(err) + } + if cfg.AdminKey != "admin-secret" { + t.Fatalf("AdminKey = %q", cfg.AdminKey) + } +} diff --git a/internal/grok/client.go b/internal/grok/client.go index 46b164c..b796cf1 100644 --- a/internal/grok/client.go +++ b/internal/grok/client.go @@ -225,11 +225,30 @@ func (c *Client) InitializeModels(ctx context.Context) error { if len(c.pool.Models()) == 0 { return errors.New("no models discovered from credential accounts") } + c.StartModelRefresh() + return nil +} + +// StartModelRefresh starts the periodic model discovery loop even when the +// credential pool is initially empty and will be bootstrapped through the +// administrator API. +func (c *Client) StartModelRefresh() { c.modelStart.Do(func() { c.modelWG.Add(1) go c.modelRefreshLoop() }) - return nil +} + +// RefreshAccountModels immediately discovers and persists one account's model +// catalog. Credential writes are serialized by the pool; avoiding the batch +// refresh mutex here keeps the caller's context deadline authoritative. +func (c *Client) RefreshAccountModels(ctx context.Context, accountID string) error { + models, err := c.fetchAccountModels(ctx, accountID, false) + if err == nil { + err = c.pool.UpdateModels(accountID, models, time.Now()) + } + c.pool.RebuildSchedulingSnapshot() + return err } func (c *Client) RefreshModels(ctx context.Context, force bool) error { diff --git a/internal/server/server.go b/internal/server/server.go index c805761..34eb300 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -10,6 +10,7 @@ import ( "fmt" "io" "log/slog" + "mime" "net/http" "strings" "time" @@ -38,6 +39,7 @@ func New(cfg config.Config) (*Server, error) { ReloadInterval: cfg.AuthsReloadInterval, RefreshConcurrency: cfg.AuthRefreshConcurrency, AccountMaxInflight: cfg.AccountMaxInflight, AffinityTTL: cfg.AffinityTTL, AffinityMaxEntries: cfg.AffinityMaxEntries, + AllowEmpty: cfg.AdminKey != "", }, httpClient) if err != nil { return nil, err @@ -47,14 +49,20 @@ func New(cfg config.Config) (*Server, error) { pool.Close() return nil, err } - modelCtx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) - err = client.InitializeModels(modelCtx) - cancel() - if err != nil { - client.Close() - pool.Close() - return nil, fmt.Errorf("initialize model catalog: %w", err) + if len(pool.AccountIDs()) > 0 { + modelCtx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) + err = client.InitializeModels(modelCtx) + cancel() + if err != nil && cfg.AdminKey == "" { + client.Close() + pool.Close() + return nil, fmt.Errorf("initialize model catalog: %w", err) + } + if err != nil { + slog.Warn("initial model discovery failed; administrator API remains available", "error", err) + } } + client.StartModelRefresh() s := &Server{cfg: cfg, pool: pool, client: client, mux: http.NewServeMux()} s.routes() return s, nil @@ -83,6 +91,11 @@ func (s *Server) routes() { s.protected("GET /v1/grok/mcp/configs", s.proxyGET("mcp/configs", false)) s.protected("GET /v1/grok/mcp/tools/list", s.proxyGET("mcp/tools/list", false)) s.protected("GET /v1/grok/feedback/config", s.proxyGET("feedback/config", true)) + if s.cfg.AdminKey != "" { + s.mux.Handle("GET /v1/admin/credentials", s.adminKeyGate(http.HandlerFunc(s.adminCredentials))) + s.mux.Handle("POST /v1/admin/credentials", s.adminKeyGate(http.HandlerFunc(s.adminCredentials))) + s.mux.Handle("DELETE /v1/admin/credentials/{id}", s.adminKeyGate(http.HandlerFunc(s.adminCredential))) + } } func (s *Server) protected(pattern string, handler http.HandlerFunc) { @@ -130,6 +143,133 @@ func (s *Server) apiKeyStatus(w http.ResponseWriter, _ *http.Request) { writeJSON(w, http.StatusOK, map[string]any{"enabled": len(s.cfg.APIKeys) > 0, "key_count": len(s.cfg.APIKeys)}) } +const ( + maxCredentialSize = 1 << 20 + maxMultipartUpload = maxCredentialSize + 64<<10 +) + +func (s *Server) adminCredentials(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodGet { + writeJSON(w, http.StatusOK, map[string]any{"object": "list", "data": s.pool.Credentials()}) + return + } + raw, status, err := readCredentialUpload(w, r) + if err != nil { + writeError(w, status, err.Error(), "invalid_request_error", http.StatusText(status)) + return + } + info, created, err := s.pool.ImportCredential(r.Context(), raw) + if err != nil { + switch { + case errors.Is(err, auth.ErrInvalidCredentialJSON): + writeError(w, http.StatusBadRequest, err.Error(), "invalid_request_error", "invalid_json") + case errors.Is(err, auth.ErrInvalidCredential): + writeError(w, http.StatusUnprocessableEntity, err.Error(), "invalid_request_error", "invalid_credential") + default: + writeError(w, http.StatusInternalServerError, "credential could not be saved", "server_error", "credential_write_failed") + } + return + } + + discovery := "succeeded" + response := map[string]any{"credential": info, "created": created} + probeCtx, cancel := context.WithTimeout(r.Context(), 30*time.Second) + probeErr := s.client.RefreshAccountModels(probeCtx, info.ID) + cancel() + if current, ok := s.pool.Credential(info.ID); ok { + response["credential"] = current + } + if probeErr != nil { + discovery = "failed" + response["warning"] = "model discovery failed; credential remains saved" + } + response["model_discovery"] = discovery + status = http.StatusOK + if created { + status = http.StatusCreated + } + slog.Info("credential uploaded", "account", info.ID, "created", created, "model_discovery", discovery) + writeJSON(w, status, response) +} + +func (s *Server) adminCredential(w http.ResponseWriter, r *http.Request) { + id := strings.TrimSpace(r.PathValue("id")) + if !validCredentialID(id) { + writeError(w, http.StatusNotFound, "credential not found", "invalid_request_error", "credential_not_found") + return + } + if err := s.pool.DeleteCredential(r.Context(), id); err != nil { + if errors.Is(err, auth.ErrCredentialNotFound) { + writeError(w, http.StatusNotFound, "credential not found", "invalid_request_error", "credential_not_found") + return + } + writeError(w, http.StatusInternalServerError, "credential could not be deleted", "server_error", "credential_delete_failed") + return + } + slog.Info("credential deleted", "account", id) + w.WriteHeader(http.StatusNoContent) +} + +func readCredentialUpload(w http.ResponseWriter, r *http.Request) ([]byte, int, error) { + mediaType, _, err := mime.ParseMediaType(r.Header.Get("Content-Type")) + if err != nil { + return nil, http.StatusUnsupportedMediaType, errors.New("Content-Type must be application/json or multipart/form-data") + } + switch mediaType { + case "application/json": + r.Body = http.MaxBytesReader(w, r.Body, maxCredentialSize) + body, readErr := io.ReadAll(r.Body) + if readErr != nil { + var tooLarge *http.MaxBytesError + if errors.As(readErr, &tooLarge) { + return nil, http.StatusRequestEntityTooLarge, errors.New("credential exceeds 1 MiB") + } + return nil, http.StatusBadRequest, errors.New("credential body could not be read") + } + return body, 0, nil + case "multipart/form-data": + r.Body = http.MaxBytesReader(w, r.Body, maxMultipartUpload) + if err := r.ParseMultipartForm(maxCredentialSize); err != nil { + var tooLarge *http.MaxBytesError + if errors.As(err, &tooLarge) { + return nil, http.StatusRequestEntityTooLarge, errors.New("credential upload exceeds the size limit") + } + return nil, http.StatusBadRequest, errors.New("invalid multipart upload") + } + if r.MultipartForm == nil { + return nil, http.StatusBadRequest, errors.New("invalid multipart upload") + } + defer r.MultipartForm.RemoveAll() + files := r.MultipartForm.File["file"] + if len(files) != 1 { + return nil, http.StatusBadRequest, errors.New("multipart upload must contain exactly one file field") + } + file, err := files[0].Open() + if err != nil { + return nil, http.StatusBadRequest, errors.New("credential file could not be opened") + } + defer file.Close() + body, err := io.ReadAll(io.LimitReader(file, maxCredentialSize+1)) + if err != nil { + return nil, http.StatusBadRequest, errors.New("credential file could not be read") + } + if len(body) > maxCredentialSize { + return nil, http.StatusRequestEntityTooLarge, errors.New("credential exceeds 1 MiB") + } + return body, 0, nil + default: + return nil, http.StatusUnsupportedMediaType, errors.New("Content-Type must be application/json or multipart/form-data") + } +} + +func validCredentialID(id string) bool { + if len(id) != 24 { + return false + } + _, err := hex.DecodeString(id) + return err == nil && id == strings.ToLower(id) +} + func (s *Server) chat(w http.ResponseWriter, r *http.Request) { timing := grok.NewRequestTiming("chat.completions") r = r.WithContext(grok.WithRequestTiming(r.Context(), timing)) @@ -446,6 +586,25 @@ func (s *Server) apiKeyGate(next http.Handler) http.Handler { }) } +func (s *Server) adminKeyGate(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Cache-Control", "no-store") + candidate := "" + if authz := r.Header.Get("Authorization"); len(authz) >= 7 && strings.EqualFold(authz[:7], "Bearer ") { + candidate = strings.TrimSpace(authz[7:]) + } + if value := strings.TrimSpace(r.Header.Get("X-Admin-Key")); value != "" { + candidate = value + } + if constantEqual(candidate, s.cfg.AdminKey) != 1 { + w.Header().Set("WWW-Authenticate", "Bearer") + writeError(w, http.StatusUnauthorized, "invalid or missing administrator key", "authentication_error", "invalid_admin_key") + return + } + next.ServeHTTP(w, r) + }) +} + func (s *Server) writeClientError(w http.ResponseWriter, err error) { var modelUnavailable *auth.ModelUnavailableError if errors.As(err, &modelUnavailable) { diff --git a/internal/server/server_test.go b/internal/server/server_test.go index e8aff8c..5577aff 100644 --- a/internal/server/server_test.go +++ b/internal/server/server_test.go @@ -2,10 +2,12 @@ package server import ( "bufio" + "bytes" "compress/gzip" "encoding/json" "fmt" "io" + "mime/multipart" "net/http" "net/http/httptest" "os" @@ -951,6 +953,153 @@ func TestPublicRoutesBypassGate(t *testing.T) { } } +func TestAdminCredentialLifecycle(t *testing.T) { + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/v1/models" { + t.Errorf("upstream path = %s", r.URL.Path) + http.NotFound(w, r) + return + } + writeJSON(w, http.StatusOK, map[string]any{"data": []any{map[string]any{"id": "grok-4"}}}) + })) + defer upstream.Close() + s := newAdminTestServer(t, upstream.URL) + defer s.Close() + h := s.Handler() + + payload := adminCredentialPayload(t, "remote-subject", "token-secret") + req := httptest.NewRequest(http.MethodPost, "/v1/admin/credentials", bytes.NewReader(payload)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("X-Admin-Key", "admin-secret") + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + if rec.Code != http.StatusCreated { + t.Fatalf("create status=%d body=%s", rec.Code, rec.Body.String()) + } + if rec.Header().Get("Cache-Control") != "no-store" { + t.Fatalf("Cache-Control = %q", rec.Header().Get("Cache-Control")) + } + var created struct { + Credential auth.CredentialInfo `json:"credential"` + Created bool `json:"created"` + Discovery string `json:"model_discovery"` + } + if err := json.Unmarshal(rec.Body.Bytes(), &created); err != nil { + t.Fatal(err) + } + if !created.Created || created.Discovery != "succeeded" || created.Credential.Status != "ready" || created.Credential.ID == "" { + t.Fatalf("create response = %#v", created) + } + + unauthorized := httptest.NewRequest(http.MethodGet, "/v1/admin/credentials", nil) + unauthorized.Header.Set("Authorization", "Bearer client-key") + rec = httptest.NewRecorder() + h.ServeHTTP(rec, unauthorized) + if rec.Code != http.StatusUnauthorized { + t.Fatalf("ordinary key status = %d", rec.Code) + } + adminOnClientAPI := httptest.NewRequest(http.MethodGet, "/v1/models", nil) + adminOnClientAPI.Header.Set("Authorization", "Bearer admin-secret") + rec = httptest.NewRecorder() + h.ServeHTTP(rec, adminOnClientAPI) + if rec.Code != http.StatusUnauthorized { + t.Fatalf("administrator key on client API status = %d", rec.Code) + } + + list := httptest.NewRequest(http.MethodGet, "/v1/admin/credentials", nil) + list.Header.Set("Authorization", "Bearer admin-secret") + rec = httptest.NewRecorder() + h.ServeHTTP(rec, list) + if rec.Code != http.StatusOK { + t.Fatalf("list status=%d body=%s", rec.Code, rec.Body.String()) + } + for _, secret := range []string{"remote-subject", "token-secret", "refresh-secret", "client-secret"} { + if strings.Contains(rec.Body.String(), secret) { + t.Fatalf("list leaked %q: %s", secret, rec.Body.String()) + } + } + + var multipartBody bytes.Buffer + writer := multipart.NewWriter(&multipartBody) + part, err := writer.CreateFormFile("file", "../../ignored.json") + if err != nil { + t.Fatal(err) + } + if _, err := part.Write(adminCredentialPayload(t, "remote-subject", "token-replaced")); err != nil { + t.Fatal(err) + } + if err := writer.Close(); err != nil { + t.Fatal(err) + } + req = httptest.NewRequest(http.MethodPost, "/v1/admin/credentials", &multipartBody) + req.Header.Set("Content-Type", writer.FormDataContentType()) + req.Header.Set("Authorization", "Bearer admin-secret") + rec = httptest.NewRecorder() + h.ServeHTTP(rec, req) + if rec.Code != http.StatusOK || !strings.Contains(rec.Body.String(), `"created":false`) { + t.Fatalf("replace status=%d body=%s", rec.Code, rec.Body.String()) + } + + deleteReq := httptest.NewRequest(http.MethodDelete, "/v1/admin/credentials/"+created.Credential.ID, nil) + deleteReq.Header.Set("X-Admin-Key", "admin-secret") + rec = httptest.NewRecorder() + h.ServeHTTP(rec, deleteReq) + if rec.Code != http.StatusNoContent { + t.Fatalf("delete status=%d body=%s", rec.Code, rec.Body.String()) + } + + inference := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"grok-4","input":"hello"}`)) + inference.Header.Set("Authorization", "Bearer client-key") + rec = httptest.NewRecorder() + h.ServeHTTP(rec, inference) + if rec.Code != http.StatusServiceUnavailable { + t.Fatalf("empty-pool inference status=%d body=%s", rec.Code, rec.Body.String()) + } +} + +func TestAdminCredentialUploadValidation(t *testing.T) { + s := newAdminTestServer(t, "http://127.0.0.1:1") + defer s.Close() + h := s.Handler() + request := func(contentType string, body io.Reader) *httptest.ResponseRecorder { + req := httptest.NewRequest(http.MethodPost, "/v1/admin/credentials", body) + req.Header.Set("Content-Type", contentType) + req.Header.Set("X-Admin-Key", "admin-secret") + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + return rec + } + if rec := request("text/plain", strings.NewReader("{}")); rec.Code != http.StatusUnsupportedMediaType { + t.Fatalf("media type status = %d", rec.Code) + } + if rec := request("application/json", strings.NewReader(`{"access_token":`)); rec.Code != http.StatusBadRequest { + t.Fatalf("invalid JSON status=%d body=%s", rec.Code, rec.Body.String()) + } + if rec := request("application/json", strings.NewReader(`{"access_token":"token"}`)); rec.Code != http.StatusUnprocessableEntity { + t.Fatalf("invalid credential status=%d body=%s", rec.Code, rec.Body.String()) + } + if rec := request("application/json", bytes.NewReader(make([]byte, maxCredentialSize+1))); rec.Code != http.StatusRequestEntityTooLarge { + t.Fatalf("oversized status=%d body=%s", rec.Code, rec.Body.String()) + } +} + +func TestAdminUploadKeepsCredentialWhenDiscoveryFails(t *testing.T) { + s := newAdminTestServer(t, "http://127.0.0.1:1") + defer s.Close() + req := httptest.NewRequest(http.MethodPost, "/v1/admin/credentials", bytes.NewReader(adminCredentialPayload(t, "offline-subject", "offline-token"))) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("X-Admin-Key", "admin-secret") + rec := httptest.NewRecorder() + s.Handler().ServeHTTP(rec, req) + if rec.Code != http.StatusCreated || !strings.Contains(rec.Body.String(), `"model_discovery":"failed"`) { + t.Fatalf("status=%d body=%s", rec.Code, rec.Body.String()) + } + credentials := s.pool.Credentials() + if len(credentials) != 1 || credentials[0].Status != "pending_models" { + t.Fatalf("persisted credentials = %#v", credentials) + } +} + func TestModelRoutesRequireAPIKey(t *testing.T) { h := newTestHandler(t, "http://127.0.0.1:1", []string{"key"}) for _, path := range []string{"/v1/models", "/v1/models/grok-4"} { @@ -1012,7 +1161,7 @@ func TestModelsEndpointAggregatesCredentialCatalogs(t *testing.T) { func TestRemovedRoutesAre404(t *testing.T) { h := newTestHandler(t, "http://127.0.0.1:1", nil) - for _, path := range []string{"/docs", "/openapi.json", "/v1/health", "/v1/auth/status"} { + for _, path := range []string{"/docs", "/openapi.json", "/v1/health", "/v1/auth/status", "/v1/admin/credentials"} { rec := httptest.NewRecorder() h.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, path, nil)) if rec.Code != http.StatusNotFound { @@ -1076,6 +1225,37 @@ func newTestHandler(t *testing.T, upstream string, keys []string) http.Handler { return newTestHandlerWithTokens(t, upstream, keys, []string{"upstream-token"}) } +func newAdminTestServer(t *testing.T, upstream string) *Server { + t.Helper() + cfg := config.Config{ + ChatProxyBaseURL: upstream, ChatProxyVersion: "v1", AuthsDir: filepath.Join(t.TempDir(), "auths"), + AuthsReloadInterval: time.Hour, AuthRefreshConcurrency: 1, AccountMaxInflight: 2, + ModelsRefreshInterval: 6 * time.Hour, RetryMaxAttempts: 1, RetryBaseDelay: time.Millisecond, + RateLimitCooldown: time.Minute, QuotaCooldown: 24 * time.Hour, + AffinityTTL: time.Hour, AffinityMaxEntries: 128, + ClientName: "grok-shell", ClientVersion: "0.2.93", ClientSurface: "tui", + ClientIdentifier: "grok-shell", TokenAuth: "xai-grok-cli", StreamCompression: "identity", + APIKeys: []string{"client-key"}, AdminKey: "admin-secret", + } + s, err := New(cfg) + if err != nil { + t.Fatal(err) + } + return s +} + +func adminCredentialPayload(t *testing.T, subject, token string) []byte { + t.Helper() + payload, err := json.Marshal(map[string]any{ + "access_token": token, "refresh_token": "refresh-secret", "client_id": "client-secret", + "sub": subject, "expired": time.Now().Add(time.Hour).UTC().Format(time.RFC3339Nano), + }) + if err != nil { + t.Fatal(err) + } + return payload +} + func newTestHandlerWithTokens(t *testing.T, upstream string, keys, tokens []string) http.Handler { return newTestHandlerWithTokensAndCompression(t, upstream, keys, tokens, "") }