Files
crank/crates/crank-core/src/correlation.rs
T

281 lines
8.1 KiB
Rust

use serde::{Deserialize, Deserializer, Serialize, de};
use uuid::Uuid;
const TRACEPARENT_VERSION: &str = "00";
const ZERO_TRACE_ID: &str = "00000000000000000000000000000000";
const ZERO_PARENT_ID: &str = "0000000000000000";
#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize)]
#[serde(transparent)]
pub struct RequestId(String);
impl RequestId {
pub const MAX_LEN: usize = 128;
pub fn generate() -> Self {
Self(Uuid::now_v7().to_string())
}
pub fn resolve(candidate: Option<&str>) -> Self {
candidate
.filter(|value| Self::is_valid(value))
.map(|value| Self(value.to_owned()))
.unwrap_or_else(Self::generate)
}
pub fn parse(value: &str) -> Result<Self, CorrelationError> {
Self::is_valid(value)
.then(|| Self(value.to_owned()))
.ok_or(CorrelationError::InvalidRequestId)
}
pub fn is_valid(value: &str) -> bool {
!value.is_empty()
&& value.len() <= Self::MAX_LEN
&& value
.bytes()
.all(|byte| byte.is_ascii_graphic() && byte != b',' && byte != b';')
}
pub fn as_str(&self) -> &str {
&self.0
}
}
impl std::fmt::Display for RequestId {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter.write_str(&self.0)
}
}
impl<'de> Deserialize<'de> for RequestId {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: Deserializer<'de>,
{
let value = String::deserialize(deserializer)?;
Self::parse(&value).map_err(de::Error::custom)
}
}
#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize)]
#[serde(transparent)]
pub struct TraceId(String);
impl TraceId {
pub const LEN: usize = 32;
pub fn generate() -> Self {
let value = Uuid::now_v7().simple().to_string();
debug_assert_ne!(value, ZERO_TRACE_ID);
Self(value)
}
pub fn parse(value: &str) -> Result<Self, CorrelationError> {
if is_lower_hex(value, Self::LEN) && value != ZERO_TRACE_ID {
Ok(Self(value.to_owned()))
} else {
Err(CorrelationError::InvalidTraceId)
}
}
pub fn as_str(&self) -> &str {
&self.0
}
}
impl std::fmt::Display for TraceId {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter.write_str(&self.0)
}
}
impl<'de> Deserialize<'de> for TraceId {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: Deserializer<'de>,
{
let value = String::deserialize(deserializer)?;
Self::parse(&value).map_err(de::Error::custom)
}
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
pub struct TraceContext {
trace_id: TraceId,
traceparent: String,
}
impl TraceContext {
pub const TRACEPARENT_LEN: usize = 55;
pub const TRACESTATE_MAX_BYTES: usize = 512;
pub const TRACESTATE_MAX_MEMBERS: usize = 32;
pub const BAGGAGE_MAX_BYTES: usize = 8_192;
pub const BAGGAGE_MAX_MEMBERS: usize = 64;
pub fn generate() -> Self {
let trace_id = TraceId::generate();
let mut parent_id = Uuid::now_v7().simple().to_string()[..16].to_owned();
if parent_id == ZERO_PARENT_ID {
parent_id.replace_range(15..16, "1");
}
// A context generated outside an SDK span must not claim that a sampler
// selected it. Ingress replaces this seed with the actual local span
// context before application code runs.
let traceparent = format!("{TRACEPARENT_VERSION}-{trace_id}-{parent_id}-00");
Self {
trace_id,
traceparent,
}
}
pub fn parse(value: &str) -> Result<Self, CorrelationError> {
if value.len() != Self::TRACEPARENT_LEN {
return Err(CorrelationError::InvalidTraceparent);
}
let bytes = value.as_bytes();
if bytes[2] != b'-' || bytes[35] != b'-' || bytes[52] != b'-' {
return Err(CorrelationError::InvalidTraceparent);
}
let version = &value[0..2];
let trace_id = &value[3..35];
let parent_id = &value[36..52];
let flags = &value[53..55];
if version != TRACEPARENT_VERSION
|| !is_lower_hex(parent_id, 16)
|| parent_id == ZERO_PARENT_ID
|| !matches!(flags, "00" | "01")
{
return Err(CorrelationError::InvalidTraceparent);
}
Ok(Self {
trace_id: TraceId::parse(trace_id).map_err(|_| CorrelationError::InvalidTraceparent)?,
traceparent: value.to_owned(),
})
}
pub fn from_span_parts(
trace_id: &str,
span_id: &str,
sampled: bool,
) -> Result<Self, CorrelationError> {
let flags = if sampled { "01" } else { "00" };
Self::parse(&format!(
"{TRACEPARENT_VERSION}-{trace_id}-{span_id}-{flags}"
))
}
pub fn continue_local(&self) -> Self {
let mut span_id = Uuid::now_v7().simple().to_string()[..16].to_owned();
if span_id == ZERO_PARENT_ID {
span_id.replace_range(15..16, "1");
}
let sampled = self.traceparent.ends_with("-01");
Self::from_span_parts(self.trace_id.as_str(), &span_id, sampled)
.expect("generated span identity is canonical")
}
pub fn trace_id(&self) -> &TraceId {
&self.trace_id
}
pub fn traceparent(&self) -> &str {
&self.traceparent
}
/// Whether this context was selected for recording by the upstream/local sampler.
pub fn is_sampled(&self) -> bool {
self.traceparent.ends_with("-01")
}
pub fn tracestate_within_budget(value: &str) -> bool {
header_list_within_budget(
value,
Self::TRACESTATE_MAX_BYTES,
Self::TRACESTATE_MAX_MEMBERS,
)
}
pub fn baggage_within_budget(value: &str) -> bool {
header_list_within_budget(value, Self::BAGGAGE_MAX_BYTES, Self::BAGGAGE_MAX_MEMBERS)
}
}
impl<'de> Deserialize<'de> for TraceContext {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: Deserializer<'de>,
{
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct WireTraceContext {
trace_id: TraceId,
traceparent: String,
}
let wire = WireTraceContext::deserialize(deserializer)?;
let context = Self::parse(&wire.traceparent).map_err(de::Error::custom)?;
if context.trace_id != wire.trace_id {
return Err(de::Error::custom(CorrelationError::InvalidTraceparent));
}
Ok(context)
}
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct CorrelationContext {
request_id: RequestId,
trace_context: TraceContext,
}
impl CorrelationContext {
pub fn new(request_id: RequestId, trace_context: TraceContext) -> Self {
Self {
request_id,
trace_context,
}
}
pub fn generate() -> Self {
Self::new(RequestId::generate(), TraceContext::generate())
}
pub fn request_id(&self) -> &RequestId {
&self.request_id
}
pub fn trace_context(&self) -> &TraceContext {
&self.trace_context
}
pub fn trace_id(&self) -> &TraceId {
self.trace_context.trace_id()
}
}
#[derive(Clone, Copy, Debug, thiserror::Error, PartialEq, Eq)]
pub enum CorrelationError {
#[error("invalid request identity")]
InvalidRequestId,
#[error("invalid trace identity")]
InvalidTraceId,
#[error("invalid trace parent")]
InvalidTraceparent,
}
fn is_lower_hex(value: &str, expected_len: usize) -> bool {
value.len() == expected_len
&& value
.bytes()
.all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte))
}
fn header_list_within_budget(value: &str, max_bytes: usize, max_members: usize) -> bool {
!value.is_empty()
&& value.len() <= max_bytes
&& value.is_ascii()
&& !value.bytes().any(|byte| byte.is_ascii_control())
&& value.split(',').count() <= max_members
&& value.split(',').all(|member| !member.trim().is_empty())
}