alibabacloud-rum 0.1.0

mod ids;
mod sw8;
mod trace_context;

use reqwest::header::HeaderMap;
use url::Url;

use crate::config::RumConfig;

pub(crate) use ids::{generate_span_id, generate_trace_id};
pub(crate) use sw8::{decode_sw8, encode_sw8, encode_sw8_with_ids};
pub(crate) use trace_context::parse_traceparent;

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(crate) enum PropagationMode {
    TraceParent,
    Sw8,
}

#[derive(Clone, Debug, Default)]
pub(crate) struct TraceInfo {
    pub trace_id: Option<String>,
    pub span_id: Option<String>,
    pub headers: Vec<(String, String)>,
}

pub(crate) fn inject_or_extract(
    headers: &mut HeaderMap,
    config: &RumConfig,
    method: &str,
    url: &Url,
    should_inject: bool,
    instance_id: &str,
) -> TraceInfo {
    if let Some(value) = headers.get("traceparent").and_then(|v| v.to_str().ok()) {
        if let Some(parsed) = parse_traceparent(value) {
            return TraceInfo {
                trace_id: Some(parsed.trace_id),
                span_id: Some(parsed.span_id),
                headers: vec![("traceparent".to_string(), value.to_string())],
            };
        }
    }
    if let Some(value) = headers.get("sw8").and_then(|v| v.to_str().ok()) {
        if let Some(parsed) = decode_sw8(value) {
            return TraceInfo {
                trace_id: Some(parsed.trace_id),
                span_id: None,
                headers: vec![("sw8".to_string(), value.to_string())],
            };
        }
    }
    if !should_inject {
        return TraceInfo::default();
    }

    match config.network.propagation_mode {
        PropagationMode::TraceParent => inject_traceparent(headers, config),
        PropagationMode::Sw8 => inject_sw8(headers, config, method, url, instance_id),
    }
}

fn inject_traceparent(headers: &mut HeaderMap, config: &RumConfig) -> TraceInfo {
    let trace_id = generate_trace_id();
    let span_id = generate_span_id();
    let traceparent = format!("00-{trace_id}-{span_id}-01");
    let tracestate = format!(
        "rum=v2,app_id={},sdk_version={},instrumentation=reqwest",
        config.app_id(),
        crate::config::SDK_VERSION
    );
    if let Ok(value) = traceparent.parse() {
        headers.insert("traceparent", value);
    }
    if let Ok(value) = tracestate.parse() {
        headers.insert("tracestate", value);
    }
    TraceInfo {
        trace_id: Some(trace_id),
        span_id: Some(span_id),
        headers: vec![
            ("traceparent".to_string(), traceparent),
            ("tracestate".to_string(), tracestate),
        ],
    }
}

fn inject_sw8(
    headers: &mut HeaderMap,
    config: &RumConfig,
    method: &str,
    url: &Url,
    instance_id: &str,
) -> TraceInfo {
    let trace_id = generate_trace_id();
    let endpoint = format!("HTTP {}", method.to_ascii_uppercase());
    let target = effective_target_address(url);
    let sw8 = encode_sw8(
        &trace_id,
        config.app_name_or_id(),
        instance_id,
        &endpoint,
        &target,
    );
    if let Ok(value) = sw8.parse() {
        headers.insert("sw8", value);
    }
    TraceInfo {
        trace_id: Some(trace_id),
        span_id: None,
        headers: vec![("sw8".to_string(), sw8)],
    }
}

pub(crate) fn effective_target_address(url: &Url) -> String {
    let host = url.host_str().unwrap_or_default();
    let port = url.port_or_known_default().unwrap_or(0);
    if port == 0 {
        host.to_string()
    } else {
        format!("{host}:{port}")
    }
}