mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 17:30:23 +08:00
Support per-format provider key auth
This commit is contained in:
@@ -230,7 +230,7 @@ fn merge_comma_header_values(left: Option<&str>, right: Option<&str>) -> Option<
|
||||
pub fn resolve_local_openai_bearer_auth(
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
) -> Option<(String, String)> {
|
||||
let auth_type = transport.key.auth_type.trim().to_ascii_lowercase();
|
||||
let auth_type = resolve_local_auth_type_for_transport_format(transport);
|
||||
if !matches!(auth_type.as_str(), "api_key" | "bearer") {
|
||||
return None;
|
||||
}
|
||||
@@ -242,7 +242,7 @@ pub fn resolve_local_openai_bearer_auth(
|
||||
pub fn resolve_local_standard_auth(
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
) -> Option<(String, String)> {
|
||||
let auth_type = transport.key.auth_type.trim().to_ascii_lowercase();
|
||||
let auth_type = resolve_local_auth_type_for_transport_format(transport);
|
||||
let secret = resolved_local_secret(transport)?;
|
||||
|
||||
match auth_type.as_str() {
|
||||
@@ -255,7 +255,7 @@ pub fn resolve_local_standard_auth(
|
||||
pub fn resolve_local_gemini_auth(
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
) -> Option<(String, String)> {
|
||||
let auth_type = transport.key.auth_type.trim().to_ascii_lowercase();
|
||||
let auth_type = resolve_local_auth_type_for_transport_format(transport);
|
||||
let secret = resolved_local_secret(transport)?;
|
||||
|
||||
match auth_type.as_str() {
|
||||
@@ -265,6 +265,30 @@ pub fn resolve_local_gemini_auth(
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn resolve_local_auth_type_for_transport_format(
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
) -> String {
|
||||
let default_auth_type = transport.key.auth_type.trim().to_ascii_lowercase();
|
||||
let api_format = aether_ai_formats::normalize_api_format_alias(&transport.endpoint.api_format);
|
||||
let Some(overrides) = transport
|
||||
.key
|
||||
.auth_type_by_format
|
||||
.as_ref()
|
||||
.and_then(serde_json::Value::as_object)
|
||||
else {
|
||||
return default_auth_type;
|
||||
};
|
||||
|
||||
overrides
|
||||
.get(&api_format)
|
||||
.or_else(|| overrides.get(transport.endpoint.api_format.trim()))
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(str::trim)
|
||||
.map(str::to_ascii_lowercase)
|
||||
.filter(|value| matches!(value.as_str(), "api_key" | "bearer"))
|
||||
.unwrap_or(default_auth_type)
|
||||
}
|
||||
|
||||
fn resolved_local_secret(transport: &GatewayProviderTransportSnapshot) -> Option<&str> {
|
||||
let secret = transport.key.decrypted_api_key.trim();
|
||||
(!secret.is_empty() && secret != PLACEHOLDER_API_KEY).then_some(secret)
|
||||
@@ -322,6 +346,8 @@ mod tests {
|
||||
auth_type: "bearer".to_string(),
|
||||
is_active: true,
|
||||
api_formats: None,
|
||||
auth_type_by_format: None,
|
||||
|
||||
allowed_models: None,
|
||||
capabilities: None,
|
||||
rate_multipliers: None,
|
||||
@@ -443,6 +469,15 @@ mod tests {
|
||||
assert!(resolve_local_standard_auth(&sample_transport()).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn local_standard_auth_rejects_empty_secret() {
|
||||
let mut transport = sample_transport();
|
||||
transport.key.auth_type = "api_key".to_string();
|
||||
transport.key.decrypted_api_key = String::new();
|
||||
|
||||
assert!(resolve_local_standard_auth(&transport).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn local_openai_bearer_auth_maps_api_key_to_bearer_authorization() {
|
||||
let mut transport = sample_transport();
|
||||
@@ -466,4 +501,35 @@ mod tests {
|
||||
Some(("authorization".to_string(), "Bearer sk-openai".to_string(),))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn local_standard_auth_uses_format_auth_type_override() {
|
||||
let mut transport = sample_transport();
|
||||
transport.key.auth_type = "api_key".to_string();
|
||||
transport.key.auth_type_by_format = Some(serde_json::json!({
|
||||
"claude:messages": "bearer"
|
||||
}));
|
||||
transport.key.decrypted_api_key = "sk-claude".to_string();
|
||||
|
||||
assert_eq!(
|
||||
resolve_local_standard_auth(&transport),
|
||||
Some(("authorization".to_string(), "Bearer sk-claude".to_string(),))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn local_gemini_auth_falls_back_to_default_when_other_format_is_overridden() {
|
||||
let mut transport = sample_transport();
|
||||
transport.endpoint.api_format = "gemini:generate_content".to_string();
|
||||
transport.key.auth_type = "api_key".to_string();
|
||||
transport.key.auth_type_by_format = Some(serde_json::json!({
|
||||
"claude:messages": "bearer"
|
||||
}));
|
||||
transport.key.decrypted_api_key = "sk-gemini".to_string();
|
||||
|
||||
assert_eq!(
|
||||
super::resolve_local_gemini_auth(&transport),
|
||||
Some(("x-goog-api-key".to_string(), "sk-gemini".to_string(),))
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -80,6 +80,8 @@ mod tests {
|
||||
auth_type: "bearer".to_string(),
|
||||
is_active: true,
|
||||
api_formats: None,
|
||||
auth_type_by_format: None,
|
||||
|
||||
allowed_models: None,
|
||||
capabilities: None,
|
||||
rate_multipliers: None,
|
||||
|
||||
@@ -119,6 +119,8 @@ mod tests {
|
||||
auth_type: "bearer".to_string(),
|
||||
is_active: true,
|
||||
api_formats: Some(vec!["claude:messages".to_string()]),
|
||||
auth_type_by_format: None,
|
||||
|
||||
allowed_models: None,
|
||||
capabilities: None,
|
||||
rate_multipliers: None,
|
||||
|
||||
@@ -196,6 +196,8 @@ mod tests {
|
||||
auth_type: "bearer".to_string(),
|
||||
is_active: true,
|
||||
api_formats: Some(vec!["claude:messages".to_string()]),
|
||||
auth_type_by_format: None,
|
||||
|
||||
allowed_models: None,
|
||||
capabilities: None,
|
||||
rate_multipliers: None,
|
||||
|
||||
@@ -136,6 +136,8 @@ mod tests {
|
||||
auth_type: "bearer".to_string(),
|
||||
is_active: true,
|
||||
api_formats: Some(vec!["claude:messages".to_string()]),
|
||||
auth_type_by_format: None,
|
||||
|
||||
allowed_models: None,
|
||||
capabilities: None,
|
||||
rate_multipliers: None,
|
||||
|
||||
@@ -232,6 +232,8 @@ mod tests {
|
||||
auth_type: "bearer".to_string(),
|
||||
is_active: true,
|
||||
api_formats: Some(vec!["claude:messages".to_string()]),
|
||||
auth_type_by_format: None,
|
||||
|
||||
allowed_models: None,
|
||||
capabilities: None,
|
||||
rate_multipliers: None,
|
||||
|
||||
@@ -298,6 +298,8 @@ mod tests {
|
||||
auth_type: "api_key".to_string(),
|
||||
is_active: true,
|
||||
api_formats: None,
|
||||
auth_type_by_format: None,
|
||||
|
||||
allowed_models: None,
|
||||
capabilities: None,
|
||||
rate_multipliers: None,
|
||||
|
||||
@@ -671,6 +671,8 @@ mod tests {
|
||||
auth_type: "bearer".to_string(),
|
||||
is_active: true,
|
||||
api_formats: None,
|
||||
auth_type_by_format: None,
|
||||
|
||||
allowed_models: None,
|
||||
capabilities: None,
|
||||
rate_multipliers: None,
|
||||
|
||||
@@ -306,6 +306,8 @@ mod tests {
|
||||
auth_type: "api_key".to_string(),
|
||||
is_active: true,
|
||||
api_formats: None,
|
||||
auth_type_by_format: None,
|
||||
|
||||
allowed_models: None,
|
||||
capabilities: None,
|
||||
rate_multipliers: None,
|
||||
|
||||
@@ -61,6 +61,7 @@ pub struct GatewayProviderTransportKey {
|
||||
pub auth_type: String,
|
||||
pub is_active: bool,
|
||||
pub api_formats: Option<Vec<String>>,
|
||||
pub auth_type_by_format: Option<serde_json::Value>,
|
||||
pub allowed_models: Option<Vec<String>>,
|
||||
pub capabilities: Option<serde_json::Value>,
|
||||
pub rate_multipliers: Option<serde_json::Value>,
|
||||
@@ -386,6 +387,8 @@ mod tests {
|
||||
"openai:chat".to_string(),
|
||||
"openai:responses".to_string(),
|
||||
]),
|
||||
auth_type_by_format: None,
|
||||
|
||||
allowed_models: Some(vec!["gpt-4.1".to_string(), "gpt-4.1-mini".to_string(),]),
|
||||
capabilities: Some(serde_json::json!({"cache_1h": true})),
|
||||
rate_multipliers: Some(serde_json::json!({"openai:chat": 0.8})),
|
||||
@@ -402,6 +405,31 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn reads_snapshot_when_provider_key_api_key_is_null() {
|
||||
let mut key = sample_key();
|
||||
key.auth_type = "service_account".to_string();
|
||||
key.encrypted_api_key = None;
|
||||
let state = TestSnapshotSource::new(
|
||||
vec![sample_provider()],
|
||||
vec![sample_endpoint()],
|
||||
vec![key],
|
||||
Some(DEVELOPMENT_ENCRYPTION_KEY.to_string()),
|
||||
);
|
||||
|
||||
let snapshot =
|
||||
read_provider_transport_snapshot(&state, "provider-1", "endpoint-1", "key-1")
|
||||
.await
|
||||
.expect("snapshot should read")
|
||||
.expect("snapshot should exist");
|
||||
|
||||
assert_eq!(snapshot.key.decrypted_api_key, "");
|
||||
assert_eq!(
|
||||
snapshot.key.decrypted_auth_config.as_deref(),
|
||||
Some("{\"refresh_token\":\"rt-1\",\"project\":\"demo\"}")
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn returns_none_when_encryption_key_is_not_configured() {
|
||||
let state = TestSnapshotSource::new(
|
||||
|
||||
@@ -54,12 +54,21 @@ pub(super) fn map_key(
|
||||
encryption_key: &str,
|
||||
fallback_encryption_keys: &[String],
|
||||
) -> Result<GatewayProviderTransportKey, DataLayerError> {
|
||||
let decrypted_api_key = decrypt_secret(
|
||||
encryption_key,
|
||||
fallback_encryption_keys,
|
||||
&key.encrypted_api_key,
|
||||
"provider_api_keys.api_key",
|
||||
)?;
|
||||
let decrypted_api_key = key
|
||||
.encrypted_api_key
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(|ciphertext| {
|
||||
decrypt_secret(
|
||||
encryption_key,
|
||||
fallback_encryption_keys,
|
||||
ciphertext,
|
||||
"provider_api_keys.api_key",
|
||||
)
|
||||
})
|
||||
.transpose()?
|
||||
.unwrap_or_default();
|
||||
let decrypted_auth_config = key
|
||||
.encrypted_auth_config
|
||||
.as_deref()
|
||||
@@ -85,6 +94,7 @@ pub(super) fn map_key(
|
||||
normalize_optional_json(key.api_formats),
|
||||
"provider_api_keys.api_formats",
|
||||
)?,
|
||||
auth_type_by_format: normalize_optional_json(key.auth_type_by_format),
|
||||
allowed_models: normalize_string_list(
|
||||
normalize_optional_json(key.allowed_models),
|
||||
"provider_api_keys.allowed_models",
|
||||
|
||||
@@ -88,6 +88,8 @@ mod tests {
|
||||
auth_type: "api_key".to_string(),
|
||||
is_active: true,
|
||||
api_formats: Some(vec!["gemini:generate_content".to_string()]),
|
||||
auth_type_by_format: None,
|
||||
|
||||
allowed_models: None,
|
||||
capabilities: None,
|
||||
rate_multipliers: None,
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
use url::Url;
|
||||
|
||||
use super::super::auth::resolve_local_auth_type_for_transport_format;
|
||||
use super::super::snapshot::GatewayProviderTransportSnapshot;
|
||||
|
||||
const VERTEX_AI_HOST: &str = "aiplatform.googleapis.com";
|
||||
@@ -32,10 +33,7 @@ pub fn is_vertex_api_key_transport_context(transport: &GatewayProviderTransportS
|
||||
.trim()
|
||||
.eq_ignore_ascii_case(super::PROVIDER_TYPE)
|
||||
{
|
||||
return transport
|
||||
.key
|
||||
.auth_type
|
||||
.trim()
|
||||
return resolve_local_auth_type_for_transport_format(transport)
|
||||
.eq_ignore_ascii_case("api_key");
|
||||
}
|
||||
|
||||
@@ -48,11 +46,7 @@ pub fn is_vertex_api_key_transport_context(transport: &GatewayProviderTransportS
|
||||
return false;
|
||||
}
|
||||
|
||||
transport
|
||||
.key
|
||||
.auth_type
|
||||
.trim()
|
||||
.eq_ignore_ascii_case("api_key")
|
||||
resolve_local_auth_type_for_transport_format(transport).eq_ignore_ascii_case("api_key")
|
||||
}
|
||||
|
||||
pub fn uses_vertex_api_key_query_auth(
|
||||
@@ -117,6 +111,8 @@ mod tests {
|
||||
auth_type: "api_key".to_string(),
|
||||
is_active: true,
|
||||
api_formats: Some(vec!["gemini:generate_content".to_string()]),
|
||||
auth_type_by_format: None,
|
||||
|
||||
allowed_models: None,
|
||||
capabilities: None,
|
||||
rate_multipliers: None,
|
||||
|
||||
@@ -203,6 +203,8 @@ mod tests {
|
||||
auth_type: "api_key".to_string(),
|
||||
is_active: true,
|
||||
api_formats: Some(vec!["gemini:generate_content".to_string()]),
|
||||
auth_type_by_format: None,
|
||||
|
||||
allowed_models: None,
|
||||
capabilities: None,
|
||||
rate_multipliers: None,
|
||||
|
||||
@@ -157,6 +157,8 @@ mod tests {
|
||||
auth_type: auth_type.to_string(),
|
||||
is_active: true,
|
||||
api_formats: None,
|
||||
auth_type_by_format: None,
|
||||
|
||||
allowed_models: None,
|
||||
capabilities: None,
|
||||
rate_multipliers: None,
|
||||
|
||||
Reference in New Issue
Block a user