Support per-format provider key auth

This commit is contained in:
fawney19
2026-04-29 15:46:50 +08:00
parent 07a319259b
commit e751289dfb
67 changed files with 1244 additions and 420 deletions

View File

@@ -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(),))
);
}
}

View File

@@ -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,

View File

@@ -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,

View File

@@ -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,

View File

@@ -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,

View File

@@ -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,

View File

@@ -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,

View File

@@ -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,

View File

@@ -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,

View File

@@ -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(

View File

@@ -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",

View File

@@ -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,

View File

@@ -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,

View File

@@ -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,

View File

@@ -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,