mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-15 00:06:24 +00:00
508 lines
16 KiB
Rust
508 lines
16 KiB
Rust
use std::collections::BTreeMap;
|
|
|
|
use axum::{
|
|
body::{to_bytes, Body},
|
|
http::{Request, StatusCode},
|
|
};
|
|
use chrono::Utc;
|
|
use serde_json::{json, Value};
|
|
use tower::ServiceExt;
|
|
|
|
use codex_proxy_rs::{
|
|
app::build_router,
|
|
auth::api_key::ApiKeyHasher,
|
|
config::{
|
|
AdminConfig, ApiConfig, AppConfig, AuthConfig, DatabaseConfig, LoggingConfig, ModelConfig,
|
|
QuotaConfig, QuotaWarningThresholds, SecurityConfig, ServerConfig, TlsConfig,
|
|
UsageStatsConfig,
|
|
},
|
|
crypto::SecretBox,
|
|
state::AppState,
|
|
storage::db::connect_sqlite,
|
|
};
|
|
|
|
fn test_config(database_url: String) -> AppConfig {
|
|
AppConfig {
|
|
server: ServerConfig {
|
|
host: "127.0.0.1".to_string(),
|
|
port: 0,
|
|
},
|
|
api: ApiConfig {
|
|
base_url: "https://chatgpt.com/backend-api".to_string(),
|
|
},
|
|
model: ModelConfig {
|
|
default_model: "gpt-5.5".to_string(),
|
|
default_reasoning_effort: None,
|
|
service_tier: None,
|
|
aliases: BTreeMap::new(),
|
|
},
|
|
auth: AuthConfig {
|
|
refresh_margin_seconds: 300,
|
|
refresh_enabled: true,
|
|
refresh_concurrency: 2,
|
|
max_concurrent_per_account: 3,
|
|
request_interval_ms: 50,
|
|
rotation_strategy: "least_used".to_string(),
|
|
tier_priority: Vec::new(),
|
|
},
|
|
quota: QuotaConfig {
|
|
refresh_interval_minutes: 5,
|
|
warning_thresholds: QuotaWarningThresholds {
|
|
primary: vec![80, 90],
|
|
secondary: vec![80, 90],
|
|
},
|
|
skip_exhausted: true,
|
|
},
|
|
usage_stats: UsageStatsConfig {
|
|
history_retention_days: None,
|
|
},
|
|
database: DatabaseConfig { url: database_url },
|
|
security: SecurityConfig {
|
|
master_key_file: "data/master.key".to_string(),
|
|
api_key_pepper_file: "data/api-key-pepper.key".to_string(),
|
|
},
|
|
tls: TlsConfig {
|
|
force_http11: false,
|
|
},
|
|
admin: AdminConfig {
|
|
session_ttl_minutes: 1440,
|
|
},
|
|
logging: LoggingConfig {
|
|
directory: "logs".to_string(),
|
|
max_file_bytes: 10_485_760,
|
|
retention_days: 14,
|
|
enabled: false,
|
|
capacity: 2_000,
|
|
capture_body: false,
|
|
},
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn admin_api_keys_should_create_list_and_authorize_v1_requests() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let db = dir.path().join("admin-api-keys.sqlite");
|
|
let url = format!("sqlite://{}", db.display());
|
|
let pool = connect_sqlite(&url).await.unwrap();
|
|
seed_admin_session(&pool, "session_1").await;
|
|
let app = build_router(AppState::with_pool_secret_and_api_key_hasher(
|
|
test_config(url),
|
|
pool,
|
|
SecretBox::new([41u8; 32]),
|
|
ApiKeyHasher::new([42u8; 32]),
|
|
));
|
|
|
|
let rejected = app
|
|
.clone()
|
|
.oneshot(
|
|
Request::builder()
|
|
.method("GET")
|
|
.uri("/v1/models")
|
|
.header("authorization", "Bearer cpr_not_stored")
|
|
.body(Body::empty())
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(rejected.status(), StatusCode::UNAUTHORIZED);
|
|
|
|
let create_response = app
|
|
.clone()
|
|
.oneshot(
|
|
Request::builder()
|
|
.method("POST")
|
|
.uri("/admin/api-keys")
|
|
.header("content-type", "application/json")
|
|
.header("cookie", "cpr_admin_session=session_1")
|
|
.header("x-request-id", "req_api_key")
|
|
.body(Body::from(r#"{"name":"cursor"}"#))
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(create_response.status(), StatusCode::OK);
|
|
let body = response_json(create_response).await;
|
|
let plaintext = body["data"]["plaintext"].as_str().unwrap().to_string();
|
|
assert!(plaintext.starts_with("cpr_"));
|
|
assert_eq!(body["requestId"], "req_api_key");
|
|
|
|
let list_response = app
|
|
.clone()
|
|
.oneshot(
|
|
Request::builder()
|
|
.method("GET")
|
|
.uri("/admin/api-keys?limit=10")
|
|
.header("cookie", "cpr_admin_session=session_1")
|
|
.body(Body::empty())
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(list_response.status(), StatusCode::OK);
|
|
let body = response_json(list_response).await;
|
|
assert_eq!(body["data"][0]["name"], "cursor");
|
|
assert!(body["data"][0].get("plaintext").is_none());
|
|
assert!(body["data"][0].get("keyHash").is_none());
|
|
|
|
let models_response = app
|
|
.oneshot(
|
|
Request::builder()
|
|
.method("GET")
|
|
.uri("/v1/models")
|
|
.header("authorization", format!("Bearer {plaintext}"))
|
|
.body(Body::empty())
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(models_response.status(), StatusCode::OK);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn admin_api_key_status_should_disable_and_reenable_client_key_authorization() {
|
|
let (app, _dir) = admin_api_keys_test_app("admin-api-key-status.sqlite").await;
|
|
let (key_id, plaintext) = create_admin_api_key(&app, "session_1", "status-key").await;
|
|
|
|
let disabled = app
|
|
.clone()
|
|
.oneshot(
|
|
Request::builder()
|
|
.method("PATCH")
|
|
.uri(format!("/admin/api-keys/{key_id}/status"))
|
|
.header("content-type", "application/json")
|
|
.header("cookie", "cpr_admin_session=session_1")
|
|
.body(Body::from(r#"{"status":"disabled"}"#))
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(disabled.status(), StatusCode::OK);
|
|
let body = response_json(disabled).await;
|
|
assert_eq!(body["data"]["enabled"], false);
|
|
|
|
let rejected = app
|
|
.clone()
|
|
.oneshot(
|
|
Request::builder()
|
|
.method("GET")
|
|
.uri("/v1/models")
|
|
.header("authorization", format!("Bearer {plaintext}"))
|
|
.body(Body::empty())
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(rejected.status(), StatusCode::UNAUTHORIZED);
|
|
|
|
let enabled = app
|
|
.clone()
|
|
.oneshot(
|
|
Request::builder()
|
|
.method("PATCH")
|
|
.uri(format!("/admin/api-keys/{key_id}/status"))
|
|
.header("content-type", "application/json")
|
|
.header("cookie", "cpr_admin_session=session_1")
|
|
.body(Body::from(r#"{"status":"active"}"#))
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(enabled.status(), StatusCode::OK);
|
|
let body = response_json(enabled).await;
|
|
assert_eq!(body["data"]["enabled"], true);
|
|
|
|
let accepted = app
|
|
.clone()
|
|
.oneshot(
|
|
Request::builder()
|
|
.method("GET")
|
|
.uri("/v1/models")
|
|
.header("authorization", format!("Bearer {plaintext}"))
|
|
.body(Body::empty())
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(accepted.status(), StatusCode::OK);
|
|
|
|
let invalid = app
|
|
.oneshot(
|
|
Request::builder()
|
|
.method("PATCH")
|
|
.uri(format!("/admin/api-keys/{key_id}/status"))
|
|
.header("content-type", "application/json")
|
|
.header("cookie", "cpr_admin_session=session_1")
|
|
.body(Body::from(r#"{"status":"expired"}"#))
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(invalid.status(), StatusCode::BAD_REQUEST);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn admin_api_key_delete_should_remove_client_key_authorization() {
|
|
let (app, _dir) = admin_api_keys_test_app("admin-api-key-delete.sqlite").await;
|
|
let (key_id, plaintext) = create_admin_api_key(&app, "session_1", "delete-key").await;
|
|
|
|
let response = app
|
|
.clone()
|
|
.oneshot(
|
|
Request::builder()
|
|
.method("DELETE")
|
|
.uri(format!("/admin/api-keys/{key_id}"))
|
|
.header("cookie", "cpr_admin_session=session_1")
|
|
.body(Body::empty())
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
let body = response_json(response).await;
|
|
assert_eq!(body["data"]["deleted"], true);
|
|
|
|
let list_response = app
|
|
.clone()
|
|
.oneshot(
|
|
Request::builder()
|
|
.method("GET")
|
|
.uri("/admin/api-keys?limit=10")
|
|
.header("cookie", "cpr_admin_session=session_1")
|
|
.body(Body::empty())
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(list_response.status(), StatusCode::OK);
|
|
let body = response_json(list_response).await;
|
|
assert_eq!(body["data"].as_array().unwrap().len(), 0);
|
|
|
|
let rejected = app
|
|
.clone()
|
|
.oneshot(
|
|
Request::builder()
|
|
.method("GET")
|
|
.uri("/v1/models")
|
|
.header("authorization", format!("Bearer {plaintext}"))
|
|
.body(Body::empty())
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(rejected.status(), StatusCode::UNAUTHORIZED);
|
|
|
|
let missing = app
|
|
.oneshot(
|
|
Request::builder()
|
|
.method("DELETE")
|
|
.uri(format!("/admin/api-keys/{key_id}"))
|
|
.header("cookie", "cpr_admin_session=session_1")
|
|
.body(Body::empty())
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(missing.status(), StatusCode::NOT_FOUND);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn admin_api_key_label_should_update_and_clear_local_client_key_name() {
|
|
let (app, _dir) = admin_api_keys_test_app("admin-api-key-label.sqlite").await;
|
|
let (key_id, _plaintext) = create_admin_api_key(&app, "session_1", "label-key").await;
|
|
|
|
let renamed = app
|
|
.clone()
|
|
.oneshot(
|
|
Request::builder()
|
|
.method("PATCH")
|
|
.uri(format!("/admin/api-keys/{key_id}/label"))
|
|
.header("content-type", "application/json")
|
|
.header("cookie", "cpr_admin_session=session_1")
|
|
.body(Body::from(r#"{"label":"automation"}"#))
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(renamed.status(), StatusCode::OK);
|
|
let body = response_json(renamed).await;
|
|
assert_eq!(body["data"]["name"], "label-key");
|
|
assert_eq!(body["data"]["label"], "automation");
|
|
|
|
let cleared = app
|
|
.clone()
|
|
.oneshot(
|
|
Request::builder()
|
|
.method("PATCH")
|
|
.uri(format!("/admin/api-keys/{key_id}/label"))
|
|
.header("content-type", "application/json")
|
|
.header("cookie", "cpr_admin_session=session_1")
|
|
.body(Body::from(r#"{"label":null}"#))
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(cleared.status(), StatusCode::OK);
|
|
let body = response_json(cleared).await;
|
|
assert_eq!(body["data"]["name"], "label-key");
|
|
assert!(body["data"]["label"].is_null());
|
|
|
|
let too_long = app
|
|
.clone()
|
|
.oneshot(
|
|
Request::builder()
|
|
.method("PATCH")
|
|
.uri(format!("/admin/api-keys/{key_id}/label"))
|
|
.header("content-type", "application/json")
|
|
.header("cookie", "cpr_admin_session=session_1")
|
|
.body(Body::from(json!({ "label": "x".repeat(65) }).to_string()))
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(too_long.status(), StatusCode::BAD_REQUEST);
|
|
|
|
let missing = app
|
|
.oneshot(
|
|
Request::builder()
|
|
.method("PATCH")
|
|
.uri("/admin/api-keys/missing/label")
|
|
.header("content-type", "application/json")
|
|
.header("cookie", "cpr_admin_session=session_1")
|
|
.body(Body::from(r#"{"label":"automation"}"#))
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(missing.status(), StatusCode::NOT_FOUND);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn admin_api_keys_batch_delete_should_remove_found_keys_and_report_missing_ids() {
|
|
let (app, _dir) = admin_api_keys_test_app("admin-api-key-batch-delete.sqlite").await;
|
|
let (key_a, plaintext_a) = create_admin_api_key(&app, "session_1", "batch-a").await;
|
|
let (key_b, plaintext_b) = create_admin_api_key(&app, "session_1", "batch-b").await;
|
|
|
|
let response = app
|
|
.clone()
|
|
.oneshot(
|
|
Request::builder()
|
|
.method("POST")
|
|
.uri("/admin/api-keys/batch-delete")
|
|
.header("content-type", "application/json")
|
|
.header("cookie", "cpr_admin_session=session_1")
|
|
.body(Body::from(
|
|
json!({
|
|
"ids": [key_a, "ghost", key_b]
|
|
})
|
|
.to_string(),
|
|
))
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
let body = response_json(response).await;
|
|
assert_eq!(body["data"]["deleted"], 2);
|
|
assert_eq!(body["data"]["notFound"], json!(["ghost"]));
|
|
|
|
for plaintext in [plaintext_a, plaintext_b] {
|
|
let rejected = app
|
|
.clone()
|
|
.oneshot(
|
|
Request::builder()
|
|
.method("GET")
|
|
.uri("/v1/models")
|
|
.header("authorization", format!("Bearer {plaintext}"))
|
|
.body(Body::empty())
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(rejected.status(), StatusCode::UNAUTHORIZED);
|
|
}
|
|
|
|
let empty = app
|
|
.oneshot(
|
|
Request::builder()
|
|
.method("POST")
|
|
.uri("/admin/api-keys/batch-delete")
|
|
.header("content-type", "application/json")
|
|
.header("cookie", "cpr_admin_session=session_1")
|
|
.body(Body::from(r#"{"ids":[]}"#))
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(empty.status(), StatusCode::BAD_REQUEST);
|
|
}
|
|
|
|
async fn admin_api_keys_test_app(db_name: &str) -> (axum::Router, tempfile::TempDir) {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let db = dir.path().join(db_name);
|
|
let url = format!("sqlite://{}", db.display());
|
|
let pool = connect_sqlite(&url).await.unwrap();
|
|
seed_admin_session(&pool, "session_1").await;
|
|
let app = build_router(AppState::with_pool_secret_and_api_key_hasher(
|
|
test_config(url),
|
|
pool,
|
|
SecretBox::new([43u8; 32]),
|
|
ApiKeyHasher::new([44u8; 32]),
|
|
));
|
|
(app, dir)
|
|
}
|
|
|
|
async fn create_admin_api_key(
|
|
app: &axum::Router,
|
|
session_id: &str,
|
|
name: &str,
|
|
) -> (String, String) {
|
|
let response = app
|
|
.clone()
|
|
.oneshot(
|
|
Request::builder()
|
|
.method("POST")
|
|
.uri("/admin/api-keys")
|
|
.header("content-type", "application/json")
|
|
.header("cookie", format!("cpr_admin_session={session_id}"))
|
|
.body(Body::from(json!({ "name": name }).to_string()))
|
|
.unwrap(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
let body = response_json(response).await;
|
|
(
|
|
body["data"]["id"].as_str().unwrap().to_string(),
|
|
body["data"]["plaintext"].as_str().unwrap().to_string(),
|
|
)
|
|
}
|
|
|
|
async fn seed_admin_session(pool: &sqlx::SqlitePool, session_id: &str) {
|
|
let now = Utc::now().to_rfc3339();
|
|
sqlx::query(
|
|
"insert into admin_users (id, password_hash, created_at, updated_at) values (?, ?, ?, ?)",
|
|
)
|
|
.bind("admin_1")
|
|
.bind("hash")
|
|
.bind(&now)
|
|
.bind(&now)
|
|
.execute(pool)
|
|
.await
|
|
.unwrap();
|
|
sqlx::query(
|
|
"insert into admin_sessions (id, user_id, expires_at, created_at) values (?, ?, ?, ?)",
|
|
)
|
|
.bind(session_id)
|
|
.bind("admin_1")
|
|
.bind("2999-01-01T00:00:00Z")
|
|
.bind(now)
|
|
.execute(pool)
|
|
.await
|
|
.unwrap();
|
|
}
|
|
|
|
async fn response_json(response: axum::response::Response) -> Value {
|
|
let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap();
|
|
serde_json::from_slice(&bytes).unwrap()
|
|
}
|