PiCloud/crates/shared/src/subscriber_token.rs

//! HMAC-signed realtime subscriber tokens (v1.1.6, design notes §5).
//!
//! A token is a compact, URL-safe, two-part string:
//!
//! ```text
//! <base64url(payload)>.<base64url(signature)>
//! ```
//!
//! where `payload` is the JSON [`TokenClaims`] and `signature` is
//! `HMAC-SHA256(app_signing_key, base64url(payload))`. Tokens are minted
//! by scripts via `pubsub::subscriber_token` (the minter lives in
//! manager-core's `PubsubServiceImpl`) and verified by the SSE subscribe
//! path (the verifier lives in manager-core's `RealtimeAuthority` impl).
//! Both sides depend on this module so the byte-for-byte contract has a
//! single home.
//!
//! There is no per-token revocation in v1.1.6 by design: HMAC bearers
//! can't be individually revoked. Rotating an app's signing key
//! invalidates every token for that app wholesale; short TTLs are the
//! safety mechanism.

use base64::engine::general_purpose::URL_SAFE_NO_PAD;
use base64::Engine as _;
use hmac::{Hmac, Mac};
use serde::{Deserialize, Serialize};
use sha2::Sha256;
use thiserror::Error;

use crate::AppId;

type HmacSha256 = Hmac<Sha256>;

/// The signed payload. `exp` / `iat` are Unix seconds.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct TokenClaims {
    pub app_id: AppId,
    pub topics: Vec<String>,
    pub exp: i64,
    pub iat: i64,
}

impl TokenClaims {
    /// Does this token grant access to `topic`?
    #[must_use]
    pub fn allows_topic(&self, topic: &str) -> bool {
        self.topics.iter().any(|t| t == topic)
    }

    /// Is the token expired relative to `now_unix` (Unix seconds)?
    #[must_use]
    pub fn is_expired(&self, now_unix: i64) -> bool {
        now_unix >= self.exp
    }
}

#[derive(Debug, Error, PartialEq, Eq)]
pub enum TokenError {
    #[error("token is malformed")]
    Malformed,
    #[error("token signature is invalid")]
    BadSignature,
    #[error("token has expired")]
    Expired,
}

/// Sign `claims` with `key`, producing the `payload.signature` string.
#[must_use]
pub fn sign(key: &[u8], claims: &TokenClaims) -> String {
    // `serde_json` on a fixed-field struct never fails to serialize.
    let payload_json = serde_json::to_vec(claims).expect("TokenClaims serialize");
    let payload_b64 = URL_SAFE_NO_PAD.encode(&payload_json);
    let sig = mac_sign(key, payload_b64.as_bytes());
    let sig_b64 = URL_SAFE_NO_PAD.encode(sig);
    format!("{payload_b64}.{sig_b64}")
}

/// Verify `token` against `key` and check expiry against `now_unix`
/// (Unix seconds). Returns the decoded [`TokenClaims`] on success.
///
/// Topic-scope checking (is the requested topic in the token's list?)
/// is the caller's responsibility via [`TokenClaims::allows_topic`] —
/// this function proves authenticity + liveness only.
///
/// # Errors
///
/// [`TokenError::Malformed`] if the shape / base64 / JSON is wrong,
/// [`TokenError::BadSignature`] if the HMAC doesn't match, or
/// [`TokenError::Expired`] if `now_unix >= exp`.
pub fn verify(key: &[u8], token: &str, now_unix: i64) -> Result<TokenClaims, TokenError> {
    let (payload_b64, sig_b64) = token.split_once('.').ok_or(TokenError::Malformed)?;

    let provided_sig = URL_SAFE_NO_PAD
        .decode(sig_b64)
        .map_err(|_| TokenError::Malformed)?;

    // Constant-time verify of the MAC over the exact payload bytes.
    let mut mac = HmacSha256::new_from_slice(key).expect("HMAC accepts any key length");
    mac.update(payload_b64.as_bytes());
    mac.verify_slice(&provided_sig)
        .map_err(|_| TokenError::BadSignature)?;

    // Signature good → decode the claims and check expiry.
    let payload_json = URL_SAFE_NO_PAD
        .decode(payload_b64)
        .map_err(|_| TokenError::Malformed)?;
    let claims: TokenClaims =
        serde_json::from_slice(&payload_json).map_err(|_| TokenError::Malformed)?;

    if claims.is_expired(now_unix) {
        return Err(TokenError::Expired);
    }
    Ok(claims)
}

fn mac_sign(key: &[u8], data: &[u8]) -> Vec<u8> {
    let mut mac = HmacSha256::new_from_slice(key).expect("HMAC accepts any key length");
    mac.update(data);
    mac.finalize().into_bytes().to_vec()
}

#[cfg(test)]
mod tests {
    use super::*;

    fn claims(app: AppId, topics: &[&str], exp: i64) -> TokenClaims {
        TokenClaims {
            app_id: app,
            topics: topics.iter().map(|s| (*s).to_string()).collect(),
            iat: 1000,
            exp,
        }
    }

    #[test]
    fn round_trip_verifies() {
        let key = b"super-secret-key-bytes-0123456789";
        let app = AppId::new();
        let c = claims(app, &["chat.room.1", "user.notify"], 5000);
        let token = sign(key, &c);
        let got = verify(key, &token, 2000).expect("valid token verifies");
        assert_eq!(got, c);
        assert!(got.allows_topic("chat.room.1"));
        assert!(!got.allows_topic("chat.room.2"));
    }

    #[test]
    fn tampered_payload_fails() {
        let key = b"super-secret-key-bytes-0123456789";
        let app = AppId::new();
        let token = sign(key, &claims(app, &["t"], 5000));
        // Flip a character in the payload half.
        let (payload, sig) = token.split_once('.').unwrap();
        let mut bytes = payload.as_bytes().to_vec();
        bytes[0] ^= 0x01;
        let tampered = format!("{}.{sig}", String::from_utf8_lossy(&bytes));
        assert_eq!(verify(key, &tampered, 2000), Err(TokenError::BadSignature));
    }

    #[test]
    fn tampered_signature_fails() {
        let key = b"super-secret-key-bytes-0123456789";
        let app = AppId::new();
        let token = sign(key, &claims(app, &["t"], 5000));
        let (payload, _sig) = token.split_once('.').unwrap();
        // A valid-base64 but wrong signature.
        let bogus = URL_SAFE_NO_PAD.encode([0u8; 32]);
        let tampered = format!("{payload}.{bogus}");
        assert_eq!(verify(key, &tampered, 2000), Err(TokenError::BadSignature));
    }

    #[test]
    fn different_key_fails() {
        let app = AppId::new();
        let token = sign(
            b"key-aaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
            &claims(app, &["t"], 5000),
        );
        assert_eq!(
            verify(b"key-bbbbbbbbbbbbbbbbbbbbbbbbbbbbbb", &token, 2000),
            Err(TokenError::BadSignature)
        );
    }

    #[test]
    fn expired_token_fails_at_expiry_check() {
        let key = b"super-secret-key-bytes-0123456789";
        let app = AppId::new();
        let token = sign(key, &claims(app, &["t"], 5000));
        // now == exp → expired (>= boundary).
        assert_eq!(verify(key, &token, 5000), Err(TokenError::Expired));
        assert_eq!(verify(key, &token, 9999), Err(TokenError::Expired));
    }

    #[test]
    fn malformed_token_fails() {
        let key = b"super-secret-key-bytes-0123456789";
        assert_eq!(verify(key, "no-dot-here", 0), Err(TokenError::Malformed));
        assert_eq!(verify(key, "a.b.c", 0), Err(TokenError::Malformed));
    }
}