modelexpress_client/
auth.rs

1// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Attaches a Kubernetes projected ServiceAccount token as `authorization: Bearer` on
5//! every RPC, re-reading on TTL/mtime change. A no-op when no token file is present.
6
7use std::path::PathBuf;
8use std::sync::{Arc, Mutex};
9use std::time::{Duration, Instant, SystemTime};
10
11use tonic::metadata::MetadataValue;
12use tonic::service::Interceptor;
13use tonic::{Request, Status};
14use tracing::warn;
15
16const DEFAULT_TOKEN_PATH: &str = "/var/run/secrets/tokens/modelexpress";
17const DEFAULT_TTL: Duration = Duration::from_secs(60);
18
19#[derive(Default)]
20struct CachedToken {
21    token: Option<String>,
22    read_at: Option<Instant>,
23    mtime: Option<SystemTime>,
24    warned_missing: bool,
25}
26
27pub struct TokenProvider {
28    path: PathBuf,
29    ttl: Duration,
30    cache: Mutex<CachedToken>,
31}
32
33impl TokenProvider {
34    #[must_use]
35    pub fn from_env() -> Self {
36        let path = std::env::var(modelexpress_common::envs::MX_AUTH_TOKEN_PATH)
37            .map(PathBuf::from)
38            .unwrap_or_else(|_| PathBuf::from(DEFAULT_TOKEN_PATH));
39        let ttl = std::env::var(modelexpress_common::envs::MX_AUTH_TOKEN_TTL_SECONDS)
40            .ok()
41            .and_then(|raw| raw.parse::<u64>().ok())
42            .map_or(DEFAULT_TTL, Duration::from_secs);
43        Self {
44            path,
45            ttl,
46            cache: Mutex::new(CachedToken::default()),
47        }
48    }
49
50    fn token(&self) -> Option<String> {
51        let now = Instant::now();
52        let mut cache = self
53            .cache
54            .lock()
55            .unwrap_or_else(std::sync::PoisonError::into_inner);
56
57        if let Some(read_at) = cache.read_at
58            && now.duration_since(read_at) < self.ttl
59        {
60            return cache.token.clone();
61        }
62
63        match std::fs::metadata(&self.path).and_then(|meta| meta.modified()) {
64            Ok(mtime) => {
65                if cache.token.is_some() && cache.mtime == Some(mtime) {
66                    cache.read_at = Some(now);
67                    return cache.token.clone();
68                }
69                match std::fs::read_to_string(&self.path) {
70                    Ok(contents) => {
71                        let token = contents.trim();
72                        cache.token = (!token.is_empty()).then(|| token.to_string());
73                        cache.read_at = Some(now);
74                        cache.mtime = Some(mtime);
75                        cache.warned_missing = false;
76                        cache.token.clone()
77                    }
78                    Err(e) => {
79                        Self::warn_once(&mut cache, &format!("failed to read token: {e}"));
80                        None
81                    }
82                }
83            }
84            Err(_) => {
85                Self::warn_once(
86                    &mut cache,
87                    "token file not found; RPCs sent without a bearer token \
88                     (server must have auth off)",
89                );
90                cache.token = None;
91                cache.mtime = None;
92                None
93            }
94        }
95    }
96
97    fn warn_once(cache: &mut CachedToken, message: &str) {
98        if !cache.warned_missing {
99            warn!("modelexpress auth: {message}");
100            cache.warned_missing = true;
101        }
102    }
103}
104
105#[derive(Clone)]
106pub struct AuthInterceptor {
107    provider: Arc<TokenProvider>,
108}
109
110impl AuthInterceptor {
111    #[must_use]
112    pub fn new(provider: Arc<TokenProvider>) -> Self {
113        Self { provider }
114    }
115}
116
117impl Interceptor for AuthInterceptor {
118    fn call(&mut self, mut request: Request<()>) -> Result<Request<()>, Status> {
119        if let Some(token) = self.provider.token() {
120            let value: MetadataValue<_> = format!("Bearer {token}")
121                .parse()
122                .map_err(|_| Status::internal("auth token is not a valid header value"))?;
123            request.metadata_mut().insert("authorization", value);
124        }
125        Ok(request)
126    }
127}
128
129#[cfg(test)]
130#[allow(clippy::expect_used)]
131mod tests {
132    use super::*;
133    use std::fs;
134    use std::io::Write;
135
136    #[test]
137    fn reads_token_from_file() {
138        let mut file = tempfile::NamedTempFile::new().expect("temp file");
139        writeln!(file, "  my-token").expect("write");
140        let provider = TokenProvider {
141            path: file.path().to_path_buf(),
142            ttl: DEFAULT_TTL,
143            cache: Mutex::new(CachedToken::default()),
144        };
145        assert_eq!(provider.token().as_deref(), Some("my-token"));
146    }
147
148    #[test]
149    fn missing_file_is_none() {
150        let provider = TokenProvider {
151            path: PathBuf::from("/no/such/token/file"),
152            ttl: DEFAULT_TTL,
153            cache: Mutex::new(CachedToken::default()),
154        };
155        assert!(provider.token().is_none());
156    }
157
158    #[test]
159    fn empty_file_is_none() {
160        let file = tempfile::NamedTempFile::new().expect("temp file");
161        let provider = TokenProvider {
162            path: file.path().to_path_buf(),
163            ttl: DEFAULT_TTL,
164            cache: Mutex::new(CachedToken::default()),
165        };
166        assert!(provider.token().is_none());
167    }
168
169    #[test]
170    fn cached_token_is_returned_until_ttl_expires() {
171        let mut file = tempfile::NamedTempFile::new().expect("temp file");
172        writeln!(file, "old-token").expect("write old token");
173        let provider = TokenProvider {
174            path: file.path().to_path_buf(),
175            ttl: DEFAULT_TTL,
176            cache: Mutex::new(CachedToken::default()),
177        };
178
179        assert_eq!(provider.token().as_deref(), Some("old-token"));
180        fs::write(file.path(), "new-token\n").expect("write new token");
181
182        assert_eq!(provider.token().as_deref(), Some("old-token"));
183    }
184
185    #[test]
186    fn cached_token_reuses_mtime_when_ttl_expires_without_file_change() {
187        let mut file = tempfile::NamedTempFile::new().expect("temp file");
188        writeln!(file, "same-token").expect("write token");
189        let provider = TokenProvider {
190            path: file.path().to_path_buf(),
191            ttl: Duration::ZERO,
192            cache: Mutex::new(CachedToken::default()),
193        };
194
195        assert_eq!(provider.token().as_deref(), Some("same-token"));
196        let first_read_at = provider
197            .cache
198            .lock()
199            .expect("cache lock")
200            .read_at
201            .expect("read timestamp");
202
203        assert_eq!(provider.token().as_deref(), Some("same-token"));
204        let second_read_at = provider
205            .cache
206            .lock()
207            .expect("cache lock")
208            .read_at
209            .expect("read timestamp");
210        assert!(second_read_at >= first_read_at);
211    }
212
213    #[test]
214    fn token_is_reread_after_ttl_when_mtime_changes() {
215        let mut file = tempfile::NamedTempFile::new().expect("temp file");
216        writeln!(file, "old-token").expect("write old token");
217        let provider = TokenProvider {
218            path: file.path().to_path_buf(),
219            ttl: Duration::ZERO,
220            cache: Mutex::new(CachedToken::default()),
221        };
222
223        assert_eq!(provider.token().as_deref(), Some("old-token"));
224        fs::write(file.path(), "new-token\n").expect("write new token");
225        provider.cache.lock().expect("cache lock").mtime = Some(SystemTime::UNIX_EPOCH);
226
227        assert_eq!(provider.token().as_deref(), Some("new-token"));
228    }
229
230    #[test]
231    fn interceptor_attaches_authorization_metadata() {
232        let mut file = tempfile::NamedTempFile::new().expect("temp file");
233        writeln!(file, "rpc-token").expect("write token");
234        let provider = Arc::new(TokenProvider {
235            path: file.path().to_path_buf(),
236            ttl: DEFAULT_TTL,
237            cache: Mutex::new(CachedToken::default()),
238        });
239        let mut interceptor = AuthInterceptor::new(provider);
240
241        let request = interceptor
242            .call(Request::new(()))
243            .expect("interceptor request");
244
245        assert_eq!(
246            request
247                .metadata()
248                .get("authorization")
249                .and_then(|value| value.to_str().ok()),
250            Some("Bearer rpc-token")
251        );
252    }
253
254    #[test]
255    fn interceptor_rejects_invalid_header_value() {
256        let mut file = tempfile::NamedTempFile::new().expect("temp file");
257        writeln!(file, "bad\nvalue").expect("write token");
258        let provider = Arc::new(TokenProvider {
259            path: file.path().to_path_buf(),
260            ttl: DEFAULT_TTL,
261            cache: Mutex::new(CachedToken::default()),
262        });
263        let mut interceptor = AuthInterceptor::new(provider);
264
265        let status = interceptor
266            .call(Request::new(()))
267            .expect_err("invalid metadata value");
268
269        assert_eq!(status.code(), tonic::Code::Internal);
270    }
271
272    #[test]
273    fn missing_file_is_rechecked_immediately() {
274        let dir = tempfile::tempdir().expect("temp dir");
275        let path = dir.path().join("late-mounted-token");
276        let provider = TokenProvider {
277            path: path.clone(),
278            ttl: DEFAULT_TTL,
279            cache: Mutex::new(CachedToken::default()),
280        };
281
282        assert!(provider.token().is_none());
283
284        fs::write(&path, "appeared-token\n").expect("write token");
285
286        assert_eq!(provider.token().as_deref(), Some("appeared-token"));
287    }
288}
modelexpress_client/auth.rs

modelexpress_client/
auth.rs
