Skip to main content

snap_control/server/
token_verifier.rs

1// Copyright 2026 Anapaya Systems
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//   http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14//! SNAP token verifier supporting both static keys and JWKS-based key resolution.
15
16use std::sync::Arc;
17
18use jsonwebtoken::{Algorithm, DecodingKey, Validation, decode, decode_header};
19use scion_sdk_token_validator::validator::Token;
20use snap_tokens::AnyClaims;
21use thiserror::Error;
22
23use crate::server::jwks_key_store::JwksKeyStore;
24
25/// Errors returned by [`SnapTokenVerifier::verify`].
26#[derive(Debug, Error)]
27pub enum SnapTokenVerifyError {
28    /// Could not decode the JWT header.
29    #[error("failed to decode JWT header: {0}")]
30    HeaderDecodeError(jsonwebtoken::errors::Error),
31    /// Token carries a `kid` that could not be resolved from the JWKS store.
32    #[error("JWKS key not found for kid '{0}'")]
33    UnknownKid(String),
34    /// JWT signature/claims validation failed.
35    #[error("token verification failed: {0}")]
36    VerificationFailed(jsonwebtoken::errors::Error),
37}
38
39/// Verifies SNAP tokens against either a statically configured key or keys fetched
40/// from a JWKS endpoint.
41///
42/// - Tokens **without** a `kid` JWT header claim are verified using the static key.
43/// - Tokens **with** a `kid` are verified using a key resolved from the `JwksKeyStore`. If no JWKS
44///   store is configured, the static key is used as a fallback.
45#[derive(Clone)]
46pub struct SnapTokenVerifier {
47    static_key: DecodingKey,
48    jwks_store: Option<Arc<JwksKeyStore>>,
49    validation: Validation,
50}
51
52impl SnapTokenVerifier {
53    /// Creates a verifier that only uses the statically configured key.
54    /// Tokens with a `kid` header claim also use this key as a fallback when no JWKS
55    /// store is configured.
56    pub fn new(static_key: DecodingKey) -> Self {
57        Self {
58            static_key,
59            jwks_store: None,
60            validation: build_validation(),
61        }
62    }
63
64    /// Attaches a JWKS store for resolving `kid`-bearing tokens.
65    pub fn with_jwks_store(mut self, store: Arc<JwksKeyStore>) -> Self {
66        self.jwks_store = Some(store);
67        self
68    }
69
70    /// Verifies a SNAP token JWT and returns the parsed claims on success.
71    ///
72    /// # Key selection
73    ///
74    /// - If the JWT header has a `kid` and a JWKS store is configured, the key is resolved from the
75    ///   JWKS store.
76    /// - Otherwise (no `kid`, or `kid` present but no JWKS store configured), the static key is
77    ///   used.
78    pub async fn verify(&self, token: &str) -> Result<AnyClaims, SnapTokenVerifyError> {
79        let header = decode_header(token).map_err(SnapTokenVerifyError::HeaderDecodeError)?;
80
81        let key = match (header.kid, &self.jwks_store) {
82            (Some(kid), Some(store)) => {
83                match store.await_key(&kid).await {
84                    Some(k) => k,
85                    None => return Err(SnapTokenVerifyError::UnknownKid(kid)),
86                }
87            }
88            _ => self.static_key.clone(),
89        };
90
91        let token_data = decode::<AnyClaims>(token, &key, &self.validation)
92            .map_err(SnapTokenVerifyError::VerificationFailed)?;
93
94        Ok(token_data.claims)
95    }
96}
97
98fn build_validation() -> Validation {
99    let mut v = Validation::new(Algorithm::EdDSA);
100    v.set_required_spec_claims(&AnyClaims::required_claims());
101    v.set_audience(&["snap"]);
102    v
103}
104
105#[cfg(test)]
106mod tests {
107    use std::time::{Duration, SystemTime, UNIX_EPOCH};
108
109    use base64::Engine;
110    use ed25519_dalek::pkcs8::EncodePrivateKey;
111    use jsonwebtoken::{Algorithm, EncodingKey, Header};
112    use scion_sdk_token_validator::validator::insecure_const_ed25519_signing_key;
113    use snap_tokens::v0::{self, SnapTokenClaims, insecure_const_snap_token_key_pair};
114    use tokio_util::sync::CancellationToken;
115
116    use super::*;
117
118    // --- helpers ---
119
120    fn static_verifier() -> SnapTokenVerifier {
121        let (_, decoding_key) = insecure_const_snap_token_key_pair();
122        SnapTokenVerifier::new(decoding_key)
123    }
124
125    fn v0_token() -> String {
126        v0::dummy_snap_token()
127    }
128
129    fn expired_v0_token() -> String {
130        let (encoding_key, _) = insecure_const_snap_token_key_pair();
131        let claims = SnapTokenClaims {
132            pssid: v0::Pssid::new(),
133            exp: 1, // far in the past
134            jti: "test".to_string(),
135        };
136        jsonwebtoken::encode(&Header::new(Algorithm::EdDSA), &claims, &encoding_key).unwrap()
137    }
138
139    fn v1_token_with_kid_and_static_key(kid: &str) -> String {
140        let (encoding_key, _) = insecure_const_snap_token_key_pair();
141        let now = SystemTime::now()
142            .duration_since(UNIX_EPOCH)
143            .unwrap()
144            .as_secs();
145        let claims = serde_json::json!({
146            "ver": 1,
147            "iss": "ssr",
148            "aud": "snap",
149            "exp": now + 3600,
150            "nbf": now,
151            "iat": now,
152            "jti": "test-jti",
153            "pssid": "AAAAAAAAAAAAAAAAAAAAAAA",
154        });
155        let mut header = Header::new(Algorithm::EdDSA);
156        header.kid = Some(kid.to_string());
157        jsonwebtoken::encode(&header, &claims, &encoding_key).unwrap()
158    }
159
160    fn v1_token_with_kid(kid: &str) -> String {
161        let signing_key = insecure_const_ed25519_signing_key();
162        let der = signing_key.to_pkcs8_der().unwrap();
163        let encoding_key = EncodingKey::from_ed_der(der.as_bytes());
164
165        let now = SystemTime::now()
166            .duration_since(UNIX_EPOCH)
167            .unwrap()
168            .as_secs();
169
170        // Build a minimal v1 claims payload.
171        // pssid = base64url-no-pad of [0x00; 17] = 23 'A' chars.
172        let claims = serde_json::json!({
173            "ver": 1,
174            "iss": "ssr",
175            "aud": "snap",
176            "exp": now + 3600,
177            "nbf": now,
178            "iat": now,
179            "jti": "test-jti",
180            "pssid": "AAAAAAAAAAAAAAAAAAAAAAA",
181        });
182
183        let mut header = Header::new(Algorithm::EdDSA);
184        header.kid = Some(kid.to_string());
185
186        jsonwebtoken::encode(&header, &claims, &encoding_key).unwrap()
187    }
188
189    fn jwks_json_for_kid(kid: &str) -> serde_json::Value {
190        let signing_key = insecure_const_ed25519_signing_key();
191        let x = base64::engine::general_purpose::URL_SAFE_NO_PAD
192            .encode(signing_key.verifying_key().as_bytes());
193        serde_json::json!({
194            "keys": [{
195                "kid": kid,
196                "kty": "OKP",
197                "use": "sig",
198                "alg": "EdDSA",
199                "crv": "Ed25519",
200                "x": x
201            }]
202        })
203    }
204
205    async fn verifier_with_jwks_server(kid: &str) -> SnapTokenVerifier {
206        use axum::{Json, Router, routing::get};
207
208        let jwks = jwks_json_for_kid(kid);
209        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
210        let addr = listener.local_addr().unwrap();
211        let app = Router::new().route(
212            "/.well-known/jwks.json",
213            get(move || {
214                let jwks = jwks.clone();
215                async move { Json(jwks) }
216            }),
217        );
218        tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
219
220        let url = format!("http://{}/.well-known/jwks.json", addr)
221            .parse()
222            .unwrap();
223        let (_, static_key) = insecure_const_snap_token_key_pair();
224        let store = JwksKeyStore::new(url, Duration::from_secs(3600), CancellationToken::new());
225        SnapTokenVerifier::new(static_key).with_jwks_store(Arc::new(store))
226    }
227
228    // --- tests ---
229
230    #[tokio::test]
231    async fn no_kid_uses_static_key_and_succeeds() {
232        let verifier = static_verifier();
233        let token = v0_token();
234        let result = verifier.verify(&token).await;
235        assert!(
236            result.is_ok(),
237            "valid V0 token should be accepted: {result:?}"
238        );
239    }
240
241    #[tokio::test]
242    async fn no_kid_invalid_signature_rejected() {
243        let verifier = static_verifier();
244        // Use a different key to sign the token (wrong key)
245        let different_key = {
246            let seed = [99u8; 32];
247            let sk = ed25519_dalek::SigningKey::from_bytes(&seed);
248            let der = sk.to_pkcs8_der().unwrap();
249            EncodingKey::from_ed_der(der.as_bytes())
250        };
251        let claims = SnapTokenClaims {
252            pssid: v0::Pssid::new(),
253            exp: (SystemTime::now() + Duration::from_secs(3600))
254                .duration_since(UNIX_EPOCH)
255                .unwrap()
256                .as_secs(),
257            jti: "test".to_string(),
258        };
259        let bad_token =
260            jsonwebtoken::encode(&Header::new(Algorithm::EdDSA), &claims, &different_key).unwrap();
261        let result = verifier.verify(&bad_token).await;
262        assert!(
263            result.is_err(),
264            "token signed with wrong key should be rejected"
265        );
266    }
267
268    #[tokio::test]
269    async fn no_kid_expired_token_rejected() {
270        let verifier = static_verifier();
271        let token = expired_v0_token();
272        let result = verifier.verify(&token).await;
273        assert!(result.is_err(), "expired token should be rejected");
274    }
275
276    #[tokio::test]
277    async fn kid_with_no_jwks_url_falls_back_to_static_key() {
278        let verifier = static_verifier(); // no JWKS store
279        let token = v1_token_with_kid_and_static_key("some-kid");
280        let result = verifier.verify(&token).await;
281        assert!(
282            result.is_ok(),
283            "token with kid but no JWKS URL should fall back to static key: {result:?}"
284        );
285    }
286
287    #[tokio::test]
288    async fn kid_resolved_via_jwks_succeeds() {
289        let kid = "ssr-key-1";
290        let verifier = verifier_with_jwks_server(kid).await;
291        let token = v1_token_with_kid(kid);
292        let result = verifier.verify(&token).await;
293        assert!(
294            result.is_ok(),
295            "V1 token with JWKS-resolved key should succeed: {result:?}"
296        );
297    }
298
299    #[tokio::test]
300    async fn unknown_kid_rejected() {
301        let verifier = verifier_with_jwks_server("other-kid").await;
302        let token = v1_token_with_kid("unknown-kid");
303        let result = verifier.verify(&token).await;
304        assert!(result.is_err(), "token with unknown kid should be rejected");
305    }
306}