qstash-rs 0.6.0

use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
use jsonwebtoken::{decode, Algorithm, DecodingKey, Validation};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};

use crate::{error::Result, Error};

/// Receiver configuration used to verify inbound QStash requests.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ReceiverConfig {
    /// Current signing key.
    pub current_signing_key: String,
    /// Next signing key.
    pub next_signing_key: String,
}

/// Verification input for [`Receiver::verify`].
#[derive(Debug, Clone, Copy)]
pub struct VerifyRequest<'a> {
    /// JWT from the `Upstash-Signature` header.
    pub signature: &'a str,
    /// Exact raw request body bytes.
    pub body: &'a [u8],
    /// Exact URL that should appear in the JWT `sub` claim.
    pub url: Option<&'a str>,
    /// Clock tolerance in seconds used for JWT time validation.
    pub clock_tolerance_seconds: u64,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
struct Claims {
    iss: String,
    sub: String,
    exp: u64,
    nbf: u64,
    iat: u64,
    jti: String,
    body: String,
}

/// Verifies `Upstash-Signature` JWTs.
#[derive(Debug, Clone)]
pub struct Receiver {
    current_signing_key: String,
    next_signing_key: String,
}

impl Receiver {
    /// Creates a receiver from the current and next signing keys.
    pub fn new(config: ReceiverConfig) -> Self {
        Self {
            current_signing_key: config.current_signing_key,
            next_signing_key: config.next_signing_key,
        }
    }

    /// Verifies an inbound QStash request.
    pub fn verify(&self, request: VerifyRequest<'_>) -> Result<()> {
        let claims = self
            .verify_with_key(&self.current_signing_key, request)
            .or_else(|_| self.verify_with_key(&self.next_signing_key, request))?;

        let _ = (&claims.iss, claims.exp, claims.nbf, claims.iat, &claims.jti);

        if let Some(url) = request.url {
            if claims.sub != url {
                return Err(Error::Signature {
                    message: format!(
                        "jwt subject mismatch: expected `{url}`, got `{}`",
                        claims.sub
                    ),
                });
            }
        }

        let body_hash = URL_SAFE_NO_PAD.encode(Sha256::digest(request.body));
        if trim_padding(&claims.body) != body_hash {
            return Err(Error::Signature {
                message: String::from("jwt body hash does not match the raw request body"),
            });
        }

        Ok(())
    }

    fn verify_with_key(&self, key: &str, request: VerifyRequest<'_>) -> Result<Claims> {
        let mut validation = Validation::new(Algorithm::HS256);
        validation.leeway = request.clock_tolerance_seconds;
        validation.set_issuer(&["Upstash"]);

        let token = decode::<Claims>(
            request.signature,
            &DecodingKey::from_secret(key.as_bytes()),
            &validation,
        )?;

        Ok(token.claims)
    }
}

fn trim_padding(value: &str) -> &str {
    value.trim_end_matches('=')
}

#[cfg(test)]
mod tests {
    use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
    use jsonwebtoken::{encode, Algorithm, EncodingKey, Header};
    use sha2::{Digest, Sha256};

    use super::{Claims, Receiver, ReceiverConfig, VerifyRequest};

    fn signature(body: &[u8], url: &str, key: &str) -> String {
        let claims = Claims {
            iss: String::from("Upstash"),
            sub: String::from(url),
            exp: (time_now() + 300) as u64,
            nbf: time_now() as u64,
            iat: time_now() as u64,
            jti: String::from("jwt_test"),
            body: URL_SAFE_NO_PAD.encode(Sha256::digest(body)),
        };

        encode(
            &Header::new(Algorithm::HS256),
            &claims,
            &EncodingKey::from_secret(key.as_bytes()),
        )
        .expect("signature should encode")
    }

    fn time_now() -> i64 {
        use std::time::{SystemTime, UNIX_EPOCH};

        SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .expect("system time should be after epoch")
            .as_secs() as i64
    }

    #[test]
    fn verify_should_accept_current_signing_key() {
        let receiver = Receiver::new(ReceiverConfig {
            current_signing_key: String::from("current"),
            next_signing_key: String::from("next"),
        });
        let body = br#"{"hello":"world"}"#;
        let url = "https://example.com/qstash";
        let token = signature(body, url, "current");

        receiver
            .verify(VerifyRequest {
                signature: &token,
                body,
                url: Some(url),
                clock_tolerance_seconds: 0,
            })
            .expect("current signing key should verify");
    }

    #[test]
    fn verify_should_fall_back_to_next_signing_key() {
        let receiver = Receiver::new(ReceiverConfig {
            current_signing_key: String::from("current"),
            next_signing_key: String::from("next"),
        });
        let body = b"hello";
        let url = "https://example.com/qstash";
        let token = signature(body, url, "next");

        receiver
            .verify(VerifyRequest {
                signature: &token,
                body,
                url: Some(url),
                clock_tolerance_seconds: 0,
            })
            .expect("next signing key should verify");
    }

    #[test]
    fn verify_should_reject_body_mismatch() {
        let receiver = Receiver::new(ReceiverConfig {
            current_signing_key: String::from("current"),
            next_signing_key: String::from("next"),
        });
        let url = "https://example.com/qstash";
        let token = signature(b"expected", url, "current");

        let error = receiver
            .verify(VerifyRequest {
                signature: &token,
                body: b"actual",
                url: Some(url),
                clock_tolerance_seconds: 0,
            })
            .expect_err("body mismatch should fail");

        assert!(error.to_string().contains("body hash"));
    }
}