perfect-postcode/server-rs/src/routes/stripe_webhook.rs

use std::sync::Arc;

use axum::body::Bytes;
use axum::extract::State;
use axum::http::{HeaderMap, StatusCode};
use axum::response::{IntoResponse, Response};
use hmac::{Hmac, KeyInit, Mac};
use sha2::Sha256;
use tracing::{info, warn};

use crate::checkout_sessions::{
    grant_license, mark_checkout_completed, mark_referral_invite_used, verify_checkout_completion,
    CheckoutCompletion,
};
use crate::state::SharedState;

type HmacSha256 = Hmac<Sha256>;

/// Verify Stripe webhook signature (v1 scheme).
fn verify_signature(payload: &[u8], sig_header: &str, secret: &str) -> bool {
    // Parse timestamp and signature from header: "t=TIMESTAMP,v1=SIGNATURE"
    let mut timestamp = None;
    let mut signatures = Vec::new();
    for part in sig_header.split(',') {
        if let Some(ts) = part.strip_prefix("t=") {
            timestamp = Some(ts);
        } else if let Some(sig) = part.strip_prefix("v1=") {
            signatures.push(sig);
        }
    }

    let Some(ts) = timestamp else {
        return false;
    };
    if signatures.is_empty() {
        return false;
    }

    // Reject webhooks older than 5 minutes to prevent replay attacks
    if let Ok(ts_secs) = ts.parse::<i64>() {
        let now = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs() as i64;
        if (now - ts_secs).abs() > 300 {
            return false;
        }
    } else {
        return false;
    }

    let mut signed_payload = Vec::with_capacity(ts.len() + 1 + payload.len());
    signed_payload.extend_from_slice(ts.as_bytes());
    signed_payload.push(b'.');
    signed_payload.extend_from_slice(payload);

    signatures.into_iter().any(|sig_hex| {
        let Ok(sig_bytes) = hex::decode(sig_hex) else {
            return false;
        };
        let Ok(mut mac) = HmacSha256::new_from_slice(secret.as_bytes()) else {
            return false;
        };
        mac.update(&signed_payload);
        mac.verify_slice(&sig_bytes).is_ok()
    })
}

/// Handle Stripe webhook events.
/// On `checkout.session.completed`, updates the user's subscription to "licensed".
pub async fn post_stripe_webhook(
    State(shared): State<Arc<SharedState>>,
    headers: HeaderMap,
    body: Bytes,
) -> Response {
    let state = shared.load_state();
    let webhook_secret = &state.stripe_webhook_secret;

    let sig_header = match headers
        .get("stripe-signature")
        .and_then(|h| h.to_str().ok())
    {
        Some(s) => s,
        None => {
            warn!("Missing Stripe-Signature header");
            return StatusCode::BAD_REQUEST.into_response();
        }
    };

    if !verify_signature(&body, sig_header, webhook_secret) {
        warn!("Invalid Stripe webhook signature");
        return StatusCode::BAD_REQUEST.into_response();
    }

    let event: serde_json::Value = match serde_json::from_slice(&body) {
        Ok(v) => v,
        Err(err) => {
            warn!("Failed to parse webhook body: {err}");
            return StatusCode::BAD_REQUEST.into_response();
        }
    };

    let event_type = event["type"].as_str().unwrap_or("");
    let event_id = event["id"].as_str().unwrap_or("");

    info!(event_id, event_type, "Received Stripe webhook");

    if event_type == "checkout.session.completed" {
        let session = &event["data"]["object"];
        match verify_checkout_completion(&state, session).await {
            Ok(CheckoutCompletion::Grant(checkout)) => {
                if let Err(err) = mark_referral_invite_used(
                    &state,
                    &checkout.referral_invite_id,
                    &checkout.user_id,
                )
                .await
                {
                    warn!(
                        user_id = %checkout.user_id,
                        reservation_id = %checkout.reservation_id,
                        referral_invite_id = %checkout.referral_invite_id,
                        "Failed to mark referral invite used after Stripe checkout: {err:?}"
                    );
                    return StatusCode::INTERNAL_SERVER_ERROR.into_response();
                }
                if let Err(err) = grant_license(&state, &checkout.user_id).await {
                    warn!(
                        user_id = %checkout.user_id,
                        reservation_id = %checkout.reservation_id,
                        "Failed to grant license after Stripe checkout: {err:?}"
                    );
                    return StatusCode::INTERNAL_SERVER_ERROR.into_response();
                }
                if let Err(err) = mark_checkout_completed(
                    &state,
                    &checkout.reservation_id,
                    checkout.paid_amount_pence,
                )
                .await
                {
                    warn!(
                        user_id = %checkout.user_id,
                        reservation_id = %checkout.reservation_id,
                        "Failed to mark checkout completed after license grant: {err:?}"
                    );
                    return StatusCode::INTERNAL_SERVER_ERROR.into_response();
                }
                info!(
                    user_id = %checkout.user_id,
                    reservation_id = %checkout.reservation_id,
                    "User subscription updated to licensed via verified Stripe checkout"
                );
            }
            Ok(CheckoutCompletion::AlreadyHandled) => {
                info!("Stripe checkout session was already handled");
            }
            Ok(CheckoutCompletion::Rejected(reason)) => {
                warn!("Rejecting Stripe checkout completion: {reason}");
            }
            Err(err) => {
                warn!("Failed to verify Stripe checkout completion: {err:?}");
                return StatusCode::INTERNAL_SERVER_ERROR.into_response();
            }
        }
    }

    StatusCode::OK.into_response()
}