kvitt/server/crates/kvitt-server/tests/push_fcm/stub.rs

//! A local stub of `oauth2.googleapis.com/token` and `fcm.googleapis.com/v1/projects/{p}/messages:send`, plus a
//! throwaway RSA service-account key generated per test (no key is committed).

use std::collections::VecDeque;
use std::net::SocketAddr;
use std::sync::{Arc, Mutex};

use aws_lc_rs::encoding::AsDer;
use aws_lc_rs::rsa::{KeyPair, KeySize};
use aws_lc_rs::signature::KeyPair as _;
use axum::Router;
use axum::body::Bytes;
use axum::extract::{Path, State};
use axum::http::{HeaderMap, StatusCode};
use axum::response::{IntoResponse, Response};
use axum::routing::post;
use base64::Engine;
use kvitt_server::push::Endpoints;
use serde_json::{Value, json};

/// One request the stub received.
#[derive(Debug, Clone)]
pub struct Received {
    pub path: String,
    pub authorization: Option<String>,
    pub body: Bytes,
}

impl Received {
    pub fn json(&self) -> Value {
        serde_json::from_slice(&self.body).unwrap()
    }
}

#[derive(Default)]
pub struct StubState {
    pub tokens: Mutex<Vec<Received>>,
    pub sends: Mutex<Vec<Received>>,
    /// Scripted send responses, used in order; afterwards 200.
    pub send_script: Mutex<VecDeque<(u16, Value)>>,
    /// Scripted token responses; afterwards a fresh token valid 3600 s.
    pub token_script: Mutex<VecDeque<(u16, Value)>>,
}

pub struct Stub {
    pub addr: SocketAddr,
    pub state: Arc<StubState>,
}

impl Stub {
    pub async fn start() -> Self {
        let state = Arc::new(StubState::default());
        let app = Router::new()
            .route("/token", post(token))
            .route("/v1/projects/{project}/{action}", post(send))
            .with_state(state.clone());
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();
        tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
        Self { addr, state }
    }

    pub fn endpoints(&self) -> Endpoints {
        Endpoints {
            token_url: format!("http://{}/token", self.addr),
            send_base_url: format!("http://{}", self.addr),
        }
    }

    pub fn script_sends(&self, responses: &[(u16, Value)]) {
        self.state
            .send_script
            .lock()
            .unwrap()
            .extend(responses.iter().cloned());
    }

    pub fn sends(&self) -> Vec<Received> {
        self.state.sends.lock().unwrap().clone()
    }

    pub fn tokens(&self) -> Vec<Received> {
        self.state.tokens.lock().unwrap().clone()
    }
}

fn record(headers: &HeaderMap, path: String, body: Bytes) -> Received {
    Received {
        path,
        authorization: headers
            .get("authorization")
            .map(|v| v.to_str().unwrap().to_owned()),
        body,
    }
}

async fn token(State(state): State<Arc<StubState>>, headers: HeaderMap, body: Bytes) -> Response {
    let n = {
        let mut tokens = state.tokens.lock().unwrap();
        tokens.push(record(&headers, "/token".into(), body));
        tokens.len()
    };
    if let Some((status, body)) = state.token_script.lock().unwrap().pop_front() {
        return (StatusCode::from_u16(status).unwrap(), axum::Json(body)).into_response();
    }
    axum::Json(json!({"access_token": format!("ya29.stub-{n}"), "expires_in": 3600, "token_type": "Bearer"}))
        .into_response()
}

async fn send(
    State(state): State<Arc<StubState>>,
    Path((project, action)): Path<(String, String)>,
    headers: HeaderMap,
    body: Bytes,
) -> Response {
    state.sends.lock().unwrap().push(record(
        &headers,
        format!("/v1/projects/{project}/{action}"),
        body,
    ));
    match state.send_script.lock().unwrap().pop_front() {
        Some((status, body)) => {
            (StatusCode::from_u16(status).unwrap(), axum::Json(body)).into_response()
        }
        None => {
            axum::Json(json!({"name": format!("projects/{project}/messages/1")})).into_response()
        }
    }
}

/// A generated service-account key file and the public key to verify its JWTs.
pub struct TestKey {
    pub json: String,
    pub public_der: Vec<u8>,
}

pub fn service_account() -> TestKey {
    let pair = KeyPair::generate(KeySize::Rsa2048).unwrap();
    let der = AsDer::<aws_lc_rs::encoding::Pkcs8V1Der>::as_der(&pair).unwrap();
    let b64 = base64::engine::general_purpose::STANDARD.encode(der.as_ref());
    let lines: Vec<&str> = b64
        .as_bytes()
        .chunks(64)
        .map(|c| std::str::from_utf8(c).unwrap())
        .collect();
    let pem = format!(
        "-----BEGIN PRIVATE KEY-----\n{}\n-----END PRIVATE KEY-----\n",
        lines.join("\n")
    );
    let json = json!({
        "type": "service_account",
        "project_id": "kvitt-1337",
        "private_key_id": "test-key-1",
        "private_key": pem,
        "client_email": "kvitt-fcm@kvitt-1337.iam.gserviceaccount.com",
        "token_uri": "https://oauth2.googleapis.com/token"
    })
    .to_string();
    TestKey {
        json,
        public_der: pair.public_key().as_ref().to_vec(),
    }
}

/// An FCM error body with an `FcmError` detail.
pub fn fcm_error(code: u16, status: &str, error_code: &str) -> Value {
    json!({"error": {"code": code, "message": "stub", "status": status, "details": [
        {"@type": "type.googleapis.com/google.firebase.fcm.v1.FcmError", "errorCode": error_code}
    ]}})
}