cookie_agent_models/recipes/
auth.rs

1use std::collections::{BTreeMap, BTreeSet};
2
3use cookie_agent_identity::{AuthFieldName, AuthParameterId};
4use serde::Serialize;
5
6use crate::authoring::{AuthDefinition, AuthOverride, SafeAuthParameterValue, SecretString};
7
8#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)]
9#[serde(rename_all = "snake_case")]
10pub enum CredentialKind {
11    ApiKey,
12    AccessToken,
13    AccessKeyId,
14    SecretAccessKey,
15    SessionToken,
16}
17
18#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)]
19pub struct AuthCredential {
20    pub name: &'static str,
21    pub kind: CredentialKind,
22    pub required: bool,
23}
24
25#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)]
26#[serde(rename_all = "snake_case")]
27pub enum WireAuth {
28    None,
29    BearerApiKey,
30    HeaderApiKey,
31    AnthropicApiKey,
32    GoogleApiKey,
33    OAuthAccessToken,
34    AwsSigV4Static,
35    AzureApiKey,
36}
37
38#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)]
39pub struct AuthMethodRecipe {
40    pub id: &'static str,
41    pub credentials: &'static [AuthCredential],
42    pub wire: WireAuth,
43    pub owned_headers: &'static [&'static str],
44}
45
46const API_KEY: &[AuthCredential] = &[AuthCredential {
47    name: "api_key",
48    kind: CredentialKind::ApiKey,
49    required: true,
50}];
51const ACCESS_TOKEN: &[AuthCredential] = &[AuthCredential {
52    name: "access_token",
53    kind: CredentialKind::AccessToken,
54    required: true,
55}];
56const AWS_STATIC: &[AuthCredential] = &[
57    AuthCredential {
58        name: "access_key_id",
59        kind: CredentialKind::AccessKeyId,
60        required: true,
61    },
62    AuthCredential {
63        name: "secret_access_key",
64        kind: CredentialKind::SecretAccessKey,
65        required: true,
66    },
67    AuthCredential {
68        name: "session_token",
69        kind: CredentialKind::SessionToken,
70        required: false,
71    },
72];
73
74const METHODS: &[AuthMethodRecipe] = &[
75    AuthMethodRecipe {
76        id: "no-auth-v1",
77        credentials: &[],
78        wire: WireAuth::None,
79        owned_headers: &[],
80    },
81    AuthMethodRecipe {
82        id: "bearer-api-key-v1",
83        credentials: API_KEY,
84        wire: WireAuth::BearerApiKey,
85        owned_headers: &["authorization"],
86    },
87    AuthMethodRecipe {
88        id: "api-key-header-v1",
89        credentials: API_KEY,
90        wire: WireAuth::HeaderApiKey,
91        owned_headers: &[],
92    },
93    AuthMethodRecipe {
94        id: "anthropic-api-key-v1",
95        credentials: API_KEY,
96        wire: WireAuth::AnthropicApiKey,
97        owned_headers: &["x-api-key"],
98    },
99    AuthMethodRecipe {
100        id: "google-api-key-header-v1",
101        credentials: API_KEY,
102        wire: WireAuth::GoogleApiKey,
103        owned_headers: &["x-goog-api-key"],
104    },
105    AuthMethodRecipe {
106        id: "oauth-access-token-v1",
107        credentials: ACCESS_TOKEN,
108        wire: WireAuth::OAuthAccessToken,
109        owned_headers: &["authorization"],
110    },
111    AuthMethodRecipe {
112        id: "aws-sigv4-credentials-v1",
113        credentials: AWS_STATIC,
114        wire: WireAuth::AwsSigV4Static,
115        owned_headers: &[
116            "authorization",
117            "host",
118            "x-amz-date",
119            "x-amz-content-sha256",
120            "x-amz-security-token",
121        ],
122    },
123    AuthMethodRecipe {
124        id: "azure-api-key-v1",
125        credentials: API_KEY,
126        wire: WireAuth::AzureApiKey,
127        owned_headers: &["api-key"],
128    },
129];
130
131#[must_use]
132pub fn auth_method(id: &str) -> Option<&'static AuthMethodRecipe> {
133    METHODS.iter().find(|method| method.id == id)
134}
135
136pub fn auth_methods() -> impl ExactSizeIterator<Item = &'static AuthMethodRecipe> {
137    METHODS.iter()
138}
139
140#[derive(Clone, Copy, Debug, Eq, PartialEq, thiserror::Error)]
141pub enum AuthValidationError {
142    #[error("unsupported_auth_method")]
143    UnsupportedAuthMethod,
144    #[error("invalid_auth_shape")]
145    InvalidShape,
146    #[error("ambiguous_api_key")]
147    AmbiguousApiKey,
148}
149
150pub fn validate_auth_override(
151    auth: &AuthOverride,
152    allowed: &[&str],
153) -> Result<&'static AuthMethodRecipe, AuthValidationError> {
154    validate(
155        auth.method.as_str(),
156        &BTreeMap::new(),
157        &auth.values,
158        allowed,
159    )
160}
161
162pub fn validate_auth_definition(
163    auth: &AuthDefinition,
164    allowed: &[&str],
165) -> Result<&'static AuthMethodRecipe, AuthValidationError> {
166    validate(
167        auth.method.as_str(),
168        &auth.parameters,
169        &auth.values,
170        allowed,
171    )
172}
173
174fn validate(
175    id: &str,
176    parameters: &BTreeMap<AuthParameterId, SafeAuthParameterValue>,
177    values: &BTreeMap<AuthFieldName, SecretString>,
178    allowed: &[&str],
179) -> Result<&'static AuthMethodRecipe, AuthValidationError> {
180    if !allowed.contains(&id) {
181        return Err(AuthValidationError::UnsupportedAuthMethod);
182    }
183    let method = auth_method(id).ok_or(AuthValidationError::UnsupportedAuthMethod)?;
184    let actual = values
185        .keys()
186        .map(AuthFieldName::as_str)
187        .collect::<BTreeSet<_>>();
188    let required = method
189        .credentials
190        .iter()
191        .filter(|field| field.required)
192        .map(|field| field.name)
193        .collect::<BTreeSet<_>>();
194    let permitted = method
195        .credentials
196        .iter()
197        .map(|field| field.name)
198        .collect::<BTreeSet<_>>();
199    if !required.is_subset(&actual) || !actual.is_subset(&permitted) {
200        return Err(AuthValidationError::InvalidShape);
201    }
202    if id == "api-key-header-v1" {
203        let header = parameters
204            .iter()
205            .find(|(name, _)| name.as_str() == "header_name")
206            .map(|(_, value)| value.as_str());
207        if parameters.len() != 1 || !matches!(header, Some("x-api-key" | "api-key")) {
208            return Err(AuthValidationError::InvalidShape);
209        }
210    } else if !parameters.is_empty() {
211        return Err(AuthValidationError::InvalidShape);
212    }
213    Ok(method)
214}
cookie_agent_models/recipes/auth.rs

cookie_agent_models/recipes/
auth.rs
