330 lines
12 KiB
Rust
330 lines
12 KiB
Rust
use rusqlite::{params, Connection};
|
||
use serde::{Deserialize, Serialize};
|
||
|
||
#[derive(Debug, Serialize, Deserialize, Clone)]
|
||
pub struct ApiKey {
|
||
pub id: String,
|
||
pub name: String,
|
||
pub provider: String,
|
||
pub base_url: Option<String>,
|
||
pub model: Option<String>,
|
||
pub notes: Option<String>,
|
||
pub created_at: Option<String>,
|
||
pub updated_at: Option<String>,
|
||
}
|
||
|
||
#[derive(Debug, Deserialize)]
|
||
pub struct SaveApiKeyInput {
|
||
pub name: String,
|
||
pub provider: String,
|
||
pub api_key: String,
|
||
pub base_url: Option<String>,
|
||
pub model: Option<String>,
|
||
pub notes: Option<String>,
|
||
}
|
||
|
||
fn row_to_api_key(row: &rusqlite::Row<'_>) -> rusqlite::Result<ApiKey> {
|
||
Ok(ApiKey {
|
||
id: row.get(0)?,
|
||
name: row.get(1)?,
|
||
provider: row.get(2)?,
|
||
base_url: row.get(3)?,
|
||
model: row.get(4)?,
|
||
notes: row.get(5)?,
|
||
created_at: row.get(6)?,
|
||
updated_at: row.get(7)?,
|
||
})
|
||
}
|
||
|
||
const SELECT_COLS: &str = "id, name, provider, base_url, model, notes, created_at, updated_at";
|
||
|
||
pub fn list_api_keys(conn: &Connection) -> Vec<ApiKey> {
|
||
let sql = format!("SELECT {SELECT_COLS} FROM api_keys ORDER BY updated_at DESC");
|
||
let mut stmt = match conn.prepare(&sql) {
|
||
Ok(s) => s,
|
||
Err(_) => return vec![],
|
||
};
|
||
stmt.query_map([], row_to_api_key)
|
||
.map(|iter| iter.filter_map(|r| r.ok()).collect())
|
||
.unwrap_or_default()
|
||
}
|
||
|
||
pub fn create_api_key(
|
||
conn: &Connection,
|
||
store: &impl crate::credential_store::CredentialStore,
|
||
input: &SaveApiKeyInput,
|
||
) -> Result<ApiKey, String> {
|
||
let id = format!("ak-{}", uuid::Uuid::new_v4());
|
||
let reference = crate::credential_store::versioned_reference(
|
||
crate::credential_store::LegacyCredentialKind::ApiKey,
|
||
&id,
|
||
);
|
||
crate::credential_store::put_verified(store, &reference, input.api_key.as_bytes())?;
|
||
let result = (|| {
|
||
let tx = conn
|
||
.unchecked_transaction()
|
||
.map_err(|_| "创建 API Key 失败".to_string())?;
|
||
tx.execute(
|
||
"INSERT INTO api_keys (id, name, provider, api_key, credential_ref, base_url, model, notes)
|
||
VALUES (?1,?2,?3,'',?4,?5,?6,?7)",
|
||
params![id, input.name, input.provider, reference, input.base_url, input.model, input.notes],
|
||
).map_err(|_| "创建 API Key 失败".to_string())?;
|
||
let key = get_api_key(&tx, &id).ok_or_else(|| "创建后查询失败".to_string())?;
|
||
tx.commit().map_err(|_| "创建 API Key 失败".to_string())?;
|
||
Ok(key)
|
||
})();
|
||
if result.is_err() {
|
||
let _ = store.delete(&reference);
|
||
}
|
||
result
|
||
}
|
||
|
||
pub fn get_api_key(conn: &Connection, id: &str) -> Option<ApiKey> {
|
||
let sql = format!("SELECT {SELECT_COLS} FROM api_keys WHERE id = ?1");
|
||
conn.query_row(&sql, params![id], row_to_api_key).ok()
|
||
}
|
||
|
||
pub fn update_api_key(
|
||
conn: &Connection,
|
||
store: &impl crate::credential_store::CredentialStore,
|
||
id: &str,
|
||
input: &SaveApiKeyInput,
|
||
) -> Result<ApiKey, String> {
|
||
let old_reference: Option<String> = conn
|
||
.query_row(
|
||
"SELECT credential_ref FROM api_keys WHERE id = ?1",
|
||
params![id],
|
||
|row| row.get(0),
|
||
)
|
||
.map_err(|_| "API Key 不存在".to_string())?;
|
||
let reference = crate::credential_store::versioned_reference(
|
||
crate::credential_store::LegacyCredentialKind::ApiKey,
|
||
id,
|
||
);
|
||
crate::credential_store::put_verified(store, &reference, input.api_key.as_bytes())?;
|
||
let result = (|| {
|
||
let tx = conn
|
||
.unchecked_transaction()
|
||
.map_err(|_| "更新 API Key 失败".to_string())?;
|
||
let rows = tx.execute(
|
||
"UPDATE api_keys SET name=?1, provider=?2, api_key='', credential_ref=?3, base_url=?4, model=?5, notes=?6, updated_at=datetime('now')
|
||
WHERE id=?7 AND credential_ref IS ?8",
|
||
params![input.name, input.provider, reference, input.base_url, input.model, input.notes, id, old_reference],
|
||
).map_err(|_| "更新 API Key 失败".to_string())?;
|
||
if rows != 1 {
|
||
return Err("更新 API Key 失败：凭证引用已变化".into());
|
||
}
|
||
let key = get_api_key(&tx, id).ok_or_else(|| "更新后查询失败".to_string())?;
|
||
tx.commit().map_err(|_| "更新 API Key 失败".to_string())?;
|
||
Ok(key)
|
||
})();
|
||
if result.is_err() {
|
||
let _ = store.delete(&reference);
|
||
} else if let Some(old) = old_reference {
|
||
let _ = store.delete(&old);
|
||
}
|
||
result
|
||
}
|
||
|
||
pub fn delete_api_key_fn(
|
||
conn: &Connection,
|
||
store: &impl crate::credential_store::CredentialStore,
|
||
id: &str,
|
||
) -> Result<(), String> {
|
||
let tx = conn
|
||
.unchecked_transaction()
|
||
.map_err(|_| "删除 API Key 失败".to_string())?;
|
||
let reference: Option<String> = tx
|
||
.query_row(
|
||
"SELECT credential_ref FROM api_keys WHERE id=?1",
|
||
params![id],
|
||
|row| row.get(0),
|
||
)
|
||
.map_err(|_| "API Key 不存在".to_string())?;
|
||
tx.execute("DELETE FROM api_keys WHERE id=?1", params![id])
|
||
.map_err(|_| "删除 API Key 失败".to_string())?;
|
||
let references: Vec<String> = reference.into_iter().collect();
|
||
let snapshots = crate::credential_store::delete_references(store, &references)?;
|
||
if tx.commit().is_err() {
|
||
crate::credential_store::restore_references(store, &snapshots)?;
|
||
return Err("删除 API Key 失败：数据库提交失败".into());
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
// ── Tauri commands ───────────────────────────────────────────────────────────
|
||
|
||
#[tauri::command]
|
||
pub fn get_api_keys() -> Result<Vec<ApiKey>, String> {
|
||
let conn = crate::db::pool().get().map_err(|e| e.to_string())?;
|
||
Ok(list_api_keys(&conn))
|
||
}
|
||
|
||
#[tauri::command]
|
||
pub fn add_api_key(input: SaveApiKeyInput) -> Result<ApiKey, String> {
|
||
let conn = crate::db::pool().get().map_err(|e| e.to_string())?;
|
||
create_api_key(
|
||
&conn,
|
||
&crate::credential_store::WindowsCredentialStore,
|
||
&input,
|
||
)
|
||
}
|
||
|
||
#[tauri::command]
|
||
pub fn edit_api_key(id: String, input: SaveApiKeyInput) -> Result<ApiKey, String> {
|
||
let conn = crate::db::pool().get().map_err(|e| e.to_string())?;
|
||
update_api_key(
|
||
&conn,
|
||
&crate::credential_store::WindowsCredentialStore,
|
||
&id,
|
||
&input,
|
||
)
|
||
}
|
||
|
||
#[tauri::command]
|
||
pub fn delete_api_key(id: String) -> Result<(), String> {
|
||
let conn = crate::db::pool().get().map_err(|e| e.to_string())?;
|
||
delete_api_key_fn(&conn, &crate::credential_store::WindowsCredentialStore, &id)
|
||
}
|
||
|
||
pub fn reveal_api_key_for_conn(
|
||
conn: &Connection,
|
||
store: &impl crate::credential_store::CredentialStore,
|
||
id: &str,
|
||
) -> Result<String, String> {
|
||
let bytes = crate::credential_store::get_or_migrate_legacy(
|
||
conn,
|
||
store,
|
||
crate::credential_store::LegacyCredentialKind::ApiKey,
|
||
id,
|
||
)?;
|
||
String::from_utf8(bytes).map_err(|_| "凭证状态 needs_repair：API Key 编码无效".to_string())
|
||
}
|
||
|
||
/// 仅供本机 UI 在用户显式请求查看或复制时调用；MCP 不公开此 command。
|
||
#[tauri::command]
|
||
pub fn reveal_api_key(id: String) -> Result<String, String> {
|
||
let conn = crate::db::pool().get().map_err(|e| e.to_string())?;
|
||
reveal_api_key_for_conn(&conn, &crate::credential_store::WindowsCredentialStore, &id)
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
use crate::credential_store::CredentialStore;
|
||
|
||
#[test]
|
||
fn list_hides_key_and_explicit_reveal_requires_existing_id() {
|
||
let conn = crate::db::conn_with_schema();
|
||
let store = crate::credential_store::MemoryCredentialStore::new();
|
||
conn.execute("INSERT INTO api_keys (id, name, provider, api_key) VALUES ('k1', 'key', 'test', 'secret')", []).unwrap();
|
||
let listed = list_api_keys(&conn);
|
||
assert!(serde_json::to_string(&listed).unwrap().contains("secret") == false);
|
||
assert_eq!(
|
||
reveal_api_key_for_conn(&conn, &store, "k1").unwrap(),
|
||
"secret"
|
||
);
|
||
assert!(reveal_api_key_for_conn(&conn, &store, "missing").is_err());
|
||
let (legacy, reference): (String, String) = conn
|
||
.query_row(
|
||
"SELECT api_key, credential_ref FROM api_keys WHERE id='k1'",
|
||
[],
|
||
|row| Ok((row.get(0)?, row.get(1)?)),
|
||
)
|
||
.unwrap();
|
||
assert!(legacy.is_empty());
|
||
assert!(store.get(&reference).unwrap().is_some());
|
||
}
|
||
|
||
#[test]
|
||
fn api_key_write_failure_preserves_legacy_and_metadata() {
|
||
let conn = crate::db::conn_with_schema();
|
||
conn.execute("INSERT INTO api_keys (id, name, provider, api_key) VALUES ('k1', 'old', 'test', 'synthetic-old-key')", []).unwrap();
|
||
let failing = crate::credential_store::MemoryCredentialStore::failing();
|
||
let input = SaveApiKeyInput {
|
||
name: "new".into(),
|
||
provider: "test".into(),
|
||
api_key: "synthetic-new-key".into(),
|
||
base_url: None,
|
||
model: None,
|
||
notes: None,
|
||
};
|
||
let error = update_api_key(&conn, &failing, "k1", &input).unwrap_err();
|
||
assert!(error.contains("needs_repair"));
|
||
assert!(!error.contains("synthetic-new-key"));
|
||
let (name, old, reference): (String, String, Option<String>) = conn
|
||
.query_row(
|
||
"SELECT name, api_key, credential_ref FROM api_keys WHERE id='k1'",
|
||
[],
|
||
|row| Ok((row.get(0)?, row.get(1)?, row.get(2)?)),
|
||
)
|
||
.unwrap();
|
||
assert_eq!(name, "old");
|
||
assert_eq!(old, "synthetic-old-key");
|
||
assert!(reference.is_none());
|
||
|
||
let store = crate::credential_store::MemoryCredentialStore::new();
|
||
update_api_key(&conn, &store, "k1", &input).unwrap();
|
||
let listed = list_api_keys(&conn);
|
||
assert!(!serde_json::to_string(&listed)
|
||
.unwrap()
|
||
.contains("synthetic-new-key"));
|
||
assert_eq!(
|
||
reveal_api_key_for_conn(&conn, &store, "k1").unwrap(),
|
||
"synthetic-new-key"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn new_api_key_keeps_only_reference_in_database_and_delete_cleans_store() {
|
||
let conn = crate::db::conn_with_schema();
|
||
let store = crate::credential_store::MemoryCredentialStore::new();
|
||
let input = SaveApiKeyInput {
|
||
name: "test".into(),
|
||
provider: "test".into(),
|
||
api_key: "synthetic-new-key".into(),
|
||
base_url: None,
|
||
model: None,
|
||
notes: None,
|
||
};
|
||
let key = create_api_key(&conn, &store, &input).unwrap();
|
||
let (legacy, reference): (String, String) = conn
|
||
.query_row(
|
||
"SELECT api_key, credential_ref FROM api_keys WHERE id=?1",
|
||
params![key.id],
|
||
|row| Ok((row.get(0)?, row.get(1)?)),
|
||
)
|
||
.unwrap();
|
||
assert!(legacy.is_empty());
|
||
assert_eq!(
|
||
store.get(&reference).unwrap().as_deref(),
|
||
Some(b"synthetic-new-key".as_slice())
|
||
);
|
||
assert_eq!(
|
||
reveal_api_key_for_conn(&conn, &store, &key.id).unwrap(),
|
||
"synthetic-new-key"
|
||
);
|
||
delete_api_key_fn(&conn, &store, &key.id).unwrap();
|
||
assert!(store.get(&reference).unwrap().is_none());
|
||
}
|
||
|
||
#[test]
|
||
fn delete_failure_keeps_api_key_and_credential_reference() {
|
||
let conn = crate::db::conn_with_schema();
|
||
let store = crate::credential_store::MemoryCredentialStore::failing_deletes();
|
||
store.put("api-key/k1/v1", b"synthetic-key").unwrap();
|
||
conn.execute(
|
||
"INSERT INTO api_keys (id, name, provider, api_key, credential_ref) VALUES ('k1', 'key', 'test', '', 'api-key/k1/v1')",
|
||
[],
|
||
).unwrap();
|
||
assert!(delete_api_key_fn(&conn, &store, "k1")
|
||
.unwrap_err()
|
||
.contains("needs_repair"));
|
||
assert!(get_api_key(&conn, "k1").is_some());
|
||
assert_eq!(
|
||
store.get("api-key/k1/v1").unwrap().as_deref(),
|
||
Some(b"synthetic-key".as_slice())
|
||
);
|
||
}
|
||
}
|