Skip to main content

snowflake_jwt/
lib.rs

1#![doc(
2    issue_tracker_base_url = "https://github.com/mycelial/snowflake-rs/issues",
3    test(no_crate_inject)
4)]
5#![doc = include_str ! ("../README.md")]
6
7use base64::Engine;
8use jsonwebtoken::{encode, Algorithm, EncodingKey, Header};
9use rsa::pkcs1::EncodeRsaPrivateKey;
10use rsa::pkcs8::{DecodePrivateKey, EncodePublicKey};
11use serde::{Deserialize, Serialize};
12use sha2::{Digest, Sha256};
13use thiserror::Error;
14use time::{Duration, OffsetDateTime};
15
16#[derive(Error, Debug)]
17pub enum JwtError {
18    #[error(transparent)]
19    Rsa(#[from] rsa::Error),
20
21    #[error(transparent)]
22    Pkcs8(#[from] rsa::pkcs8::Error),
23
24    #[error(transparent)]
25    Spki(#[from] rsa::pkcs8::spki::Error),
26
27    #[error(transparent)]
28    Pkcs1(#[from] rsa::pkcs1::Error),
29
30    #[error(transparent)]
31    Utf8(#[from] std::string::FromUtf8Error),
32
33    #[error(transparent)]
34    Der(#[from] rsa::pkcs1::der::Error),
35
36    #[error(transparent)]
37    JwtEncoding(#[from] jsonwebtoken::errors::Error),
38}
39
40#[derive(Debug, Serialize, Deserialize)]
41struct Claims {
42    iss: String,
43    sub: String,
44    #[serde(with = "jwt_numeric_date")]
45    iat: OffsetDateTime,
46    #[serde(with = "jwt_numeric_date")]
47    exp: OffsetDateTime,
48}
49
50impl Claims {
51    /// If a token should always be equal to its representation after serializing and deserializing
52    /// again, this function must be used for construction. `OffsetDateTime` contains a microsecond
53    /// field but JWT timestamps are defined as UNIX timestamps (seconds). This function normalizes
54    /// the timestamps.
55    pub fn new(iss: String, sub: String, iat: OffsetDateTime, exp: OffsetDateTime) -> Self {
56        // normalize the timestamps by stripping of microseconds
57        let iat = iat
58            .date()
59            .with_hms_milli(iat.hour(), iat.minute(), iat.second(), 0)
60            .unwrap()
61            .assume_utc();
62        let exp = exp
63            .date()
64            .with_hms_milli(exp.hour(), exp.minute(), exp.second(), 0)
65            .unwrap()
66            .assume_utc();
67
68        Self { iss, sub, iat, exp }
69    }
70}
71
72mod jwt_numeric_date {
73    //! Custom serialization of OffsetDateTime to conform with the JWT spec (RFC 7519 section 2, "Numeric Date")
74    use serde::{self, Deserialize, Deserializer, Serializer};
75    use time::OffsetDateTime;
76
77    /// Serializes an OffsetDateTime to a Unix timestamp (milliseconds since 1970/1/1T00:00:00T)
78    pub fn serialize<S>(date: &OffsetDateTime, serializer: S) -> Result<S::Ok, S::Error>
79    where
80        S: Serializer,
81    {
82        let timestamp = date.unix_timestamp();
83        serializer.serialize_i64(timestamp)
84    }
85
86    /// Attempts to deserialize an i64 and use as a Unix timestamp
87    pub fn deserialize<'de, D>(deserializer: D) -> Result<OffsetDateTime, D::Error>
88    where
89        D: Deserializer<'de>,
90    {
91        OffsetDateTime::from_unix_timestamp(i64::deserialize(deserializer)?)
92            .map_err(|_| serde::de::Error::custom("invalid Unix timestamp value"))
93    }
94}
95
96fn pubkey_fingerprint(pubkey: &[u8]) -> String {
97    let mut hasher = Sha256::new();
98    hasher.update(pubkey);
99
100    base64::engine::general_purpose::STANDARD.encode(hasher.finalize())
101}
102
103pub fn generate_jwt_token(
104    private_key_pem: &str,
105    // Snowflake expects uppercase <account identifier>.<username>
106    full_identifier: &str,
107) -> Result<String, JwtError> {
108    // Reading a private key:
109    // rsa-2048.p8 -> public key -> der bytes -> hash
110    let pkey = rsa::RsaPrivateKey::from_pkcs8_pem(private_key_pem)?;
111    let pubk = pkey.to_public_key().to_public_key_der()?;
112    let iss = format!(
113        "{}.SHA256:{}",
114        full_identifier,
115        pubkey_fingerprint(pubk.as_bytes())
116    );
117
118    let iat = OffsetDateTime::now_utc();
119    let exp = iat + Duration::days(1);
120
121    let claims = Claims::new(iss, full_identifier.to_owned(), iat, exp);
122    let ek = EncodingKey::from_rsa_der(pkey.to_pkcs1_der()?.as_bytes());
123
124    let res = encode(&Header::new(Algorithm::RS256), &claims, &ek)?;
125    Ok(res)
126}