feat: add postgres auth sync and stripe webhook verification

auth-api/Cargo.toml

@@ -17,3 +17,8 @@ tracing-subscriber.workspace = true
 fastrand.workspace = true
 metrics.workspace = true
 metrics-exporter-prometheus.workspace = true
+tokio-postgres.workspace = true
+hmac.workspace = true
+sha2.workspace = true
+hex.workspace = true
+uuid.workspace = true

auth-api/src/main.rs

@@ -1,24 +1,32 @@
-use std::{net::SocketAddr, sync::Arc};
+use std::{net::SocketAddr, sync::Arc, time::Instant};
 
 use anyhow::{Context, Result};
 use axum::{
     Json, Router,
     extract::State,
-    http::StatusCode,
+    http::{HeaderMap, StatusCode},
     response::IntoResponse,
     routing::{get, post},
 };
-use chrono::{Duration, Utc};
+use chrono::{Duration, TimeZone, Utc};
+use hmac::{Hmac, Mac};
 use jsonwebtoken::{Algorithm, DecodingKey, EncodingKey, Header, Validation, decode, encode};
 use metrics_exporter_prometheus::{PrometheusBuilder, PrometheusHandle};
 use redis::AsyncCommands;
 use serde::{Deserialize, Serialize};
+use sha2::Sha256;
+use tokio_postgres::{Client as PgClient, NoTls};
 use tracing::{info, warn};
+use uuid::Uuid;
+
+type HmacSha256 = Hmac<Sha256>;
 
 #[derive(Clone)]
 struct AppState {
     jwt_secret: Arc<String>,
+    stripe_webhook_secret: Option<Arc<String>>,
     redis: Option<redis::aio::ConnectionManager>,
+    pg: Option<Arc<PgClient>>,
     metrics: PrometheusHandle,
 }
 
@@ -67,6 +75,17 @@ struct StripeWebhookEvent {
     tier: String,
     #[serde(default = "default_max_tunnels")]
     max_tunnels: u32,
+    #[serde(default = "default_subscription_status")]
+    status: String,
+    provider_customer_id: Option<String>,
+    provider_subscription_id: Option<String>,
+    current_period_end: Option<i64>,
+}
+
+#[derive(Debug, Clone)]
+struct PlanEntitlement {
+    tier: String,
+    max_tunnels: u32,
 }
 
 #[tokio::main]
@@ -80,9 +99,11 @@ async fn main() -> Result<()> {
 
     let bind = std::env::var("AUTH_BIND").unwrap_or_else(|_| "0.0.0.0:8080".into());
     let jwt_secret = std::env::var("JWT_SECRET").unwrap_or_else(|_| "dev-secret-change-me".into());
+    let stripe_webhook_secret = std::env::var("STRIPE_WEBHOOK_SECRET").ok().map(Arc::new);
     let metrics = PrometheusBuilder::new()
         .install_recorder()
         .context("install prometheus recorder")?;
+
     let redis = if let Ok(url) = std::env::var("REDIS_URL") {
         let client = redis::Client::open(url.clone()).context("open redis client")?;
         match redis::aio::ConnectionManager::new(client).await {
@@ -99,9 +120,13 @@ async fn main() -> Result<()> {
         None
     };
 
+    let pg = connect_postgres().await?;
+
     let state = AppState {
         jwt_secret: Arc::new(jwt_secret),
+        stripe_webhook_secret,
         redis,
+        pg,
         metrics,
     };
 
@@ -122,6 +147,22 @@ async fn main() -> Result<()> {
     Ok(())
 }
 
+async fn connect_postgres() -> Result<Option<Arc<PgClient>>> {
+    let Some(url) = std::env::var("DATABASE_URL").ok() else {
+        return Ok(None);
+    };
+    let (client, conn) = tokio_postgres::connect(&url, NoTls)
+        .await
+        .context("connect postgres")?;
+    tokio::spawn(async move {
+        if let Err(e) = conn.await {
+            warn!(error = %e, "postgres connection task ended");
+        }
+    });
+    info!("auth-api connected to postgres");
+    Ok(Some(Arc::new(client)))
+}
+
 async fn healthz() -> &'static str {
     "ok"
 }
@@ -135,7 +176,7 @@ async fn issue_dev_token(
     State(state): State<AppState>,
     Json(req): Json<DevTokenRequest>,
 ) -> Result<Json<TokenResponse>, ApiError> {
-    let started = std::time::Instant::now();
+    let started = Instant::now();
     let now = Utc::now();
     let exp = now + Duration::hours(24);
     let claims = Claims {
@@ -154,20 +195,7 @@ async fn issue_dev_token(
     )
     .map_err(ApiError::internal)?;
 
-    if let Some(mut redis) = state.redis.clone() {
-        let key = format!("auth:jwt:jti:{}", claims.jti);
-        let ttl = (claims.exp as i64 - now.timestamp()).max(1);
-        let payload = serde_json::json!({
-            "user_id": claims.sub,
-            "plan_tier": claims.tier,
-            "max_tunnels": claims.max_tunnels
-        })
-        .to_string();
-        let _: () = redis
-            .set_ex(key, payload, ttl as u64)
-            .await
-            .map_err(ApiError::internal)?;
-    }
+    sync_jwt_cache(&state, &claims).await?;
     metrics::counter!("auth_jwt_issued_total").increment(1);
     metrics::histogram!("auth_issue_token_latency_ms")
         .record(started.elapsed().as_secs_f64() * 1000.0);
@@ -183,72 +211,232 @@ async fn validate_token(
     State(state): State<AppState>,
     Json(req): Json<ValidateRequest>,
 ) -> Result<Json<ValidateResponse>, ApiError> {
-    let started = std::time::Instant::now();
+    let started = Instant::now();
     let decoded = decode::<Claims>(
         &req.token,
         &DecodingKey::from_secret(state.jwt_secret.as_bytes()),
         &Validation::new(Algorithm::HS256),
     );
 
-    match decoded {
+    let response = match decoded {
         Ok(tok) => {
-            let c = tok.claims;
-            if let Some(mut redis) = state.redis.clone() {
-                let key = format!("plan:user:{}", c.sub);
-                let payload = serde_json::json!({
-                    "tier": c.tier,
-                    "max_tunnels": c.max_tunnels,
-                    "source": "auth-api"
-                })
-                .to_string();
-                let _: () = redis.set_ex(key, payload, 300).await.map_err(ApiError::internal)?;
-            }
+            let claims = tok.claims;
+            let ent = match load_plan_for_user(&state, &claims.sub).await? {
+                Some(db_plan) => db_plan,
+                None => PlanEntitlement {
+                    tier: claims.tier.clone(),
+                    max_tunnels: claims.max_tunnels,
+                },
+            };
+            sync_plan_cache(&state, &claims.sub, &ent, "auth-validate").await?;
             metrics::counter!("auth_token_validate_total", "result" => "valid").increment(1);
-            metrics::histogram!("auth_validate_token_latency_ms")
-                .record(started.elapsed().as_secs_f64() * 1000.0);
-            Ok(Json(ValidateResponse {
+            ValidateResponse {
                 valid: true,
-                user_id: Some(c.sub),
-                tier: Some(c.tier),
-                max_tunnels: Some(c.max_tunnels),
-            }))
+                user_id: Some(claims.sub),
+                tier: Some(ent.tier),
+                max_tunnels: Some(ent.max_tunnels),
+            }
         }
-        Err(_) => Ok(Json(ValidateResponse {
-            valid: false,
-            user_id: None,
-            tier: None,
-            max_tunnels: None,
-        })).map(|resp| {
+        Err(_) => {
             metrics::counter!("auth_token_validate_total", "result" => "invalid").increment(1);
-            metrics::histogram!("auth_validate_token_latency_ms")
-                .record(started.elapsed().as_secs_f64() * 1000.0);
-            resp
-        }),
-    }
+            ValidateResponse {
+                valid: false,
+                user_id: None,
+                tier: None,
+                max_tunnels: None,
+            }
+        }
+    };
+
+    metrics::histogram!("auth_validate_token_latency_ms")
+        .record(started.elapsed().as_secs_f64() * 1000.0);
+    Ok(Json(response))
 }
 
-#[tracing::instrument(skip(state))]
+#[tracing::instrument(skip(state, headers, body))]
 async fn stripe_webhook(
     State(state): State<AppState>,
-    Json(event): Json<StripeWebhookEvent>,
+    headers: HeaderMap,
+    body: String,
 ) -> Result<impl IntoResponse, ApiError> {
-    let started = std::time::Instant::now();
+    let started = Instant::now();
+    verify_stripe_signature(state.stripe_webhook_secret.as_deref().map(|s| s.as_str()), &headers, &body)?;
+
+    let event: StripeWebhookEvent = serde_json::from_str(&body).map_err(ApiError::bad_request)?;
+    let ent = PlanEntitlement {
+        tier: event.tier.clone(),
+        max_tunnels: event.max_tunnels,
+    };
+
+    if let Some(pg) = &state.pg {
+        upsert_subscription_from_webhook(pg, &event).await?;
+    }
+    sync_plan_cache(&state, &event.user_id, &ent, "stripe_webhook").await?;
+
+    metrics::counter!("stripe_webhook_events_total", "event_type" => event.event_type.clone())
+        .increment(1);
+    metrics::histogram!("stripe_webhook_latency_ms")
+        .record(started.elapsed().as_secs_f64() * 1000.0);
+    Ok(StatusCode::NO_CONTENT)
+}
+
+async fn sync_jwt_cache(state: &AppState, claims: &Claims) -> Result<(), ApiError> {
     if let Some(mut redis) = state.redis.clone() {
-        let key = format!("plan:user:{}", event.user_id);
+        let key = format!("auth:jwt:jti:{}", claims.jti);
+        let ttl = (claims.exp as i64 - Utc::now().timestamp()).max(1);
         let payload = serde_json::json!({
-            "tier": event.tier,
-            "max_tunnels": event.max_tunnels,
-            "source": "stripe_webhook",
-            "last_event_type": event.event_type,
-            "updated_at": Utc::now().timestamp(),
+            "user_id": claims.sub,
+            "plan_tier": claims.tier,
+            "max_tunnels": claims.max_tunnels
+        })
+        .to_string();
+        let _: () = redis
+            .set_ex(key, payload, ttl as u64)
+            .await
+            .map_err(ApiError::internal)?;
+    }
+    Ok(())
+}
+
+async fn sync_plan_cache(
+    state: &AppState,
+    user_id: &str,
+    ent: &PlanEntitlement,
+    source: &str,
+) -> Result<(), ApiError> {
+    if let Some(mut redis) = state.redis.clone() {
+        let key = format!("plan:user:{user_id}");
+        let payload = serde_json::json!({
+            "tier": ent.tier,
+            "max_tunnels": ent.max_tunnels,
+            "source": source,
+            "updated_at": Utc::now().timestamp()
         })
         .to_string();
         let _: () = redis.set_ex(key, payload, 300).await.map_err(ApiError::internal)?;
     }
-    metrics::counter!("stripe_webhook_events_total", "event_type" => event.event_type.clone()).increment(1);
-    metrics::histogram!("stripe_webhook_latency_ms")
-        .record(started.elapsed().as_secs_f64() * 1000.0);
-    Ok(StatusCode::NO_CONTENT)
+    Ok(())
+}
+
+async fn load_plan_for_user(state: &AppState, user_id: &str) -> Result<Option<PlanEntitlement>, ApiError> {
+    let Some(pg) = &state.pg else {
+        return Ok(None);
+    };
+    let _ = Uuid::parse_str(user_id).map_err(ApiError::bad_request)?;
+    let row = pg
+        .query_opt(
+            r#"
+            select p.id as plan_id, p.max_tunnels
+            from subscriptions s
+            join plans p on p.id = s.plan_id
+            where s.user_id = $1::uuid and s.status in ('active', 'trialing')
+            order by s.updated_at desc
+            limit 1
+            "#,
+            &[&user_id],
+        )
+        .await
+        .map_err(ApiError::internal)?;
+
+    Ok(row.map(|r| PlanEntitlement {
+        tier: r.get::<_, String>("plan_id"),
+        max_tunnels: r.get::<_, i32>("max_tunnels") as u32,
+    }))
+}
+
+async fn upsert_subscription_from_webhook(pg: &PgClient, event: &StripeWebhookEvent) -> Result<(), ApiError> {
+    let _ = Uuid::parse_str(&event.user_id).map_err(ApiError::bad_request)?;
+
+    let period_end = event
+        .current_period_end
+        .and_then(|ts| Utc.timestamp_opt(ts, 0).single())
+        .map(|dt| dt.naive_utc());
+
+    let provider_sub_id = event
+        .provider_subscription_id
+        .clone()
+        .unwrap_or_else(|| format!("sub-dev-{}", fastrand::u64(..)));
+
+    pg.execute(
+        r#"
+        insert into subscriptions (
+            user_id, plan_id, provider, provider_customer_id, provider_subscription_id, status,
+            current_period_end, updated_at
+        )
+        values ($1::uuid, $2, 'stripe', $3, $4, $5, $6, now())
+        on conflict (provider_subscription_id)
+        do update set
+            user_id = excluded.user_id,
+            plan_id = excluded.plan_id,
+            provider_customer_id = excluded.provider_customer_id,
+            status = excluded.status,
+            current_period_end = excluded.current_period_end,
+            updated_at = now()
+        "#,
+        &[
+            &event.user_id,
+            &event.tier,
+            &event.provider_customer_id,
+            &provider_sub_id,
+            &event.status,
+            &period_end,
+        ],
+    )
+    .await
+    .map_err(ApiError::internal)?;
+
+    Ok(())
+}
+
+fn verify_stripe_signature(
+    secret: Option<&str>,
+    headers: &HeaderMap,
+    body: &str,
+) -> Result<(), ApiError> {
+    let Some(secret) = secret else {
+        return Ok(());
+    };
+
+    let header = headers
+        .get("stripe-signature")
+        .and_then(|h| h.to_str().ok())
+        .ok_or_else(|| ApiError::unauthorized("missing stripe-signature header"))?;
+
+    let mut timestamp: Option<i64> = None;
+    let mut sigs: Vec<&str> = Vec::new();
+    for part in header.split(',') {
+        let mut kv = part.trim().splitn(2, '=');
+        let k = kv.next().unwrap_or_default();
+        let v = kv.next().unwrap_or_default();
+        match k {
+            "t" => timestamp = v.parse::<i64>().ok(),
+            "v1" => sigs.push(v),
+            _ => {}
+        }
+    }
+
+    let ts = timestamp.ok_or_else(|| ApiError::unauthorized("invalid stripe signature timestamp"))?;
+    let now = Utc::now().timestamp();
+    if (now - ts).abs() > 300 {
+        return Err(ApiError::unauthorized("stale stripe signature"));
+    }
+
+    let signed_payload = format!("{ts}.{body}");
+    let mut mac = HmacSha256::new_from_slice(secret.as_bytes()).map_err(ApiError::internal)?;
+    mac.update(signed_payload.as_bytes());
+    let expected = mac.finalize().into_bytes();
+
+    for sig in sigs {
+        let Ok(bytes) = hex::decode(sig) else { continue; };
+        let mut verify_mac = HmacSha256::new_from_slice(secret.as_bytes()).map_err(ApiError::internal)?;
+        verify_mac.update(signed_payload.as_bytes());
+        if verify_mac.verify_slice(&bytes).is_ok() {
+            return Ok(());
+        }
+    }
+
+    let _ = expected;
+    Err(ApiError::unauthorized("stripe signature verification failed"))
 }
 
 #[derive(Debug)]
@@ -264,6 +452,20 @@ impl ApiError {
             message: e.to_string(),
         }
     }
+
+    fn bad_request<E: std::fmt::Display>(e: E) -> Self {
+        Self {
+            status: StatusCode::BAD_REQUEST,
+            message: e.to_string(),
+        }
+    }
+
+    fn unauthorized<E: Into<String>>(msg: E) -> Self {
+        Self {
+            status: StatusCode::UNAUTHORIZED,
+            message: msg.into(),
+        }
+    }
 }
 
 impl IntoResponse for ApiError {
@@ -279,3 +481,7 @@ fn default_tier() -> String {
 fn default_max_tunnels() -> u32 {
     1
 }
+
+fn default_subscription_status() -> String {
+    "active".to_string()
+}