espresso_telemetry/
token.rs

1use std::{
2    path::Path,
3    time::{SystemTime, UNIX_EPOCH},
4};
5
6use anyhow::{Context, anyhow};
7use ark_serialize::{CanonicalDeserialize, CanonicalSerialize};
8use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD as B64};
9use jf_signature::{
10    SignatureScheme,
11    bls_over_bn254::{BLSOverBN254CurveSignatureScheme, SignKey, Signature, VerKey},
12};
13use serde::{Deserialize, Serialize};
14use tagged_base64::TaggedBase64;
15use thiserror::Error;
16
17#[derive(Debug, Error)]
18pub enum TokenParseError {
19    #[error("invalid JWT format, expected header.payload.signature")]
20    InvalidFormat,
21    #[error("unsupported JWT: alg={alg}, typ={typ}")]
22    UnsupportedAlgorithm { alg: String, typ: String },
23    #[error("invalid header: {0}")]
24    InvalidHeader(String),
25    #[error("invalid payload: {0}")]
26    InvalidPayload(String),
27    #[error("invalid signature: {0}")]
28    InvalidSignature(String),
29}
30
31#[derive(Debug, Error)]
32pub enum TokenVerifyError {
33    #[error("BLS signature verification failed")]
34    InvalidSignature,
35    #[error("token timestamp is in the future")]
36    FutureTimestamp,
37    #[error("token expired: age {age}s exceeds max {max_age}s")]
38    Expired { age: u64, max_age: u64 },
39    #[error("system clock error: {0}")]
40    SystemClock(String),
41}
42
43/// JWT header for BLS-BN254 tokens.
44#[derive(Serialize, Deserialize)]
45struct Header {
46    alg: String,
47    typ: String,
48}
49
50/// JWT payload with standard claims.
51#[derive(Serialize, Deserialize)]
52struct Payload {
53    sub: String,
54    iat: u64,
55    #[serde(skip_serializing_if = "Option::is_none", default)]
56    node_name: Option<String>,
57    #[serde(skip_serializing_if = "Option::is_none", default)]
58    company_name: Option<String>,
59}
60
61/// Parsed but unverified JWT token with BLS-BN254 signature.
62pub struct UnauthenticatedToken {
63    pubkey: VerKey,
64    payload: Payload,
65    signing_input: String,
66    signature: Signature,
67}
68
69impl UnauthenticatedToken {
70    pub fn pubkey_str(&self) -> &str {
71        &self.payload.sub
72    }
73
74    pub fn iat(&self) -> u64 {
75        self.payload.iat
76    }
77
78    pub fn node_name(&self) -> Option<&str> {
79        self.payload.node_name.as_deref()
80    }
81
82    pub fn company_name(&self) -> Option<&str> {
83        self.payload.company_name.as_deref()
84    }
85}
86
87/// A token whose BLS signature and timestamp have been verified.
88/// Only constructable via `UnauthenticatedToken::verify`.
89#[derive(Debug)]
90pub struct Token {
91    pubkey_str: String,
92    node_name: Option<String>,
93    company_name: Option<String>,
94}
95
96impl Token {
97    pub fn pubkey_str(&self) -> &str {
98        &self.pubkey_str
99    }
100
101    pub fn node_name(&self) -> Option<&str> {
102        self.node_name.as_deref()
103    }
104
105    pub fn company_name(&self) -> Option<&str> {
106        self.company_name.as_deref()
107    }
108}
109
110impl UnauthenticatedToken {
111    /// Generate a JWT: sign the current timestamp with the BLS key.
112    pub fn generate(signing_key: &SignKey) -> anyhow::Result<Self> {
113        Self::generate_with(signing_key, None, None)
114    }
115
116    /// Generate a JWT with optional `node_name` and `company_name` claims.
117    pub fn generate_with(
118        signing_key: &SignKey,
119        node_name: Option<&str>,
120        company_name: Option<&str>,
121    ) -> anyhow::Result<Self> {
122        let pubkey = VerKey::from(signing_key);
123        let pubkey_str = TaggedBase64::from(&pubkey).to_string();
124        let iat = now_unix_secs()?;
125
126        let header = Header {
127            alg: "BLS-BN254".to_string(),
128            typ: "JWT".to_string(),
129        };
130        let payload = Payload {
131            sub: pubkey_str,
132            iat,
133            node_name: node_name.map(str::to_owned),
134            company_name: company_name.map(str::to_owned),
135        };
136
137        let signing_input = format!(
138            "{}.{}",
139            B64.encode(serde_json::to_vec(&header)?),
140            B64.encode(serde_json::to_vec(&payload)?)
141        );
142
143        let signature = BLSOverBN254CurveSignatureScheme::sign(
144            &(),
145            signing_key,
146            signing_input.as_bytes(),
147            &mut rand::thread_rng(),
148        )?;
149
150        Ok(Self {
151            pubkey,
152            payload,
153            signing_input,
154            signature,
155        })
156    }
157
158    /// Parse a JWT string (header.payload.signature).
159    pub fn parse(s: &str) -> Result<Self, TokenParseError> {
160        let parts: Vec<&str> = s.splitn(3, '.').collect();
161        if parts.len() != 3 {
162            return Err(TokenParseError::InvalidFormat);
163        }
164
165        let header_bytes = B64
166            .decode(parts[0])
167            .map_err(|e| TokenParseError::InvalidHeader(e.to_string()))?;
168        let header: Header = serde_json::from_slice(&header_bytes)
169            .map_err(|e| TokenParseError::InvalidHeader(e.to_string()))?;
170        if header.alg != "BLS-BN254" || header.typ != "JWT" {
171            return Err(TokenParseError::UnsupportedAlgorithm {
172                alg: header.alg,
173                typ: header.typ,
174            });
175        }
176
177        let payload_bytes = B64
178            .decode(parts[1])
179            .map_err(|e| TokenParseError::InvalidPayload(e.to_string()))?;
180        let payload: Payload = serde_json::from_slice(&payload_bytes)
181            .map_err(|e| TokenParseError::InvalidPayload(e.to_string()))?;
182
183        let pubkey: VerKey = TaggedBase64::parse(&payload.sub)
184            .map_err(|e| TokenParseError::InvalidPayload(e.to_string()))?
185            .try_into()
186            .map_err(|_| TokenParseError::InvalidPayload("invalid BLS public key".into()))?;
187
188        let sig_bytes = B64
189            .decode(parts[2])
190            .map_err(|e| TokenParseError::InvalidSignature(e.to_string()))?;
191        let signature = Signature::deserialize_compressed(&sig_bytes[..])
192            .map_err(|e| TokenParseError::InvalidSignature(e.to_string()))?;
193
194        let signing_input = format!("{}.{}", parts[0], parts[1]);
195
196        Ok(Self {
197            pubkey,
198            payload,
199            signing_input,
200            signature,
201        })
202    }
203
204    /// Encode as JWT string (header.payload.signature).
205    pub fn encode(&self) -> String {
206        let mut sig_bytes = Vec::new();
207        self.signature
208            .serialize_compressed(&mut sig_bytes)
209            .expect("signature serialization should not fail");
210
211        format!("{}.{}", self.signing_input, B64.encode(&sig_bytes))
212    }
213
214    /// Verify signature and timestamp, returning an authenticated Token on success.
215    pub fn verify(self, max_age_secs: u64) -> Result<Token, TokenVerifyError> {
216        BLSOverBN254CurveSignatureScheme::verify(
217            &(),
218            &self.pubkey,
219            self.signing_input.as_bytes(),
220            &self.signature,
221        )
222        .map_err(|_| TokenVerifyError::InvalidSignature)?;
223
224        let now = now_unix_secs().map_err(|e| TokenVerifyError::SystemClock(e.to_string()))?;
225        if self.payload.iat > now + 60 {
226            return Err(TokenVerifyError::FutureTimestamp);
227        }
228        let age = now.saturating_sub(self.payload.iat);
229        if age > max_age_secs {
230            return Err(TokenVerifyError::Expired {
231                age,
232                max_age: max_age_secs,
233            });
234        }
235
236        Ok(Token {
237            pubkey_str: self.payload.sub,
238            node_name: self.payload.node_name,
239            company_name: self.payload.company_name,
240        })
241    }
242}
243
244pub fn parse_bls_signing_key(s: &str) -> anyhow::Result<SignKey> {
245    let tb = TaggedBase64::parse(s.trim()).context("invalid tagged-base64")?;
246    let key: SignKey = tb
247        .try_into()
248        .map_err(|_| anyhow!("failed to convert tagged-base64 to SignKey"))?;
249    Ok(key)
250}
251
252pub fn load_bls_signing_key(path: &Path) -> anyhow::Result<SignKey> {
253    let contents = std::fs::read_to_string(path)
254        .with_context(|| format!("failed to read BLS key file: {}", path.display()))?;
255    let key_str = contents
256        .lines()
257        .find(|line| !line.trim().is_empty())
258        .ok_or_else(|| anyhow!("BLS key file is empty"))?;
259    parse_bls_signing_key(key_str)
260}
261
262fn now_unix_secs() -> anyhow::Result<u64> {
263    let d = SystemTime::now()
264        .duration_since(UNIX_EPOCH)
265        .map_err(|e| anyhow::anyhow!("system clock is before UNIX_EPOCH: {e}"))?;
266    Ok(d.as_secs())
267}
268
269#[cfg(test)]
270mod tests {
271    use super::*;
272
273    fn gen_signing_key() -> SignKey {
274        BLSOverBN254CurveSignatureScheme::key_gen(&(), &mut rand::thread_rng())
275            .unwrap()
276            .0
277    }
278
279    #[test]
280    fn roundtrip_encode_parse() {
281        let sk = gen_signing_key();
282        let token = UnauthenticatedToken::generate(&sk).unwrap();
283        let encoded = token.encode();
284
285        assert_eq!(encoded.matches('.').count(), 2);
286
287        let parsed = UnauthenticatedToken::parse(&encoded).unwrap();
288        let authed = parsed.verify(60).unwrap();
289        assert!(!authed.pubkey_str().is_empty());
290    }
291
292    #[test]
293    fn jwt_header_contains_alg() {
294        let sk = gen_signing_key();
295        let token = UnauthenticatedToken::generate(&sk).unwrap();
296        let encoded = token.encode();
297        let header_b64 = encoded.split('.').next().unwrap();
298        let header_json = B64.decode(header_b64).unwrap();
299        let header: serde_json::Value = serde_json::from_slice(&header_json).unwrap();
300        assert_eq!(header["alg"], "BLS-BN254");
301        assert_eq!(header["typ"], "JWT");
302    }
303
304    #[test]
305    fn verify_valid_token() {
306        let sk = gen_signing_key();
307        let token = UnauthenticatedToken::generate(&sk).unwrap();
308        token.verify(60).unwrap();
309    }
310
311    fn sign_jwt(sk: &SignKey, iat: u64) -> String {
312        let pubkey_str = TaggedBase64::from(&VerKey::from(sk)).to_string();
313        let header = Header {
314            alg: "BLS-BN254".into(),
315            typ: "JWT".into(),
316        };
317        let payload = Payload {
318            sub: pubkey_str,
319            iat,
320            node_name: None,
321            company_name: None,
322        };
323        let signing_input = format!(
324            "{}.{}",
325            B64.encode(serde_json::to_vec(&header).unwrap()),
326            B64.encode(serde_json::to_vec(&payload).unwrap())
327        );
328        let sig = BLSOverBN254CurveSignatureScheme::sign(
329            &(),
330            sk,
331            signing_input.as_bytes(),
332            &mut rand::thread_rng(),
333        )
334        .unwrap();
335        let mut sig_bytes = Vec::new();
336        sig.serialize_compressed(&mut sig_bytes).unwrap();
337        format!("{}.{}", signing_input, B64.encode(&sig_bytes))
338    }
339
340    #[test]
341    fn verify_expired_token() {
342        let sk = gen_signing_key();
343        let jwt = sign_jwt(&sk, 1000);
344        let err = UnauthenticatedToken::parse(&jwt)
345            .unwrap()
346            .verify(60)
347            .unwrap_err();
348        assert!(matches!(err, TokenVerifyError::Expired { .. }));
349    }
350
351    #[test]
352    fn verify_future_dated_token() {
353        let sk = gen_signing_key();
354        let jwt = sign_jwt(&sk, now_unix_secs().unwrap() + 10000);
355        let err = UnauthenticatedToken::parse(&jwt)
356            .unwrap()
357            .verify(86400)
358            .unwrap_err();
359        assert!(matches!(err, TokenVerifyError::FutureTimestamp));
360    }
361
362    #[test]
363    fn parse_invalid_format() {
364        assert!(matches!(
365            UnauthenticatedToken::parse("garbage"),
366            Err(TokenParseError::InvalidFormat)
367        ));
368        assert!(matches!(
369            UnauthenticatedToken::parse("only.two"),
370            Err(TokenParseError::InvalidFormat)
371        ));
372    }
373
374    #[test]
375    fn generate_without_claims_payload_is_byte_compat() {
376        // Decode the payload segment and check it has only `sub` + `iat`.
377        // This proves omitted claims aren't serialized, so existing tokens
378        // remain bytewise compatible.
379        let sk = gen_signing_key();
380        let token = UnauthenticatedToken::generate(&sk).unwrap();
381        let encoded = token.encode();
382        let payload_b64 = encoded.split('.').nth(1).unwrap();
383        let payload_bytes = B64.decode(payload_b64).unwrap();
384        let payload: serde_json::Value = serde_json::from_slice(&payload_bytes).unwrap();
385        let obj = payload.as_object().unwrap();
386        assert_eq!(obj.len(), 2, "expected only sub+iat, got {obj:?}");
387        assert!(obj.contains_key("sub"));
388        assert!(obj.contains_key("iat"));
389    }
390
391    #[test]
392    fn generate_with_both_claims_roundtrip() {
393        let sk = gen_signing_key();
394        let token =
395            UnauthenticatedToken::generate_with(&sk, Some("node-01"), Some("acme")).unwrap();
396        let encoded = token.encode();
397
398        let parsed = UnauthenticatedToken::parse(&encoded).unwrap();
399        let authed = parsed.verify(60).unwrap();
400        assert_eq!(authed.node_name(), Some("node-01"));
401        assert_eq!(authed.company_name(), Some("acme"));
402    }
403
404    #[test]
405    fn generate_with_node_name_only() {
406        let sk = gen_signing_key();
407        let token = UnauthenticatedToken::generate_with(&sk, Some("node-42"), None).unwrap();
408        let encoded = token.encode();
409
410        let payload_b64 = encoded.split('.').nth(1).unwrap();
411        let payload_bytes = B64.decode(payload_b64).unwrap();
412        let payload: serde_json::Value = serde_json::from_slice(&payload_bytes).unwrap();
413        let obj = payload.as_object().unwrap();
414        assert!(!obj.contains_key("company_name"));
415
416        let authed = UnauthenticatedToken::parse(&encoded)
417            .unwrap()
418            .verify(60)
419            .unwrap();
420        assert_eq!(authed.node_name(), Some("node-42"));
421        assert_eq!(authed.company_name(), None);
422    }
423}
espresso_telemetry/token.rs

espresso_telemetry/
token.rs
