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::{
    complete_verified_checkout, reverse_license_for_payment_intent, 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);

    // Verify every candidate signature without short-circuiting, so the total
    // time taken doesn't depend on which (if any) signature matched.
    let mut matched = false;
    for sig_hex in signatures {
        let Ok(sig_bytes) = hex::decode(sig_hex) else {
            continue;
        };
        let Ok(mut mac) = HmacSha256::new_from_slice(secret.as_bytes()) else {
            continue;
        };
        mac.update(&signed_payload);
        // verify_slice itself is constant-time.
        if mac.verify_slice(&sig_bytes).is_ok() {
            matched = true;
        }
    }
    matched
}

fn payment_intent_id_from_object(object: &serde_json::Value) -> Option<&str> {
    object["payment_intent"]
        .as_str()
        .filter(|id| is_safe_stripe_id(id))
}

fn is_safe_stripe_id(id: &str) -> bool {
    !id.is_empty()
        && id.len() <= 128
        && id
            .bytes()
            .all(|b| b.is_ascii_alphanumeric() || b == b'_' || b == b'-')
}

fn reversal_event_is_actionable(event_type: &str, object: &serde_json::Value) -> bool {
    match event_type {
        "charge.refunded" => {
            object["refunded"].as_bool().unwrap_or(false)
                || object["amount_refunded"].as_u64().unwrap_or(0) > 0
        }
        "charge.refund.updated" | "refund.created" | "refund.updated" => {
            matches!(object["status"].as_str(), Some("succeeded"))
        }
        "charge.dispute.created" | "charge.dispute.funds_withdrawn" => true,
        "charge.dispute.closed" => matches!(object["status"].as_str(), Some("lost")),
        _ => false,
    }
}

/// 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) = complete_verified_checkout(&state, &checkout).await {
                    warn!(
                        user_id = %checkout.user_id,
                        reservation_id = %checkout.reservation_id,
                        "Failed to complete verified Stripe checkout: {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();
            }
        }
    } else if matches!(
        event_type,
        "charge.refunded"
            | "charge.refund.updated"
            | "refund.created"
            | "refund.updated"
            | "charge.dispute.created"
            | "charge.dispute.closed"
            | "charge.dispute.funds_withdrawn"
    ) {
        let object = &event["data"]["object"];
        let Some(payment_intent_id) = payment_intent_id_from_object(object) else {
            warn!(
                event_id,
                event_type, "Stripe reversal event missing payment intent id"
            );
            return StatusCode::OK.into_response();
        };
        if !reversal_event_is_actionable(event_type, object) {
            info!(
                payment_intent_id,
                event_type, "Ignoring non-final Stripe reversal event"
            );
            return StatusCode::OK.into_response();
        }
        match reverse_license_for_payment_intent(&state, payment_intent_id, event_type).await {
            Ok(Some(user_id)) => {
                info!(
                    user_id,
                    payment_intent_id, event_type, "Processed Stripe payment reversal event"
                );
            }
            Ok(None) => {
                warn!(
                    payment_intent_id,
                    event_type, "Stripe reversal event had no matching checkout reservation"
                );
            }
            Err(err) => {
                warn!(
                    payment_intent_id,
                    event_type, "Failed to process Stripe payment reversal event: {err:?}"
                );
                return StatusCode::INTERNAL_SERVER_ERROR.into_response();
            }
        }
    }

    StatusCode::OK.into_response()
}