rs-zero 0.2.3

#![cfg(feature = "observability")]

use axum::{Router, routing::get};
use rs_zero::observability::{
    CorrelationContext, MetricsRegistry, current_trace_id, metrics_router, request_id_from_headers,
    span_id_from_traceparent, trace_id_from_traceparent, traceparent_from_headers,
};
use rs_zero::rest::{ApiResponse, RestConfig, RestMetricsConfig, RestMiddlewareConfig, RestServer};
use tower::ServiceExt;

#[tokio::test]
async fn http_request_id_is_available_for_trace_correlation() {
    let registry = MetricsRegistry::new();
    let config = RestConfig {
        metrics_registry: Some(registry.clone()),
        middlewares: RestMiddlewareConfig {
            metrics: RestMetricsConfig { enabled: true },
            ..RestMiddlewareConfig::default()
        },
        ..RestConfig::default()
    };
    let app = RestServer::new(
        config,
        Router::new()
            .route("/ready", get(|| async { ApiResponse::success("ok") }))
            .merge(metrics_router(registry.clone())),
    )
    .into_router();

    let response = app
        .oneshot(
            axum::http::Request::builder()
                .uri("/ready")
                .header("x-request-id", "trace-req-1")
                .header(
                    "traceparent",
                    "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
                )
                .body(axum::body::Body::empty())
                .expect("request"),
        )
        .await
        .expect("response");

    assert_eq!(
        response.headers().get("x-request-id").expect("request id"),
        "trace-req-1"
    );
    assert_eq!(
        request_id_from_headers(response.headers()).as_deref(),
        Some("trace-req-1")
    );
    assert!(registry.render_prometheus().contains("route=\"/ready\""));
    assert!(current_trace_id().is_none());
}

#[tokio::test]
async fn request_id_extension_is_available_to_handlers() {
    use axum::extract::Extension;
    use rs_zero::observability::CurrentRequestId;

    async fn handler(Extension(request_id): Extension<CurrentRequestId>) -> String {
        request_id.0
    }

    let registry = MetricsRegistry::new();
    let config = RestConfig {
        metrics_registry: Some(registry.clone()),
        middlewares: RestMiddlewareConfig {
            metrics: RestMetricsConfig { enabled: true },
            ..RestMiddlewareConfig::default()
        },
        ..RestConfig::default()
    };
    let app = RestServer::new(config, Router::new().route("/ready", get(handler))).into_router();

    let response = app
        .oneshot(
            axum::http::Request::builder()
                .uri("/ready")
                .header("x-request-id", "req-handler-1")
                .body(axum::body::Body::empty())
                .expect("request"),
        )
        .await
        .expect("response");
    let body = axum::body::to_bytes(response.into_body(), usize::MAX)
        .await
        .expect("body");

    assert_eq!(&body[..], b"req-handler-1");
}

#[test]
fn http_traceparent_is_parsed_without_forging_trace_context() {
    let request = axum::http::Request::builder()
        .uri("/ready")
        .header(
            "traceparent",
            "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
        )
        .body(())
        .expect("request");
    let traceparent = traceparent_from_headers(request.headers()).expect("traceparent");

    assert_eq!(
        trace_id_from_traceparent(&traceparent),
        Some("4bf92f3577b34da6a3ce929d0e0e4736")
    );
    assert!(current_trace_id().is_none());
}

#[test]
fn http_correlation_context_uses_low_cardinality_route_fields() {
    let request = axum::http::Request::builder()
        .uri("/users/42?token=raw-secret")
        .header("x-request-id", "req-http-1")
        .header(
            "traceparent",
            "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
        )
        .body(())
        .expect("request");

    let context = CorrelationContext::from_http_headers(
        Some("users"),
        "GET",
        Some("/users/:id"),
        request.headers(),
    );
    let pairs = context.as_pairs();

    assert!(pairs.contains(&("service".to_string(), "users".to_string())));
    assert!(pairs.contains(&("transport".to_string(), "http".to_string())));
    assert!(pairs.contains(&("route".to_string(), "/users/:id".to_string())));
    assert!(pairs.contains(&("method".to_string(), "GET".to_string())));
    assert!(pairs.contains(&("request_id".to_string(), "req-http-1".to_string())));
    assert!(pairs.contains(&(
        "trace_id".to_string(),
        "4bf92f3577b34da6a3ce929d0e0e4736".to_string()
    )));
    assert!(pairs.contains(&("span_id".to_string(), "00f067aa0ba902b7".to_string())));
    assert!(!pairs.iter().any(|(_, value)| value.contains("42")));
    assert!(!pairs.iter().any(|(_, value)| value.contains("raw-secret")));
}

#[test]
fn raw_http_paths_are_not_used_as_correlation_routes() {
    let request = axum::http::Request::builder()
        .uri("/users/42?api_key=raw-secret")
        .body(())
        .expect("request");

    let context =
        CorrelationContext::from_http_headers(Some("users"), "GET", None, request.headers());

    assert_eq!(context.route(), "unknown");
    assert!(
        !context
            .as_pairs()
            .iter()
            .any(|(_, value)| value == "/users/42")
    );
}

#[cfg(feature = "otlp")]
#[test]
fn http_traceparent_extracts_valid_opentelemetry_parent() {
    use opentelemetry::global;
    use opentelemetry::trace::TraceContextExt;
    use opentelemetry_sdk::propagation::TraceContextPropagator;
    use rs_zero::observability::opentelemetry_context_from_headers;

    global::set_text_map_propagator(TraceContextPropagator::new());
    let request = axum::http::Request::builder()
        .uri("/ready")
        .header(
            "traceparent",
            "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
        )
        .body(())
        .expect("request");

    let context = opentelemetry_context_from_headers(request.headers()).expect("context");

    assert_eq!(
        context.span().span_context().trace_id().to_string(),
        "4bf92f3577b34da6a3ce929d0e0e4736"
    );
}

#[cfg(feature = "rpc")]
#[test]
fn rpc_request_id_metadata_is_available_for_trace_correlation() {
    use rs_zero::observability::{
        insert_traceparent_metadata, request_id_from_metadata, traceparent_from_metadata,
    };
    use rs_zero::rpc::request_id_interceptor;
    use tonic::{Request, service::Interceptor};

    let mut interceptor = request_id_interceptor();
    let request = interceptor.call(Request::new(())).expect("request");

    assert!(request_id_from_metadata(request.metadata()).is_some());

    let mut request = Request::new(());
    insert_traceparent_metadata(
        request.metadata_mut(),
        "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
    )
    .expect("traceparent");
    assert!(traceparent_from_metadata(request.metadata()).is_some());
}

#[cfg(feature = "rpc")]
#[test]
fn rpc_correlation_context_uses_metadata_and_method_pattern() {
    use rs_zero::observability::insert_traceparent_metadata;
    use tonic::Request;

    let mut request = Request::new(());
    request
        .metadata_mut()
        .insert("x-request-id", "req-rpc-1".parse().expect("request id"));
    insert_traceparent_metadata(
        request.metadata_mut(),
        "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
    )
    .expect("traceparent");

    let context =
        CorrelationContext::from_rpc_metadata("billing.Billing", "GetInvoice", request.metadata());
    let pairs = context.as_pairs();

    assert!(pairs.contains(&("service".to_string(), "billing.Billing".to_string())));
    assert!(pairs.contains(&("transport".to_string(), "grpc".to_string())));
    assert!(pairs.contains(&("route".to_string(), "GetInvoice".to_string())));
    assert!(pairs.contains(&("method".to_string(), "GetInvoice".to_string())));
    assert!(pairs.contains(&("request_id".to_string(), "req-rpc-1".to_string())));
    assert_eq!(
        span_id_from_traceparent(context.traceparent().expect("traceparent")),
        Some("00f067aa0ba902b7")
    );
    assert!(!pairs.iter().any(|(_, value)| value.contains("token")));
}

#[cfg(feature = "rpc")]
#[tokio::test]
async fn rpc_observe_helper_reads_correlation_from_metadata() {
    use rs_zero::observability::{insert_traceparent_metadata, observe_rpc_unary_with_metadata};
    use tonic::Request;

    let registry = MetricsRegistry::new();
    let mut request = Request::new(());
    request.metadata_mut().insert(
        "x-request-id",
        "req-rpc-metadata-1".parse().expect("request id"),
    );
    insert_traceparent_metadata(
        request.metadata_mut(),
        "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
    )
    .expect("traceparent");

    let value = observe_rpc_unary_with_metadata(
        Some(&registry),
        "hello.Hello",
        "GetHello",
        request.metadata(),
        async { Ok::<_, tonic::Status>("ok") },
    )
    .await
    .expect("rpc");

    assert_eq!(value, "ok");
    let text = registry.render_prometheus();
    assert!(text.contains("service=\"hello.Hello\",method=\"GetHello\",code=\"Ok\""));
}

#[cfg(all(feature = "otlp", feature = "rpc"))]
#[test]
fn rpc_traceparent_extracts_valid_opentelemetry_parent() {
    use opentelemetry::global;
    use opentelemetry::trace::TraceContextExt;
    use opentelemetry_sdk::propagation::TraceContextPropagator;
    use rs_zero::observability::{
        insert_traceparent_metadata, opentelemetry_context_from_metadata,
    };
    use tonic::Request;

    global::set_text_map_propagator(TraceContextPropagator::new());
    let mut request = Request::new(());
    insert_traceparent_metadata(
        request.metadata_mut(),
        "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
    )
    .expect("traceparent");

    let context = opentelemetry_context_from_metadata(request.metadata()).expect("context");

    assert_eq!(
        context.span().span_context().trace_id().to_string(),
        "4bf92f3577b34da6a3ce929d0e0e4736"
    );
}