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

snap_control/server/
token_verifier.rs
