mentedb_embedding/
aws_credentials.rs

1//! AWS credentials for the hand-rolled SigV4 signers (Bedrock embeddings here,
2//! Bedrock extraction in `mentedb-extraction`).
3//!
4//! Resolution order, checked once at construction:
5//! 1. Static keys: `AWS_ACCESS_KEY_ID` / `AWS_SECRET_ACCESS_KEY` (and
6//!    `AWS_SESSION_TOKEN` for temporary credentials).
7//! 2. The container credentials endpoint used by ECS task roles and EKS Pod
8//!    Identity: `AWS_CONTAINER_CREDENTIALS_FULL_URI` (with
9//!    `AWS_CONTAINER_AUTHORIZATION_TOKEN` or `_TOKEN_FILE` when set), else
10//!    `AWS_CONTAINER_CREDENTIALS_RELATIVE_URI` on `169.254.170.2`.
11//!
12//! Endpoint credentials are temporary, so they are cached and fetched again
13//! shortly before they expire. Profiles and EC2 instance metadata are not
14//! resolved; export keys or run in a container for those.
15
16use std::path::PathBuf;
17use std::sync::{Arc, Mutex};
18use std::time::Duration;
19
20use chrono::{DateTime, Utc};
21
22use mentedb_core::error::{MenteError, MenteResult};
23
24const ECS_ENDPOINT_HOST: &str = "http://169.254.170.2";
25
26/// A resolved set of AWS credentials.
27#[derive(Clone)]
28pub struct AwsCredentials {
29    pub access_key_id: String,
30    pub secret_access_key: String,
31    pub session_token: Option<String>,
32}
33
34impl AwsCredentials {
35    /// Load from `AWS_ACCESS_KEY_ID` / `AWS_SECRET_ACCESS_KEY` /
36    /// `AWS_SESSION_TOKEN` (the last is required for SSO/temporary credentials).
37    pub fn from_env() -> MenteResult<Self> {
38        let access_key_id = std::env::var("AWS_ACCESS_KEY_ID")
39            .map_err(|_| MenteError::Storage("AWS_ACCESS_KEY_ID not set".to_string()))?;
40        let secret_access_key = std::env::var("AWS_SECRET_ACCESS_KEY")
41            .map_err(|_| MenteError::Storage("AWS_SECRET_ACCESS_KEY not set".to_string()))?;
42        Ok(Self {
43            access_key_id,
44            secret_access_key,
45            session_token: std::env::var("AWS_SESSION_TOKEN").ok(),
46        })
47    }
48}
49
50/// Settings for fetching credentials from the container endpoint.
51#[derive(Clone, Debug)]
52pub struct CredentialRefreshConfig {
53    /// Fetch new credentials when the cached ones expire within this window.
54    pub refresh_before_expiry: Duration,
55    /// Timeout for one call to the credentials endpoint (a link-local address).
56    pub endpoint_timeout: Duration,
57}
58
59impl Default for CredentialRefreshConfig {
60    fn default() -> Self {
61        Self {
62            refresh_before_expiry: Duration::from_secs(5 * 60),
63            endpoint_timeout: Duration::from_secs(2),
64        }
65    }
66}
67
68#[derive(Clone)]
69enum ContainerAuth {
70    None,
71    Token(String),
72    TokenFile(PathBuf),
73}
74
75enum Source {
76    Static(AwsCredentials),
77    Container {
78        url: String,
79        auth: ContainerAuth,
80        agent: ureq::Agent,
81        refresh: CredentialRefreshConfig,
82        cache: Mutex<Option<(AwsCredentials, DateTime<Utc>)>>,
83    },
84}
85
86/// Where request signing gets its credentials from. Cheap to clone and safe to
87/// share across threads; endpoint credentials are cached between calls.
88#[derive(Clone)]
89pub struct AwsCredentialSource {
90    inner: Arc<Source>,
91}
92
93#[derive(serde::Deserialize)]
94#[serde(rename_all = "PascalCase")]
95struct EndpointCredentials {
96    access_key_id: String,
97    secret_access_key: String,
98    token: Option<String>,
99    expiration: String,
100}
101
102impl AwsCredentialSource {
103    /// Resolve the source from the environment (see the module docs). Fails
104    /// when neither static keys nor a container endpoint are configured. No
105    /// network call is made here.
106    pub fn from_env() -> MenteResult<Self> {
107        Self::from_env_with(CredentialRefreshConfig::default())
108    }
109
110    /// [`Self::from_env`] with explicit endpoint refresh settings.
111    pub fn from_env_with(refresh: CredentialRefreshConfig) -> MenteResult<Self> {
112        if let Ok(creds) = AwsCredentials::from_env() {
113            return Ok(Self::fixed(creds));
114        }
115        let var = |k: &str| std::env::var(k).ok().filter(|v| !v.is_empty());
116        let auth = if let Some(file) = var("AWS_CONTAINER_AUTHORIZATION_TOKEN_FILE") {
117            ContainerAuth::TokenFile(PathBuf::from(file))
118        } else if let Some(token) = var("AWS_CONTAINER_AUTHORIZATION_TOKEN") {
119            ContainerAuth::Token(token)
120        } else {
121            ContainerAuth::None
122        };
123        if let Some(url) = var("AWS_CONTAINER_CREDENTIALS_FULL_URI") {
124            return Ok(Self::container(url, auth, refresh));
125        }
126        if let Some(path) = var("AWS_CONTAINER_CREDENTIALS_RELATIVE_URI") {
127            return Ok(Self::container(
128                format!("{ECS_ENDPOINT_HOST}{path}"),
129                auth,
130                refresh,
131            ));
132        }
133        Err(MenteError::Storage(
134            "no AWS credentials: set AWS_ACCESS_KEY_ID / AWS_SECRET_ACCESS_KEY, or run with \
135             an ECS task role or EKS Pod Identity (AWS_CONTAINER_CREDENTIALS_RELATIVE_URI / \
136             AWS_CONTAINER_CREDENTIALS_FULL_URI)"
137                .to_string(),
138        ))
139    }
140
141    /// Always sign with these credentials.
142    pub fn fixed(creds: AwsCredentials) -> Self {
143        Self {
144            inner: Arc::new(Source::Static(creds)),
145        }
146    }
147
148    fn container(url: String, auth: ContainerAuth, refresh: CredentialRefreshConfig) -> Self {
149        let agent = ureq::Agent::config_builder()
150            .timeout_global(Some(refresh.endpoint_timeout))
151            .build()
152            .into();
153        Self {
154            inner: Arc::new(Source::Container {
155                url,
156                auth,
157                agent,
158                refresh,
159                cache: Mutex::new(None),
160            }),
161        }
162    }
163
164    /// Credentials to sign the next request with. Endpoint credentials come
165    /// from the cache unless they expire within the refresh window.
166    pub fn credentials(&self) -> MenteResult<AwsCredentials> {
167        match self.inner.as_ref() {
168            Source::Static(creds) => Ok(creds.clone()),
169            Source::Container {
170                url,
171                auth,
172                agent,
173                refresh,
174                cache,
175            } => {
176                let mut cached = cache.lock().unwrap_or_else(|e| e.into_inner());
177                let window = chrono::Duration::from_std(refresh.refresh_before_expiry)
178                    .unwrap_or_else(|_| chrono::Duration::minutes(5));
179                if let Some((creds, expires)) = cached.as_ref()
180                    && *expires - window > Utc::now()
181                {
182                    return Ok(creds.clone());
183                }
184                let fresh = fetch_endpoint(agent, url, auth)?;
185                *cached = Some(fresh.clone());
186                Ok(fresh.0)
187            }
188        }
189    }
190}
191
192fn fetch_endpoint(
193    agent: &ureq::Agent,
194    url: &str,
195    auth: &ContainerAuth,
196) -> MenteResult<(AwsCredentials, DateTime<Utc>)> {
197    let token = match auth {
198        ContainerAuth::None => None,
199        ContainerAuth::Token(t) => Some(t.clone()),
200        ContainerAuth::TokenFile(path) => Some(
201            std::fs::read_to_string(path)
202                .map_err(|e| {
203                    MenteError::Storage(format!(
204                        "reading AWS container authorization token {}: {e}",
205                        path.display()
206                    ))
207                })?
208                .trim()
209                .to_string(),
210        ),
211    };
212    let mut req = agent.get(url).header("accept", "application/json");
213    if let Some(token) = &token {
214        req = req.header("authorization", token);
215    }
216    let mut resp = req.call().map_err(|e| {
217        MenteError::Storage(format!("AWS container credentials request failed: {e}"))
218    })?;
219    let parsed: EndpointCredentials = resp
220        .body_mut()
221        .read_json()
222        .map_err(|e| MenteError::Storage(format!("AWS container credentials parse failed: {e}")))?;
223    let expires = DateTime::parse_from_rfc3339(&parsed.expiration)
224        .map_err(|e| {
225            MenteError::Storage(format!(
226                "AWS container credentials expiration '{}' is not RFC 3339: {e}",
227                parsed.expiration
228            ))
229        })?
230        .with_timezone(&Utc);
231    Ok((
232        AwsCredentials {
233            access_key_id: parsed.access_key_id,
234            secret_access_key: parsed.secret_access_key,
235            session_token: parsed.token,
236        },
237        expires,
238    ))
239}
240
241#[cfg(test)]
242mod tests {
243    use super::*;
244    use std::io::{BufRead, BufReader, Write};
245    use std::net::TcpListener;
246    use std::sync::atomic::{AtomicUsize, Ordering};
247
248    /// A one-route credentials endpoint that counts requests and records the
249    /// authorization header it saw. Each response expires `ttl` from now.
250    fn endpoint(ttl: chrono::Duration) -> (String, Arc<AtomicUsize>, Arc<Mutex<Option<String>>>) {
251        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
252        let addr = listener.local_addr().unwrap();
253        let hits = Arc::new(AtomicUsize::new(0));
254        let seen_auth = Arc::new(Mutex::new(None));
255        let (h, a) = (hits.clone(), seen_auth.clone());
256        std::thread::spawn(move || {
257            for stream in listener.incoming() {
258                let mut stream = stream.unwrap();
259                let mut reader = BufReader::new(stream.try_clone().unwrap());
260                let mut line = String::new();
261                loop {
262                    line.clear();
263                    if reader.read_line(&mut line).unwrap() == 0 || line == "\r\n" {
264                        break;
265                    }
266                    if let Some(v) = line.to_ascii_lowercase().strip_prefix("authorization: ") {
267                        *a.lock().unwrap() = Some(v.trim().to_string());
268                    }
269                }
270                let n = h.fetch_add(1, Ordering::SeqCst) + 1;
271                let body = serde_json::json!({
272                    "AccessKeyId": format!("AKID{n}"),
273                    "SecretAccessKey": "secret",
274                    "Token": "session",
275                    "Expiration": (Utc::now() + ttl).to_rfc3339(),
276                })
277                .to_string();
278                let _ = write!(
279                    stream,
280                    "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}",
281                    body.len()
282                );
283            }
284        });
285        (format!("http://{addr}/creds"), hits, seen_auth)
286    }
287
288    #[test]
289    fn endpoint_credentials_are_cached_until_near_expiry() {
290        let (url, hits, _) = endpoint(chrono::Duration::hours(6));
291        let src = AwsCredentialSource::container(
292            url,
293            ContainerAuth::None,
294            CredentialRefreshConfig::default(),
295        );
296        let first = src.credentials().unwrap();
297        let second = src.credentials().unwrap();
298        assert_eq!(first.access_key_id, "AKID1");
299        assert_eq!(second.access_key_id, "AKID1");
300        assert_eq!(first.session_token.as_deref(), Some("session"));
301        assert_eq!(
302            hits.load(Ordering::SeqCst),
303            1,
304            "second call served from cache"
305        );
306    }
307
308    #[test]
309    fn credentials_inside_the_refresh_window_are_fetched_again() {
310        // Expires in 2 minutes, inside the default 5 minute window.
311        let (url, hits, _) = endpoint(chrono::Duration::minutes(2));
312        let src = AwsCredentialSource::container(
313            url,
314            ContainerAuth::None,
315            CredentialRefreshConfig::default(),
316        );
317        assert_eq!(src.credentials().unwrap().access_key_id, "AKID1");
318        assert_eq!(src.credentials().unwrap().access_key_id, "AKID2");
319        assert_eq!(hits.load(Ordering::SeqCst), 2);
320    }
321
322    #[test]
323    fn authorization_token_is_sent() {
324        let (url, _, seen) = endpoint(chrono::Duration::hours(1));
325        let src = AwsCredentialSource::container(
326            url,
327            ContainerAuth::Token("pod-token".to_string()),
328            CredentialRefreshConfig::default(),
329        );
330        src.credentials().unwrap();
331        assert_eq!(seen.lock().unwrap().as_deref(), Some("pod-token"));
332    }
333
334    #[test]
335    fn unreachable_endpoint_is_an_error_not_a_hang() {
336        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
337        let url = format!("http://{}/creds", listener.local_addr().unwrap());
338        drop(listener);
339        let src = AwsCredentialSource::container(
340            url,
341            ContainerAuth::None,
342            CredentialRefreshConfig::default(),
343        );
344        assert!(src.credentials().is_err());
345    }
346}
mentedb_embedding/aws_credentials.rs

mentedb_embedding/
aws_credentials.rs
