crank/apps/admin-api/tests/integration/request_context.rs

use std::{
    io,
    sync::{Arc, Mutex},
};

use admin_api::request_context::{REQUEST_ID_HEADER, apply_request_context};
use axum::{
    Router,
    body::Body,
    http::{HeaderMap, Request, StatusCode},
    routing::get,
};
use crank_observability::{
    ObservabilityConfig, RedactionLimits, ServiceIdentity, inject_current_trace_context,
};
use opentelemetry::{global, trace::TracerProvider as _};
use opentelemetry_sdk::{propagation::TraceContextPropagator, trace::SdkTracerProvider};
use tower::ServiceExt;
use tracing::instrument::WithSubscriber;
use tracing_subscriber::{fmt::MakeWriter, layer::SubscriberExt};
use uuid::Version;

#[tokio::test(flavor = "current_thread")]
async fn logs_request_completion_and_rejects_untrusted_values() {
    let writer = SharedLogWriter::default();
    let subscriber = crank_observability::build_subscriber(
        ObservabilityConfig::new(
            ServiceIdentity::try_new("admin-api", "test", "test").unwrap(),
            "info",
            RedactionLimits::default(),
        ),
        writer.clone(),
    )
    .unwrap();
    let dispatch = tracing::Dispatch::new(subscriber);
    let app = probe_app();

    let response = app
        .clone()
        .oneshot(
            Request::builder()
                .uri("/probe")
                .header(REQUEST_ID_HEADER.as_str(), "req_admin_trace_123")
                .body(Body::empty())
                .unwrap(),
        )
        .with_subscriber(dispatch.clone())
        .await
        .unwrap();

    assert_eq!(response.status(), StatusCode::OK);
    assert_eq!(
        response.headers()[REQUEST_ID_HEADER.as_str()]
            .to_str()
            .unwrap(),
        "req_admin_trace_123"
    );
    let event: serde_json::Value = writer
        .output()
        .lines()
        .map(|line| serde_json::from_str(line).unwrap())
        .find(|event: &serde_json::Value| event["event"] == "admin.request.completed")
        .unwrap();
    assert_eq!(event["request_id"], "req_admin_trace_123");
    assert_eq!(event["fields"]["status"], 200);

    let invalid_response = app
        .oneshot(
            Request::builder()
                .uri("/probe")
                .header(REQUEST_ID_HEADER.as_str(), "bad,value")
                .header("traceparent", "canary-invalid-traceparent")
                .body(Body::empty())
                .unwrap(),
        )
        .with_subscriber(dispatch)
        .await
        .unwrap();
    let generated = invalid_response.headers()[REQUEST_ID_HEADER.as_str()]
        .to_str()
        .unwrap();
    assert_eq!(
        uuid::Uuid::parse_str(generated).unwrap().get_version(),
        Some(Version::SortRand)
    );
    assert!(!writer.output().contains("canary-invalid-traceparent"));
}

#[tokio::test(flavor = "current_thread")]
async fn covers_valid_invalid_and_absent_traceparent() {
    global::set_text_map_propagator(TraceContextPropagator::new());
    let provider = SdkTracerProvider::builder().build();
    let tracer = provider.tracer("admin-request-context-test");
    let subscriber =
        tracing_subscriber::registry().with(tracing_opentelemetry::layer().with_tracer(tracer));
    let dispatch = tracing::Dispatch::new(subscriber);
    let app = trace_probe_app();

    let valid = observed_trace_id(
        app.clone()
            .oneshot(
                Request::builder()
                    .uri("/trace")
                    .header(
                        "traceparent",
                        "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01",
                    )
                    .header(REQUEST_ID_HEADER.as_str(), "request-id-is-separate")
                    .body(Body::empty())
                    .unwrap(),
            )
            .with_subscriber(dispatch.clone())
            .await
            .unwrap(),
    );
    let invalid = observed_trace_id(
        app.clone()
            .oneshot(
                Request::builder()
                    .uri("/trace")
                    .header("traceparent", "canary-invalid-traceparent")
                    .body(Body::empty())
                    .unwrap(),
            )
            .with_subscriber(dispatch.clone())
            .await
            .unwrap(),
    );
    let absent = observed_trace_id(
        app.oneshot(
            Request::builder()
                .uri("/trace")
                .body(Body::empty())
                .unwrap(),
        )
        .with_subscriber(dispatch)
        .await
        .unwrap(),
    );

    assert_eq!(valid, "0af7651916cd43dd8448eb211c80319c");
    assert_ne!(invalid, valid);
    assert_ne!(absent, valid);
    assert_ne!(invalid, absent);
    provider.shutdown().unwrap();
}

fn probe_app() -> Router {
    Router::new()
        .route("/probe", get(|| async { "ok" }))
        .layer(axum::middleware::from_fn(apply_request_context))
}

fn trace_probe_app() -> Router {
    Router::new()
        .route("/trace", get(observed_traceparent))
        .layer(axum::middleware::from_fn(apply_request_context))
}

async fn observed_traceparent() -> HeaderMap {
    let mut trace_headers = HeaderMap::new();
    inject_current_trace_context(&mut trace_headers);
    let mut response_headers = HeaderMap::new();
    if let Some(traceparent) = trace_headers.remove("traceparent") {
        response_headers.insert("x-observed-traceparent", traceparent);
    }
    response_headers
}

fn observed_trace_id(response: axum::response::Response) -> String {
    let traceparent = response.headers()["x-observed-traceparent"]
        .to_str()
        .unwrap();
    traceparent[3..35].to_owned()
}

#[derive(Clone, Default)]
struct SharedLogWriter {
    buffer: Arc<Mutex<Vec<u8>>>,
}

impl SharedLogWriter {
    fn output(&self) -> String {
        String::from_utf8(self.buffer.lock().unwrap().clone()).unwrap()
    }
}

impl<'a> MakeWriter<'a> for SharedLogWriter {
    type Writer = SharedLogGuard;

    fn make_writer(&'a self) -> Self::Writer {
        SharedLogGuard {
            buffer: Arc::clone(&self.buffer),
        }
    }
}

struct SharedLogGuard {
    buffer: Arc<Mutex<Vec<u8>>>,
}

impl io::Write for SharedLogGuard {
    fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
        self.buffer.lock().unwrap().extend_from_slice(bytes);
        Ok(bytes.len())
    }

    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}