codecraft/server/src/push.rs

//! Web Push notifications (VAPID, RFC 8030/8291/8292) without external services.
//!
//! The server sends a push to a user's subscribed devices when an agent asks
//! for permission, asks a question, or finishes a longer run – but only while
//! none of the user's app windows is visible.

use crate::app::App;
use crate::types::Event;
use aes_gcm::aead::{Aead, KeyInit};
use anyhow::{Context, Result, anyhow, bail};
use base64::Engine;
use base64::engine::general_purpose::URL_SAFE_NO_PAD as B64;
use hkdf::Hkdf;
use p256::ecdsa::signature::Signer;
use p256::elliptic_curve::sec1::ToSec1Point;
use p256::{PublicKey, SecretKey};
use serde_json::{Value, json};
use sha2::Sha256;
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::{Arc, Mutex, OnceLock};
use std::time::{Duration, Instant};

/// Runs shorter than this don't trigger a "finished" notification.
const MIN_RUN_FOR_DONE: Duration = Duration::from_secs(20);

pub struct Push {
    key_path: PathBuf,
    key: OnceLock<SecretKey>,
    /// user id -> number of visible app windows (reported over the WebSocket)
    visible: Mutex<HashMap<String, usize>>,
    /// session id -> start of the current run
    runs: Mutex<HashMap<String, Instant>>,
}

fn random_bytes<const N: usize>() -> [u8; N] {
    rand::random()
}

impl Push {
    pub fn new(data_dir: &std::path::Path) -> Self {
        Self { key_path: data_dir.join("vapid.key"), key: OnceLock::new(), visible: Mutex::default(), runs: Mutex::default() }
    }

    /// The VAPID key pair, created on first use and kept on the data volume.
    fn key(&self) -> &SecretKey {
        self.key.get_or_init(|| {
            if let Ok(bytes) = std::fs::read(&self.key_path) {
                if let Ok(k) = SecretKey::from_slice(&bytes) {
                    return k;
                }
            }
            let k = loop {
                if let Ok(k) = SecretKey::from_slice(&random_bytes::<32>()) {
                    break k;
                }
            };
            let _ = std::fs::write(&self.key_path, k.to_bytes());
            #[cfg(unix)]
            {
                use std::os::unix::fs::PermissionsExt;
                let _ = std::fs::set_permissions(&self.key_path, std::fs::Permissions::from_mode(0o600));
            }
            k
        })
    }

    /// Public key for PushManager.subscribe (base64url, uncompressed point).
    pub fn public_key(&self) -> String {
        B64.encode(self.key().public_key().as_affine().to_uncompressed_point())
    }

    pub fn set_visible(&self, user_id: &str, delta: isize) {
        let mut v = self.visible.lock().unwrap();
        let n = v.entry(user_id.to_string()).or_default();
        *n = (*n as isize + delta).max(0) as usize;
    }

    fn is_visible(&self, user_id: &str) -> bool {
        self.visible.lock().unwrap().get(user_id).is_some_and(|n| *n > 0)
    }
}

/// RFC 8291 "aes128gcm" encryption of one record.
fn encrypt(ua_public: &[u8], auth: &[u8], plaintext: &[u8]) -> Result<Vec<u8>> {
    let as_secret = loop {
        if let Ok(k) = SecretKey::from_slice(&random_bytes::<32>()) {
            break k;
        }
    };
    encrypt_with(ua_public, auth, plaintext, &as_secret, random_bytes::<16>())
}

fn encrypt_with(ua_public: &[u8], auth: &[u8], plaintext: &[u8], as_secret: &SecretKey, salt: [u8; 16]) -> Result<Vec<u8>> {
    let ua = PublicKey::from_sec1_bytes(ua_public).map_err(|_| anyhow!("invalid p256dh key"))?;
    let as_public = as_secret.public_key().as_affine().to_uncompressed_point();
    let shared = p256::ecdh::diffie_hellman(as_secret.to_nonzero_scalar(), ua.as_affine());

    // IKM from the ECDH secret and the subscription's auth secret.
    let mut key_info = b"WebPush: info\0".to_vec();
    key_info.extend_from_slice(ua_public);
    key_info.extend_from_slice(&as_public);
    let mut ikm = [0u8; 32];
    Hkdf::<Sha256>::new(Some(auth), shared.raw_secret_bytes()).expand(&key_info, &mut ikm).map_err(|_| anyhow!("hkdf"))?;

    let prk = Hkdf::<Sha256>::new(Some(&salt), &ikm);
    let mut cek = [0u8; 16];
    let mut nonce = [0u8; 12];
    prk.expand(b"Content-Encoding: aes128gcm\0", &mut cek).map_err(|_| anyhow!("hkdf"))?;
    prk.expand(b"Content-Encoding: nonce\0", &mut nonce).map_err(|_| anyhow!("hkdf"))?;

    let mut record = plaintext.to_vec();
    record.push(2); // padding delimiter: last record
    let cipher = aes_gcm::Aes128Gcm::new_from_slice(&cek).map_err(|_| anyhow!("aes key"))?;
    let ct = cipher.encrypt(&nonce.into(), record.as_slice()).map_err(|_| anyhow!("encryption failed"))?;

    let mut out = Vec::with_capacity(16 + 4 + 1 + 65 + ct.len());
    out.extend_from_slice(&salt);
    out.extend_from_slice(&4096u32.to_be_bytes());
    out.push(as_public.len() as u8);
    out.extend_from_slice(&as_public);
    out.extend_from_slice(&ct);
    Ok(out)
}

/// VAPID authorization header (RFC 8292) for an endpoint.
fn vapid_header(key: &SecretKey, endpoint: &str, subject: &str) -> Result<String> {
    let url = reqwest::Url::parse(endpoint)?;
    let aud = url.origin().ascii_serialization();
    let exp = chrono::Utc::now().timestamp() + 12 * 3600;
    let header = B64.encode(br#"{"typ":"JWT","alg":"ES256"}"#);
    let claims = B64.encode(serde_json::to_vec(&json!({"aud": aud, "exp": exp, "sub": subject}))?);
    let signing_input = format!("{header}.{claims}");
    let sk = p256::ecdsa::SigningKey::from(key);
    let sig: p256::ecdsa::Signature = sk.sign(signing_input.as_bytes());
    let jwt = format!("{signing_input}.{}", B64.encode(sig.to_bytes()));
    let k = B64.encode(key.public_key().as_affine().to_uncompressed_point());
    Ok(format!("vapid t={jwt}, k={k}"))
}

/// Send a notification to all devices of a user: (delivered, error messages).
pub async fn send_to_user(app: &App, user_id: &str, payload: &Value) -> (usize, Vec<String>) {
    let subs = app.store.push_subscriptions(user_id).unwrap_or_default();
    let mut sent = 0;
    let mut errors = Vec::new();
    for sub in subs {
        match send_one(app, &sub.endpoint, &sub.p256dh, &sub.auth, payload).await {
            Ok(()) => sent += 1,
            Err(e) => {
                let gone = e.to_string().contains("410") || e.to_string().contains("404");
                tracing::warn!("push to {}…: {e:#}", sub.endpoint.chars().take(40).collect::<String>());
                errors.push(format!("{e:#}").chars().take(200).collect());
                if gone {
                    let _ = app.store.push_remove(&sub.endpoint);
                }
            }
        }
    }
    (sent, errors)
}

async fn send_one(app: &App, endpoint: &str, p256dh: &str, auth: &str, payload: &Value) -> Result<()> {
    let ua = B64.decode(p256dh.trim_end_matches('=')).context("p256dh")?;
    let auth = B64.decode(auth.trim_end_matches('=')).context("auth")?;
    let body = encrypt(&ua, &auth, &serde_json::to_vec(payload)?)?;
    let cfg = app.config();
    let subject = if cfg.auth.public_url.starts_with("https://") { cfg.auth.public_url.clone() } else { "mailto:admin@localhost".to_string() };
    let resp = crate::provider::http_client()
        .post(endpoint)
        .header("Authorization", vapid_header(app.push.key(), endpoint, &subject)?)
        .header("Content-Encoding", "aes128gcm")
        .header("Content-Type", "application/octet-stream")
        .header("TTL", "86400")
        .header("Urgency", "high")
        .timeout(Duration::from_secs(20))
        .body(body)
        .send()
        .await?;
    if !resp.status().is_success() {
        bail!("HTTP {}: {}", resp.status().as_u16(), resp.text().await.unwrap_or_default().chars().take(200).collect::<String>());
    }
    Ok(())
}

/// Owner and title of the top-level session.
fn session_info(app: &App, session_id: &str) -> Option<(String, String, String)> {
    let root = app.root_session_id(session_id);
    let s = app.store.get_session(&root).ok().flatten()?;
    let title = if s.title.is_empty() { "CodeCraft".to_string() } else { s.title.clone() };
    Some((s.owner_id?, title, root))
}

/// Background task: turn agent events into notifications.
pub fn spawn_listener(app: Arc<App>) {
    tokio::spawn(async move {
        let mut rx = app.subscribe();
        loop {
            let ev = match rx.recv().await {
                Ok(ev) => ev,
                Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
                Err(_) => break,
            };
            let (session_id, kind, body) = match &ev {
                Event::PermissionRequest { session_id, title, tool, .. } => (session_id.clone(), "permission", format!("{tool}: {title}")),
                Event::QuestionRequest { session_id, questions, .. } => {
                    let q = questions.first().map(|q| q.question.clone()).unwrap_or_default();
                    (session_id.clone(), "question", q)
                }
                Event::SessionStatus { session_id, busy: true } => {
                    app.push.runs.lock().unwrap().insert(session_id.clone(), Instant::now());
                    continue;
                }
                Event::SessionStatus { session_id, busy: false } => {
                    let started = app.push.runs.lock().unwrap().remove(session_id);
                    if !started.is_some_and(|t| t.elapsed() >= MIN_RUN_FOR_DONE) {
                        continue;
                    }
                    // Subagents report through their parent.
                    if app.store.get_session(session_id).ok().flatten().is_none_or(|s| s.parent_id.is_some()) {
                        continue;
                    }
                    (session_id.clone(), "done", String::new())
                }
                _ => continue,
            };
            let Some((owner, title, root)) = session_info(&app, &session_id) else { continue };
            if app.push.is_visible(&owner) {
                continue;
            }
            let payload = json!({
                "kind": kind,
                "title": title,
                "body": body.chars().take(180).collect::<String>(),
                "tag": format!("{root}-{kind}"),
                "url": format!("/#/s/{root}"),
            });
            let app2 = app.clone();
            tokio::spawn(async move {
                send_to_user(&app2, &owner, &payload).await;
            });
        }
    });
}

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

    /// Decrypt with the user agent's keys (RFC 8291) to check the encryption.
    #[test]
    fn roundtrip() {
        let ua_secret = SecretKey::from_slice(&[7u8; 32]).unwrap();
        let ua_public = ua_secret.public_key().as_affine().to_uncompressed_point();
        let auth = [9u8; 16];
        let msg = br#"{"title":"hi"}"#;
        let body = encrypt(&ua_public, &auth, msg).unwrap();

        let salt = &body[..16];
        let idlen = body[20] as usize;
        let as_public = &body[21..21 + idlen];
        let ct = &body[21 + idlen..];
        let as_pk = PublicKey::from_sec1_bytes(as_public).unwrap();
        let shared = p256::ecdh::diffie_hellman(ua_secret.to_nonzero_scalar(), as_pk.as_affine());
        let mut info = b"WebPush: info\0".to_vec();
        info.extend_from_slice(&ua_public);
        info.extend_from_slice(as_public);
        let mut ikm = [0u8; 32];
        Hkdf::<Sha256>::new(Some(&auth), shared.raw_secret_bytes()).expand(&info, &mut ikm).unwrap();
        let prk = Hkdf::<Sha256>::new(Some(salt), &ikm);
        let (mut cek, mut nonce) = ([0u8; 16], [0u8; 12]);
        prk.expand(b"Content-Encoding: aes128gcm\0", &mut cek).unwrap();
        prk.expand(b"Content-Encoding: nonce\0", &mut nonce).unwrap();
        let pt = aes_gcm::Aes128Gcm::new_from_slice(&cek).unwrap().decrypt(&nonce.into(), ct).unwrap();
        assert_eq!(&pt[..pt.len() - 1], msg);
        assert_eq!(pt[pt.len() - 1], 2);
    }

    /// Example from RFC 8291, section 5.
    #[test]
    fn rfc8291_vector() {
        let d = |s: &str| B64.decode(s).unwrap();
        let as_secret = SecretKey::from_slice(&d("yfWPiYE-n46HLnH0KqZOF1fJJU3MYrct3AELtAQ-oRw")).unwrap();
        let ua_public = d("BCVxsr7N_eNgVRqvHtD0zTZsEc6-VV-JvLexhqUzORcxaOzi6-AYWXvTBHm4bjyPjs7Vd8pZGH6SRpkNtoIAiw4");
        let auth = d("BTBZMqHH6r4Tts7J_aSIgg");
        let salt: [u8; 16] = d("DGv6ra1nlYgDCS1FRnbzlw").try_into().unwrap();
        let out = encrypt_with(&ua_public, &auth, b"When I grow up, I want to be a watermelon", &as_secret, salt).unwrap();
        assert_eq!(
            B64.encode(out),
            "DGv6ra1nlYgDCS1FRnbzlwAAEABBBP4z9KsN6nGRTbVYI_c7VJSPQTBtkgcy27mlmlMoZIIgDll6e3vCYLocInmYWAmS6TlzAC8wEqKK6PBru3jl7A_yl95bQpu6cVPTpK4Mqgkf1CXztLVBSt2Ks3oZwbuwXPXLWyouBWLVWGNWQexSgSxsj_Qulcy4a-fN"
        );
    }

    #[test]
    fn vapid_jwt_verifies() {
        use p256::ecdsa::signature::Verifier;
        let key = SecretKey::from_slice(&[3u8; 32]).unwrap();
        let h = vapid_header(&key, "https://fcm.googleapis.com/fcm/send/abc", "mailto:a@b.c").unwrap();
        let jwt = h.strip_prefix("vapid t=").unwrap().split(", k=").next().unwrap();
        let (input, sig) = jwt.rsplit_once('.').unwrap();
        let sig = p256::ecdsa::Signature::from_slice(&B64.decode(sig).unwrap()).unwrap();
        p256::ecdsa::VerifyingKey::from(key.public_key()).verify(input.as_bytes(), &sig).unwrap();
        let claims: Value = serde_json::from_slice(&B64.decode(input.split('.').nth(1).unwrap()).unwrap()).unwrap();
        assert_eq!(claims["aud"], "https://fcm.googleapis.com");
    }
}