diff --git a/migrations/0004_client_api_key_labels.sql b/migrations/0004_client_api_key_labels.sql new file mode 100644 index 000000000..354a3edd7 --- /dev/null +++ b/migrations/0004_client_api_key_labels.sql @@ -0,0 +1 @@ +alter table client_api_keys add column label text; diff --git a/src/app.rs b/src/app.rs index ca7270835..07f71a41d 100644 --- a/src/app.rs +++ b/src/app.rs @@ -8,11 +8,11 @@ use tower_http::trace::TraceLayer; use crate::{ http::{ admin::{ - accounts, api_keys, batch_delete_accounts, batch_update_account_status, create_api_key, - delete_account, delete_account_cookies, delete_api_key, get_account_cookies, - import_accounts, login, logs, refresh_models, set_account_cookies, settings, - update_account_label, update_account_status, update_api_key_status, usage_stats, - usage_stats_summary, + accounts, api_keys, batch_delete_accounts, batch_delete_api_keys, + batch_update_account_status, create_api_key, delete_account, delete_account_cookies, + delete_api_key, get_account_cookies, import_accounts, login, logs, refresh_models, + set_account_cookies, settings, update_account_label, update_account_status, + update_api_key_label, update_api_key_status, usage_stats, usage_stats_summary, }, health::health, middleware::attach_request_id, @@ -63,7 +63,12 @@ pub fn build_router(state: AppState) -> Router { ) .route("/admin/accounts/import", post(import_accounts)) .route("/admin/api-keys", get(api_keys).post(create_api_key)) + .route("/admin/api-keys/batch-delete", post(batch_delete_api_keys)) .route("/admin/api-keys/{key_id}", delete(delete_api_key)) + .route( + "/admin/api-keys/{key_id}/label", + patch(update_api_key_label), + ) .route( "/admin/api-keys/{key_id}/status", patch(update_api_key_status), diff --git a/src/auth/api_key_repository.rs b/src/auth/api_key_repository.rs index b91fcc804..2faaf6dc4 100644 --- a/src/auth/api_key_repository.rs +++ b/src/auth/api_key_repository.rs @@ -24,6 +24,7 @@ pub type ClientApiKeyRepositoryResult = Result, pub prefix: String, pub enabled: bool, pub created_at: String, @@ -60,6 +61,7 @@ impl ClientApiKeyRepository { Ok(StoredClientApiKey { id, name: name.to_string(), + label: None, prefix: generated.prefix.clone(), enabled: true, created_at: now, @@ -77,7 +79,7 @@ impl ClientApiKeyRepository { let (created_at, id) = decode_cursor(&cursor).ok_or(ClientApiKeyRepositoryError::InvalidCursor)?; sqlx::query( - "select id, name, prefix, enabled, created_at, last_used_at from client_api_keys where created_at < ? or (created_at = ? and id < ?) order by created_at desc, id desc limit ?", + "select id, name, label, prefix, enabled, created_at, last_used_at from client_api_keys where created_at < ? or (created_at = ? and id < ?) order by created_at desc, id desc limit ?", ) .bind(&created_at) .bind(created_at) @@ -87,7 +89,7 @@ impl ClientApiKeyRepository { .await? } else { sqlx::query( - "select id, name, prefix, enabled, created_at, last_used_at from client_api_keys order by created_at desc, id desc limit ?", + "select id, name, label, prefix, enabled, created_at, last_used_at from client_api_keys order by created_at desc, id desc limit ?", ) .bind(fetch_limit) .fetch_all(&self.pool) @@ -111,6 +113,16 @@ impl ClientApiKeyRepository { Ok(Page { items, next_cursor }) } + pub async fn get(&self, id: &str) -> ClientApiKeyRepositoryResult> { + let row = sqlx::query( + "select id, name, label, prefix, enabled, created_at, last_used_at from client_api_keys where id = ?", + ) + .bind(id) + .fetch_optional(&self.pool) + .await?; + Ok(row.map(|row| key_from_row(&row))) + } + pub async fn verify_and_touch( &self, plaintext: &str, @@ -147,6 +159,22 @@ impl ClientApiKeyRepository { Ok(result.rows_affected() > 0) } + pub async fn set_label( + &self, + id: &str, + label: Option, + ) -> ClientApiKeyRepositoryResult> { + let result = sqlx::query("update client_api_keys set label = ? where id = ?") + .bind(label) + .bind(id) + .execute(&self.pool) + .await?; + if result.rows_affected() == 0 { + return Ok(None); + } + self.get(id).await + } + pub async fn delete(&self, id: &str) -> ClientApiKeyRepositoryResult { let result = sqlx::query("delete from client_api_keys where id = ?") .bind(id) @@ -160,6 +188,7 @@ fn key_from_row(row: &sqlx::sqlite::SqliteRow) -> StoredClientApiKey { StoredClientApiKey { id: row.get("id"), name: row.get("name"), + label: row.get("label"), prefix: row.get("prefix"), enabled: row.get::("enabled") != 0, created_at: row.get("created_at"), diff --git a/src/http/admin.rs b/src/http/admin.rs index ed667a5c2..2ae15d459 100644 --- a/src/http/admin.rs +++ b/src/http/admin.rs @@ -325,6 +325,7 @@ pub struct CreateApiKeyRequest { pub struct ClientApiKeyData { pub id: String, pub name: String, + pub label: Option, pub prefix: String, pub enabled: bool, pub created_at: String, @@ -336,6 +337,7 @@ pub struct ClientApiKeyData { pub struct CreatedClientApiKeyData { pub id: String, pub name: String, + pub label: Option, pub prefix: String, pub enabled: bool, pub created_at: String, @@ -349,6 +351,12 @@ pub struct UpdateClientApiKeyStatusRequest { pub status: String, } +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct UpdateClientApiKeyLabelRequest { + pub label: Option, +} + #[derive(Debug, Serialize)] #[serde(rename_all = "camelCase")] pub struct UpdateClientApiKeyStatusData { @@ -356,6 +364,19 @@ pub struct UpdateClientApiKeyStatusData { pub enabled: bool, } +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct BatchDeleteClientApiKeysRequest { + pub ids: Vec, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct BatchDeleteClientApiKeysData { + pub deleted: u32, + pub not_found: Vec, +} + #[derive(Debug, Serialize)] #[serde(rename_all = "camelCase")] pub struct DeleteClientApiKeyData { @@ -367,6 +388,7 @@ impl From for ClientApiKeyData { Self { id: key.id, name: key.name, + label: key.label, prefix: key.prefix, enabled: key.enabled, created_at: key.created_at, @@ -380,6 +402,7 @@ impl CreatedClientApiKeyData { Self { id: key.id, name: key.name, + label: key.label, prefix: key.prefix, enabled: key.enabled, created_at: key.created_at, @@ -1820,6 +1843,136 @@ pub async fn create_api_key( } } +pub async fn batch_delete_api_keys( + State(state): State, + Extension(request_id): Extension, + headers: HeaderMap, + Json(payload): Json, +) -> Response { + let request_id = request_id.as_str().to_string(); + if payload.ids.is_empty() { + return AdminResponse::new( + StatusCode::BAD_REQUEST, + AdminEnvelope::new(40001, "API key ids are required", (), request_id), + ) + .into_response(); + } + let Some(pool) = state.db() else { + return AdminResponse::new( + StatusCode::INTERNAL_SERVER_ERROR, + AdminEnvelope::new(50001, "Database is not initialized", (), request_id), + ) + .into_response(); + }; + if let Err(response) = require_admin_session(pool, &headers, &request_id).await { + return response; + } + let Some(repo) = state.client_api_key_repository() else { + return AdminResponse::new( + StatusCode::INTERNAL_SERVER_ERROR, + AdminEnvelope::new( + 50001, + "API key repository is not initialized", + (), + request_id, + ), + ) + .into_response(); + }; + + let mut deleted = 0u32; + let mut not_found = Vec::new(); + for key_id in payload.ids { + match repo.delete(&key_id).await { + Ok(true) => deleted += 1, + Ok(false) => not_found.push(key_id), + Err(_) => { + return AdminResponse::new( + StatusCode::INTERNAL_SERVER_ERROR, + AdminEnvelope::new(50001, "Failed to delete API key", (), request_id), + ) + .into_response(); + } + } + } + + AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok( + BatchDeleteClientApiKeysData { deleted, not_found }, + request_id, + ), + ) + .into_response() +} + +pub async fn update_api_key_label( + State(state): State, + Extension(request_id): Extension, + headers: HeaderMap, + Path(key_id): Path, + Json(payload): Json, +) -> Response { + let request_id = request_id.as_str().to_string(); + if payload + .label + .as_ref() + .is_some_and(|label| label.chars().count() > 64) + { + return AdminResponse::new( + StatusCode::BAD_REQUEST, + AdminEnvelope::new( + 40001, + "API key label must be 64 characters or fewer", + (), + request_id, + ), + ) + .into_response(); + } + let Some(pool) = state.db() else { + return AdminResponse::new( + StatusCode::INTERNAL_SERVER_ERROR, + AdminEnvelope::new(50001, "Database is not initialized", (), request_id), + ) + .into_response(); + }; + if let Err(response) = require_admin_session(pool, &headers, &request_id).await { + return response; + } + let Some(repo) = state.client_api_key_repository() else { + return AdminResponse::new( + StatusCode::INTERNAL_SERVER_ERROR, + AdminEnvelope::new( + 50001, + "API key repository is not initialized", + (), + request_id, + ), + ) + .into_response(); + }; + + // label 是后台显示备注，不能影响 cpr_ key 的 hash、prefix 或启用状态。 + match repo.set_label(&key_id, payload.label).await { + Ok(Some(key)) => AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(ClientApiKeyData::from(key), request_id), + ) + .into_response(), + Ok(None) => AdminResponse::new( + StatusCode::NOT_FOUND, + AdminEnvelope::new(40401, "API key not found", (), request_id), + ) + .into_response(), + Err(_) => AdminResponse::new( + StatusCode::INTERNAL_SERVER_ERROR, + AdminEnvelope::new(50001, "Failed to update API key label", (), request_id), + ) + .into_response(), + } +} + pub async fn update_api_key_status( State(state): State, Extension(request_id): Extension, diff --git a/tests/admin_api_keys_route_test.rs b/tests/admin_api_keys_route_test.rs index b05214001..0fa0d01a7 100644 --- a/tests/admin_api_keys_route_test.rs +++ b/tests/admin_api_keys_route_test.rs @@ -305,6 +305,137 @@ async fn admin_api_key_delete_should_remove_client_key_authorization() { 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);