1use std::collections::{HashMap, VecDeque};
26use std::sync::Mutex;
27use std::time::{SystemTime, UNIX_EPOCH};
28
29use base64::engine::general_purpose::URL_SAFE_NO_PAD;
30use base64::Engine as _;
31use hmac::{Hmac, Mac};
32use sha2::Sha256;
33
34use super::{AuthContext, AuthRequest, Authenticate};
35use crate::RpcError;
36
37type HmacSha256 = Hmac<Sha256>;
38
39pub const PROOF_HEADER: &str = "vgi-proxy-proof";
41pub const PROOF_REQUIRED_HEADER: &str = "VGI-Proxy-Proof-Required";
43
44const PROOF_VERSION: &str = "v1";
45const DOMAIN_PREFIX: &[u8] = b"vgi.proxy.proof.v1";
46const DERIVE_LABEL: &[u8] = b"vgi.proxy.proof.v1/";
47const MAX_HEADER_LEN: usize = 512;
48const SECRET_LEN: usize = 32;
49const CLAIMS_PREFIX: &str = "vgi_proxy_proof";
50const DEFAULT_REPLAY_CAPACITY: usize = 100_000;
51
52#[derive(Debug, Clone, Copy, PartialEq, Eq)]
54pub enum ProofMode {
55 Off,
57 Allow,
59 Require,
61}
62
63#[derive(Debug, Clone)]
65pub struct ProofConfig {
66 pub mode: ProofMode,
67 pub origin_id: String,
70 pub secrets: HashMap<String, ([u8; SECRET_LEN], String)>,
72 pub skew_seconds: i64,
74 pub replay_capacity: usize,
76 pub enable_replay_cache: bool,
78}
79
80impl ProofConfig {
81 pub fn new(
83 mode: ProofMode,
84 origin_id: impl Into<String>,
85 secrets: HashMap<String, ([u8; SECRET_LEN], String)>,
86 ) -> Self {
87 Self {
88 mode,
89 origin_id: origin_id.into(),
90 secrets,
91 skew_seconds: 30,
92 replay_capacity: DEFAULT_REPLAY_CAPACITY,
93 enable_replay_cache: true,
94 }
95 }
96
97 pub fn with_skew_seconds(mut self, skew: i64) -> Self {
99 self.skew_seconds = skew;
100 self
101 }
102
103 pub fn without_replay_cache(mut self) -> Self {
105 self.enable_replay_cache = false;
106 self
107 }
108}
109
110#[derive(Debug, Clone, PartialEq, Eq)]
112pub struct ProofError {
113 pub reason: &'static str,
114}
115
116impl ProofError {
117 fn new(reason: &'static str) -> Self {
118 Self { reason }
119 }
120}
121
122fn is_kid(s: &str) -> bool {
126 !s.is_empty()
127 && s.len() <= 64
128 && s.bytes()
129 .all(|b| b.is_ascii_alphanumeric() || b == b'_' || b == b'-')
130}
131
132fn is_ts(s: &str) -> bool {
133 !s.is_empty() && s.len() <= 20 && s.bytes().all(|b| b.is_ascii_digit())
134}
135
136fn is_nonce(s: &str) -> bool {
137 s.len() == 22
138 && s.bytes()
139 .all(|b| b.is_ascii_alphanumeric() || b == b'_' || b == b'-')
140}
141
142fn is_mac(s: &str) -> bool {
143 s.len() == 43
147 && s.bytes()
148 .all(|b| b.is_ascii_alphanumeric() || b == b'_' || b == b'-')
149}
150
151fn is_origin(s: &str) -> bool {
152 !s.is_empty()
153 && s.len() <= 255
154 && s.bytes()
155 .all(|b| b.is_ascii_alphanumeric() || matches!(b, b'.' | b'_' | b':' | b'/' | b'-'))
156}
157
158pub fn derive_proof_secret(
163 base_key: &[u8; SECRET_LEN],
164 proxy_id: &str,
165 origin_id: &str,
166) -> Result<[u8; SECRET_LEN], RpcError> {
167 if !is_origin(proxy_id) || !is_origin(origin_id) {
168 return Err(RpcError::value_error("invalid proxy_id or origin_id"));
169 }
170 let mut msg = Vec::with_capacity(DERIVE_LABEL.len() + proxy_id.len() + origin_id.len() + 1);
171 msg.extend_from_slice(DERIVE_LABEL);
172 msg.extend_from_slice(proxy_id.as_bytes());
173 msg.push(0);
176 msg.extend_from_slice(origin_id.as_bytes());
177
178 let mut mac = HmacSha256::new_from_slice(base_key).expect("hmac accepts any key length");
179 mac.update(&msg);
180 let out = mac.finalize().into_bytes();
181 let mut secret = [0u8; SECRET_LEN];
182 secret.copy_from_slice(&out);
183 Ok(secret)
184}
185
186fn canonical_string(kid: &str, ts: &str, nonce: &str, origin_id: &str) -> Vec<u8> {
191 let mut out = Vec::with_capacity(
192 DOMAIN_PREFIX.len() + kid.len() + ts.len() + nonce.len() + origin_id.len() + 4,
193 );
194 out.extend_from_slice(DOMAIN_PREFIX);
195 for part in [kid, ts, nonce, origin_id] {
196 out.push(0);
197 out.extend_from_slice(part.as_bytes());
198 }
199 out
200}
201
202pub fn mint_proof(
204 secret: &[u8; SECRET_LEN],
205 kid: &str,
206 origin_id: &str,
207 now: i64,
208 nonce: &str,
209) -> Result<String, RpcError> {
210 if !is_kid(kid) || !is_origin(origin_id) {
211 return Err(RpcError::value_error("invalid kid or origin_id"));
212 }
213 let ts = now.to_string();
214 let mut mac = HmacSha256::new_from_slice(secret).expect("hmac accepts any key length");
215 mac.update(&canonical_string(kid, &ts, nonce, origin_id));
216 let sig = URL_SAFE_NO_PAD.encode(mac.finalize().into_bytes());
217 Ok(format!("{PROOF_VERSION}.{kid}.{ts}.{nonce}.{sig}"))
218}
219
220fn unix_now() -> i64 {
221 SystemTime::now()
222 .duration_since(UNIX_EPOCH)
223 .map(|d| d.as_secs() as i64)
224 .unwrap_or(0)
225}
226
227pub fn verify_proof(
232 token: &str,
233 cfg: &ProofConfig,
234 cache: Option<&NonceCache>,
235 now: i64,
236) -> Result<Vec<(String, String)>, ProofError> {
237 if token.len() > MAX_HEADER_LEN {
238 return Err(ProofError::new("malformed"));
239 }
240 let parts: Vec<&str> = token.split('.').collect();
241 if parts.len() != 5 {
242 return Err(ProofError::new("malformed"));
243 }
244 let (version, kid, ts_raw, nonce, mac_b64) = (parts[0], parts[1], parts[2], parts[3], parts[4]);
245 if version != PROOF_VERSION
246 || !is_kid(kid)
247 || !is_ts(ts_raw)
248 || !is_nonce(nonce)
249 || !is_mac(mac_b64)
250 {
251 return Err(ProofError::new("malformed"));
252 }
253
254 let (secret, label) = cfg
255 .secrets
256 .get(kid)
257 .ok_or_else(|| ProofError::new("unknown_kid"))?;
258
259 let ts: i64 = ts_raw.parse().map_err(|_| ProofError::new("malformed"))?;
262 let age = now - ts;
263 if age > cfg.skew_seconds {
264 return Err(ProofError::new("expired"));
265 }
266 if -age > cfg.skew_seconds {
267 return Err(ProofError::new("not_yet_valid"));
268 }
269
270 let received = URL_SAFE_NO_PAD
271 .decode(mac_b64)
272 .map_err(|_| ProofError::new("malformed"))?;
273 let mut mac = HmacSha256::new_from_slice(secret).expect("hmac accepts any key length");
274 mac.update(&canonical_string(kid, ts_raw, nonce, &cfg.origin_id));
275 mac.verify_slice(&received)
278 .map_err(|_| ProofError::new("bad_mac"))?;
279
280 if let Some(cache) = cache {
281 if !cache.check_and_add(nonce, now) {
282 return Err(ProofError::new("replayed"));
283 }
284 }
285
286 Ok(vec![
287 ("verified".into(), "true".into()),
288 ("proxy".into(), label.clone()),
289 ("kid".into(), kid.to_string()),
290 ("origin_id".into(), cfg.origin_id.clone()),
291 ("reason".into(), "ok".into()),
292 ])
293}
294
295#[derive(Debug)]
301pub struct NonceCache {
302 inner: Mutex<NonceCacheInner>,
303 ttl: i64,
304 capacity: usize,
305}
306
307#[derive(Debug)]
308struct NonceCacheInner {
309 order: VecDeque<(String, i64)>,
310 seen: HashMap<String, ()>,
311}
312
313impl NonceCache {
314 pub fn new(ttl: i64, capacity: usize) -> Self {
316 Self {
317 inner: Mutex::new(NonceCacheInner {
318 order: VecDeque::new(),
319 seen: HashMap::new(),
320 }),
321 ttl,
322 capacity,
323 }
324 }
325
326 pub fn check_and_add(&self, nonce: &str, now: i64) -> bool {
331 let mut inner = self.inner.lock().expect("nonce cache poisoned");
332 while let Some((n, expires)) = inner.order.front().cloned() {
335 if expires > now {
336 break;
337 }
338 inner.order.pop_front();
339 inner.seen.remove(&n);
340 }
341 if inner.seen.contains_key(nonce) {
342 return false;
343 }
344 while inner.order.len() >= self.capacity {
348 if let Some((n, _)) = inner.order.pop_front() {
349 inner.seen.remove(&n);
350 }
351 }
352 inner.order.push_back((nonce.to_string(), now + self.ttl));
353 inner.seen.insert(nonce.to_string(), ());
354 true
355 }
356
357 pub fn len(&self) -> usize {
359 self.inner.lock().expect("nonce cache poisoned").seen.len()
360 }
361
362 pub fn is_empty(&self) -> bool {
364 self.len() == 0
365 }
366}
367
368pub fn proof_authenticate(
378 cfg: ProofConfig,
379 inner: Option<Authenticate>,
380) -> Result<Authenticate, RpcError> {
381 if cfg.mode == ProofMode::Off {
382 return Err(RpcError::value_error(
383 "proof_authenticate called with mode=Off; install no gate instead",
384 ));
385 }
386 if !is_origin(&cfg.origin_id) {
387 return Err(RpcError::value_error(
388 "origin_id is required and must be valid",
389 ));
390 }
391 if cfg.secrets.is_empty() {
392 return Err(RpcError::value_error("at least one secret is required"));
393 }
394 for kid in cfg.secrets.keys() {
395 if !is_kid(kid) {
396 return Err(RpcError::value_error("invalid kid in secrets"));
397 }
398 }
399 if cfg.skew_seconds <= 0 {
400 return Err(RpcError::value_error("skew_seconds must be positive"));
401 }
402
403 let cache = cfg
404 .enable_replay_cache
405 .then(|| NonceCache::new(cfg.skew_seconds, cfg.replay_capacity));
406 let required = cfg.mode == ProofMode::Require;
407
408 Ok(std::sync::Arc::new(move |req: &AuthRequest<'_>| {
409 let claims = match verify_request(req, &cfg, cache.as_ref()) {
410 Ok(c) => c,
411 Err(err) => {
412 if required {
413 tracing_reason(err.reason);
417 return Err(RpcError::auth_failure(
421 crate::unauthorized::AuthReason::ProxyRequired,
422 "proxy proof required",
423 ));
424 }
425 vec![
426 ("verified".to_string(), "false".to_string()),
427 ("proxy".to_string(), String::new()),
428 ("kid".to_string(), String::new()),
429 ("origin_id".to_string(), cfg.origin_id.clone()),
430 ("reason".to_string(), err.reason.to_string()),
431 ]
432 }
433 };
434
435 let proxy_label = claims
436 .iter()
437 .find(|(k, _)| k == "proxy")
438 .map(|(_, v)| v.clone())
439 .unwrap_or_default();
440
441 let mut merged: BTreeMapAlias = match &inner {
444 Some(f) => f(req)?.claims,
445 None => Default::default(),
446 };
447 let (domain, authenticated, principal) = match &inner {
448 Some(f) => {
449 let ctx = f(req)?;
450 (ctx.domain, ctx.authenticated, ctx.principal)
451 }
452 None => (CLAIMS_PREFIX.to_string(), true, proxy_label.clone()),
453 };
454 for (k, v) in claims {
455 merged.insert(format!("{CLAIMS_PREFIX}.{k}"), v);
456 }
457 Ok(AuthContext {
458 domain,
459 authenticated,
460 principal,
461 claims: merged,
462 })
463 }))
464}
465
466type BTreeMapAlias = std::collections::BTreeMap<String, String>;
467
468fn tracing_reason(reason: &str) {
469 tracing::warn!(proof_reason = reason, "proxy proof rejected");
470}
471
472fn verify_request(
473 req: &AuthRequest<'_>,
474 cfg: &ProofConfig,
475 cache: Option<&NonceCache>,
476) -> Result<Vec<(String, String)>, ProofError> {
477 let raw = req
478 .header(PROOF_HEADER)
479 .ok_or_else(|| ProofError::new("no_proof"))?;
480 if raw.is_empty() {
481 return Err(ProofError::new("no_proof"));
482 }
483 if raw.contains(',') {
484 return Err(ProofError::new("malformed"));
485 }
486 verify_proof(raw, cfg, cache, unix_now())
487}
488
489pub fn parse_proof_secrets(
495 raw: &str,
496) -> Result<HashMap<String, ([u8; SECRET_LEN], String)>, RpcError> {
497 let mut out = HashMap::new();
498 for chunk in raw.split(',') {
499 let item = chunk.trim();
500 if item.is_empty() {
501 continue;
502 }
503 let (kid, hex_secret) = item
504 .split_once(':')
505 .ok_or_else(|| RpcError::value_error("expected 'kid:hex'"))?;
506 if !is_kid(kid) {
507 return Err(RpcError::value_error("invalid kid"));
508 }
509 if hex_secret.len() != SECRET_LEN * 2 {
510 return Err(RpcError::value_error("secret must be 64 hex chars"));
511 }
512 let mut secret = [0u8; SECRET_LEN];
513 for (i, byte) in secret.iter_mut().enumerate() {
514 *byte = u8::from_str_radix(&hex_secret[i * 2..i * 2 + 2], 16)
515 .map_err(|_| RpcError::value_error("secret is not valid hex"))?;
516 }
517 out.insert(kid.to_string(), (secret, kid.to_string()));
518 }
519 if out.is_empty() {
520 return Err(RpcError::value_error("no secrets parsed"));
521 }
522 Ok(out)
523}
524
525#[cfg(test)]
526mod tests {
527 use super::*;
528
529 const GOLDEN_TOKEN: &str = "v1.conformance-proxy.1700000000.Q0ZPUk1BTkNFTk9OQ0UxMQ.XQ2QBf35oajjaP7HIas3OfyEvNhyXTTptbrxWFxWk3I";
534 const GOLDEN_ORIGIN: &str = "conformance-origin";
535 const GOLDEN_KID: &str = "conformance-proxy";
536 const GOLDEN_TIME: i64 = 1700000000;
537 const GOLDEN_DERIVED: &str = "af85db125b8270bc0a0971736340dc8476ba70e1fad472b72b68ba739bd1cd94";
538
539 fn secret() -> [u8; 32] {
540 [0x11u8; 32]
541 }
542
543 fn config() -> ProofConfig {
544 let mut secrets = HashMap::new();
545 secrets.insert(GOLDEN_KID.to_string(), (secret(), GOLDEN_KID.to_string()));
546 ProofConfig::new(ProofMode::Require, GOLDEN_ORIGIN, secrets)
547 }
548
549 fn to_hex(bytes: &[u8]) -> String {
550 bytes.iter().map(|b| format!("{b:02x}")).collect()
551 }
552
553 #[test]
554 fn verifies_python_minted_token() {
555 let claims = verify_proof(GOLDEN_TOKEN, &config(), None, GOLDEN_TIME)
556 .expect("cross-language token must verify");
557 assert!(claims.contains(&("proxy".to_string(), GOLDEN_KID.to_string())));
558 }
559
560 #[test]
561 fn mint_matches_python() {
562 let token = mint_proof(
563 &secret(),
564 GOLDEN_KID,
565 GOLDEN_ORIGIN,
566 GOLDEN_TIME,
567 "Q0ZPUk1BTkNFTk9OQ0UxMQ",
568 )
569 .unwrap();
570 assert_eq!(token, GOLDEN_TOKEN, "Rust mint diverged from Python");
571 }
572
573 #[test]
574 fn derivation_matches_python() {
575 let mut base = [0u8; 32];
576 for (i, b) in base.iter_mut().enumerate() {
577 *b = i as u8;
578 }
579 let got = derive_proof_secret(&base, "prod-use1", "worker-a").unwrap();
580 assert_eq!(to_hex(&got), GOLDEN_DERIVED);
581 }
582
583 #[test]
584 fn derivation_separator_is_unambiguous() {
585 let base = [0u8; 32];
586 let a = derive_proof_secret(&base, "ab", "c.d").unwrap();
587 let b = derive_proof_secret(&base, "a", "b.c.d").unwrap();
588 assert_ne!(a, b, "component boundaries can be shifted");
589 }
590
591 #[test]
592 fn malformed_tokens_rejected() {
593 let cfg = config();
594 for token in [
595 "",
596 "garbage",
597 "v1.a.b.c",
598 "v1.a.b.c.d.e",
599 "v2.conformance-proxy.1.Q0ZPUk1BTkNFTk9OQ0UxMQ.AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA",
600 "v1.bad!kid.1.Q0ZPUk1BTkNFTk9OQ0UxMQ.AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA",
601 "v1.conformance-proxy.xyz.Q0ZPUk1BTkNFTk9OQ0UxMQ.AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA",
602 "v1.conformance-proxy.1.short.AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA",
603 "v1.conformance-proxy.1.Q0ZPUk1BTkNFTk9OQ0UxMQ.!!!",
604 ] {
605 let err = verify_proof(token, &cfg, None, GOLDEN_TIME).unwrap_err();
606 assert_eq!(err.reason, "malformed", "token {token:?}");
607 }
608 }
609
610 #[test]
611 fn unknown_kid_rejected() {
612 let mut cfg = config();
613 cfg.secrets = HashMap::new();
614 cfg.secrets
615 .insert("other".to_string(), (secret(), "other".to_string()));
616 assert_eq!(
617 verify_proof(GOLDEN_TOKEN, &cfg, None, GOLDEN_TIME)
618 .unwrap_err()
619 .reason,
620 "unknown_kid"
621 );
622 }
623
624 #[test]
625 fn wrong_origin_rejected() {
626 let mut cfg = config();
628 cfg.origin_id = "some-other-worker".to_string();
629 assert_eq!(
630 verify_proof(GOLDEN_TOKEN, &cfg, None, GOLDEN_TIME)
631 .unwrap_err()
632 .reason,
633 "bad_mac"
634 );
635 }
636
637 #[test]
638 fn time_window_is_two_sided() {
639 let cfg = config();
640 assert_eq!(
643 verify_proof(GOLDEN_TOKEN, &cfg, None, GOLDEN_TIME + 91)
644 .unwrap_err()
645 .reason,
646 "expired"
647 );
648 assert_eq!(
649 verify_proof(GOLDEN_TOKEN, &cfg, None, GOLDEN_TIME - 91)
650 .unwrap_err()
651 .reason,
652 "not_yet_valid"
653 );
654 assert!(verify_proof(GOLDEN_TOKEN, &cfg, None, GOLDEN_TIME + 20).is_ok());
655 }
656
657 #[test]
658 fn mac_framing_must_be_separated() {
659 let mut bad = Vec::from(DOMAIN_PREFIX);
662 bad.extend_from_slice(GOLDEN_KID.as_bytes());
663 bad.extend_from_slice(b"1700000000");
664 bad.extend_from_slice(b"Q0ZPUk1BTkNFTk9OQ0UxMQ");
665 bad.extend_from_slice(GOLDEN_ORIGIN.as_bytes());
666 let mut mac = HmacSha256::new_from_slice(&secret()).unwrap();
667 mac.update(&bad);
668 let sig = URL_SAFE_NO_PAD.encode(mac.finalize().into_bytes());
669 let token = format!("v1.{GOLDEN_KID}.1700000000.Q0ZPUk1BTkNFTk9OQ0UxMQ.{sig}");
670 assert_eq!(
671 verify_proof(&token, &config(), None, GOLDEN_TIME)
672 .unwrap_err()
673 .reason,
674 "bad_mac"
675 );
676 }
677
678 #[test]
679 fn replay_rejected() {
680 let cache = NonceCache::new(30, 100);
681 let cfg = config();
682 assert!(verify_proof(GOLDEN_TOKEN, &cfg, Some(&cache), GOLDEN_TIME).is_ok());
683 assert_eq!(
684 verify_proof(GOLDEN_TOKEN, &cfg, Some(&cache), GOLDEN_TIME)
685 .unwrap_err()
686 .reason,
687 "replayed"
688 );
689 }
690
691 #[test]
692 fn nonce_cache_capacity_is_hard() {
693 let cache = NonceCache::new(3600, 10);
696 for i in 0..500 {
697 cache.check_and_add(&format!("nonce-{i}"), GOLDEN_TIME);
698 }
699 assert!(
700 cache.len() <= 10,
701 "capacity cap not enforced: {}",
702 cache.len()
703 );
704 }
705
706 #[test]
707 fn nonce_cache_expires() {
708 let cache = NonceCache::new(30, 100);
709 assert!(cache.check_and_add("n1", 1000));
710 assert!(!cache.check_and_add("n1", 1000));
711 assert!(
712 cache.check_and_add("n1", 1031),
713 "entry should expire past the TTL"
714 );
715 }
716
717 #[test]
718 fn off_mode_refuses_to_build() {
719 let mut cfg = config();
720 cfg.mode = ProofMode::Off;
721 assert!(proof_authenticate(cfg, None).is_err());
722 }
723
724 #[test]
725 fn parse_secrets_round_trip() {
726 let parsed = parse_proof_secrets(&format!("prod-use1:{}", "11".repeat(32))).unwrap();
727 assert_eq!(parsed["prod-use1"].1, "prod-use1");
728 for bad in ["prod-use1", "prod-use1:zz", "bad!kid:11", ""] {
729 assert!(parse_proof_secrets(bad).is_err(), "accepted {bad:?}");
730 }
731 }
732}