From 28058944f280f2605ea9bb6e2dc8500f6b2a8d33 Mon Sep 17 00:00:00 2001 From: Floriane TUERNAL SABOTINOV Date: Wed, 10 Sep 2025 14:48:28 +0200 Subject: [PATCH] cleanup --- .../bin-common/src/opentelemetry/context.rs | 73 +++++++++++++-- .../src/types/text_request/mod.rs | 21 +---- .../connection_handler/authenticated.rs | 89 +++++-------------- .../websocket/connection_handler/fresh.rs | 89 ++++++++++++------- .../websocket/connection_handler/mod.rs | 6 +- sdk/rust/nym-sdk/examples/surb_reply.rs | 59 ++---------- 6 files changed, 155 insertions(+), 182 deletions(-) diff --git a/common/bin-common/src/opentelemetry/context.rs b/common/bin-common/src/opentelemetry/context.rs index 35f3357c01..cfe2a0d2c4 100644 --- a/common/bin-common/src/opentelemetry/context.rs +++ b/common/bin-common/src/opentelemetry/context.rs @@ -1,5 +1,8 @@ -use opentelemetry::propagation::{Injector, Extractor}; -use opentelemetry::trace::TraceId; +use opentelemetry::{Context, TraceFlags}; +use opentelemetry::propagation::{Injector, Extractor, TextMapPropagator}; +use opentelemetry::trace::{SpanContext, TraceContextExt, TraceId}; +use opentelemetry_sdk::propagation::TraceContextPropagator; +use opentelemetry_sdk::trace::IdGenerator; use std::collections::HashMap; /// Make a Carrier for context propagation @@ -8,12 +11,25 @@ pub struct ContextCarrier { } impl ContextCarrier { - pub fn new() -> Self { + pub fn new_empty() -> Self { ContextCarrier { data: HashMap::new(), } } + pub fn new_with_current_context(context: Context) -> Self { + let propagator = TraceContextPropagator::new(); + let mut carrier = ContextCarrier::new_empty(); + propagator.inject_context(&context, &mut carrier); + carrier + } + + pub fn new_with_extracted_context(external_context: &ContextCarrier) -> Self { + let extractor = TraceContextPropagator::new(); + let context = extractor.extract(external_context); + ContextCarrier::new_with_current_context(context) + } + pub fn from_map(data: HashMap) -> Self { ContextCarrier { data } } @@ -50,10 +66,49 @@ impl Extractor for ContextCarrier { } } -// TODO use generics to extract context from request - // fn extract_context_from_request(request: ) -> Context { - // global::get_text_map_propagator(|propagator| { - // propagator.extract(&HeaderExtractor(&request.headers)) - // }) - // } +pub struct AsyncSpanContextExt { + pub context_carrier: ContextCarrier, + pub root_span: tracing::Span, + pub trace_id: Option, +} +impl AsyncSpanContextExt { + pub fn new() -> Self { + AsyncSpanContextExt { + context_carrier: ContextCarrier::new_empty(), + root_span: tracing::Span::none(), + trace_id: None, + } + } + + pub fn with_context_carrier(mut self, carrier: ContextCarrier) -> Self { + self.context_carrier = carrier; + self.trace_id = self.context_carrier.extract_trace_id(); + self + } + + pub fn with_extracted_context(mut self, external_context: &ContextCarrier) -> Self { + self.context_carrier = ContextCarrier::new_with_extracted_context(external_context); + self.trace_id = self.context_carrier.extract_trace_id(); + self + } + + pub fn set_root_span(mut self, span: tracing::Span) -> Self { + self.root_span = span; + self + } + } + +pub fn new_span_context_with_id(trace_id: TraceId) -> Context { + let id_gen = opentelemetry_sdk::trace::RandomIdGenerator::default(); + let span_id = id_gen.new_span_id(); + let span_context = SpanContext::new( + trace_id, + span_id, + TraceFlags::SAMPLED, + true, + Default::default(), + ); + + Context::current().with_remote_span_context(span_context) +} \ No newline at end of file diff --git a/common/gateway-requests/src/types/text_request/mod.rs b/common/gateway-requests/src/types/text_request/mod.rs index 762ee71d6d..9bf89c52fe 100644 --- a/common/gateway-requests/src/types/text_request/mod.rs +++ b/common/gateway-requests/src/types/text_request/mod.rs @@ -8,6 +8,7 @@ use crate::{ AUTHENTICATE_V2_PROTOCOL_VERSION, CREDENTIAL_UPDATE_V2_PROTOCOL_VERSION, INITIAL_PROTOCOL_VERSION, }; +use nym_bin_common::opentelemetry::context::ContextCarrier; use nym_credentials_interface::CredentialSpendingData; use nym_crypto::asymmetric::ed25519; use nym_sphinx::DestinationAddressBytes; @@ -17,9 +18,6 @@ use std::collections::HashMap; use std::str::FromStr; use tracing::{instrument, warn}; use tungstenite::Message; -use nym_bin_common::opentelemetry::context::ContextCarrier; -use opentelemetry::propagation::TextMapPropagator; -use opentelemetry_sdk::propagation::TraceContextPropagator; pub mod authenticate; @@ -137,17 +135,8 @@ impl ClientControlRequest { let nonce = shared_key.random_nonce_or_iv(); let ciphertext = shared_key.encrypt_naive(address.as_bytes_ref(), Some(&nonce))?; - use nym_bin_common::opentelemetry::context::ContextCarrier; - use opentelemetry::propagation::TextMapPropagator; - use opentelemetry_sdk::propagation::TraceContextPropagator; - let context = opentelemetry::Context::current(); - let propagator = TraceContextPropagator::new(); - let mut carrier = ContextCarrier::new(); - propagator.inject_context(&context, &mut carrier); - let context_carrier = carrier.into_map(); - - tracing::error!("context_carrier is {:?}", context_carrier); + let context_carrier = ContextCarrier::new_with_current_context(context).into_map(); Ok(ClientControlRequest::Authenticate { protocol_version, @@ -167,11 +156,7 @@ impl ClientControlRequest { let protocol_version = AUTHENTICATE_V2_PROTOCOL_VERSION; let otel_context = opentelemetry::Context::current(); - let propagator = TraceContextPropagator::new(); - let mut carrier =ContextCarrier::new(); - propagator.inject_context(&otel_context, &mut carrier); - let context_carrier = carrier.into_map(); - tracing::error!("carrier is {:?}", context_carrier); + let context_carrier = ContextCarrier::new_with_current_context(otel_context).into_map(); Ok(ClientControlRequest::AuthenticateV2(Box::new( AuthenticateRequest::new( diff --git a/gateway/src/node/client_handling/websocket/connection_handler/authenticated.rs b/gateway/src/node/client_handling/websocket/connection_handler/authenticated.rs index 49d2644417..6d334bd4ac 100644 --- a/gateway/src/node/client_handling/websocket/connection_handler/authenticated.rs +++ b/gateway/src/node/client_handling/websocket/connection_handler/authenticated.rs @@ -14,11 +14,12 @@ use futures::{ future::{FusedFuture, OptionFuture}, FutureExt, StreamExt, }; -use nym_bin_common::opentelemetry::context::ContextCarrier; +use nym_bin_common::opentelemetry::context::{new_span_context_with_id, AsyncSpanContextExt, ContextCarrier}; use nym_credential_verification::CredentialVerifier; use nym_credential_verification::{ bandwidth_storage_manager::BandwidthStorageManager, ClientBandwidth, }; +use nym_crypto::shared_key::new_ephemeral_shared_key; use nym_gateway_requests::{ types::{BinaryRequest, ServerResponse}, ClientControlRequest, ClientRequest, GatewayRequestsError, SensitiveServerResponse, @@ -31,13 +32,13 @@ use nym_node_metrics::events::MetricsEvent; use nym_sphinx::forwarding::packet::MixPacket; use nym_statistics_common::{gateways::GatewaySessionEvent, types::SessionType}; use nym_validator_client::coconut::EcashApiError; -use opentelemetry_sdk::{propagation::TraceContextPropagator, trace::IdGenerator}; -use opentelemetry::{propagation::TextMapPropagator, trace::{SpanContext, TraceContextExt, Tracer}, TraceFlags}; +use opentelemetry::{propagation::TextMapPropagator, trace::TraceContextExt}; +use opentelemetry_sdk::propagation::TraceContextPropagator; use rand::{random, CryptoRng, Rng}; use tracing_opentelemetry::OpenTelemetrySpanExt; use std::{process, time::Duration}; use thiserror::Error; -use tokio::{io::{AsyncRead, AsyncWrite}, time::error::Elapsed}; +use tokio::{io::{AsyncRead, AsyncWrite}, time::{error::Elapsed, interval}}; use tokio_tungstenite::tungstenite::{protocol::Message, Error as WsError}; use tracing::*; @@ -598,74 +599,30 @@ impl AuthenticatedHandler { S: AsyncRead + AsyncWrite + Unpin, { trace!("Started listening for ALL incoming requests..."); - let remote_span = if let Some(ctx) = &self.client.otel_context { - let carrier = ContextCarrier::from_map(ctx.clone()); - let extracted_trace_id = carrier.extract_trace_id(); - warn!("=== Listen for requests: Extracted trace id from client: {:?} ===", extracted_trace_id); + let context_ext = if let Some(ctx) = &self.client.otel_context { + //retrieve the context from the client + let context_ext = AsyncSpanContextExt::new() + .with_context_carrier(ContextCarrier::from_map(ctx.clone())); - let id_generator = opentelemetry_sdk::trace::RandomIdGenerator::default(); - let span_id = id_generator.new_span_id(); - let trace_id = extracted_trace_id.unwrap_or_else(|| { - warn!("Failed to extract trace id from the provided otel context carrier"); - let context = opentelemetry::Context::current(); - context.span().span_context().trace_id() - }); + // extract the parent context from the carrier + let extractor = TraceContextPropagator::new(); + let parent_cx = extractor.extract(&context_ext.context_carrier); + let trace_id = parent_cx.span().span_context().trace_id(); - let span_context = SpanContext::new( - trace_id, - span_id, - TraceFlags::SAMPLED, - true, - Default::default(), - ); - let cx = opentelemetry::Context::current().with_remote_span_context(span_context); - let _context_guard = cx.clone().attach(); - let span = info_span!("starting point with otel", trace_id = %trace_id); - span.set_parent(cx); + let span_cx = new_span_context_with_id(trace_id); + let _guard = span_cx.clone().attach(); - Some(span) + let span = info_span!("Started listening for incomming requests with otel", %trace_id); + span.set_parent(span_cx); + let context_ext = context_ext.set_root_span(span); + context_ext } else { - None + AsyncSpanContextExt::new() }; - let child_span = if let Some(ref remote_span) = remote_span { - info_span!(parent: remote_span, "listen_for_requests with otel") - } else { - info_span!("listen_for_requests without otel") - }; - let _child_span_guard = child_span.enter(); + let _guard = context_ext.root_span.enter(); - // Note: the commented out code below is an alternative way of propagating the context - // it does not work because of the way we spawn the task - // i.e. `tokio::spawn(async move { ... })` prevents the context from being propagated - // since the `async move` creates a new future that does not capture the current context - // so for now we just manually extract the trace id from the client and create a new span with it - - // if we ever change the way we spawn the task, we can try using the code below instead - // let span = if let Some(remote_ctx) = &self.client.otel_context { - // let carrier = ContextCarrier::from_map(remote_ctx.clone()); - // let extracted_trace_id = carrier.extract_trace_id().unwrap_or_else(|| { - // warn!("Failed to extract trace id from the provided otel context carrier"); - // let context = opentelemetry::Context::current(); - // context.span().span_context().trace_id() - // }); - // warn!("=== Listen for requests: Extracted trace id from client: {:?} ===", extracted_trace_id); - // let tracer = opentelemetry::global::tracer("nym-ws-listener"); - // tracer - // .span_builder("websocket listen_for_requests") - // .with_trace_id(extracted_trace_id) - // .start(&tracer) - // } else { - // let tracer = opentelemetry::global::tracer("nym-ws-listener"); - // tracer - // .span_builder("websocket listen_for_requests") - // .start(&tracer) - // }; - // let context = opentelemetry::Context::current_with_span(span); - // let _guard = context.clone().attach(); - // warn!("=== Context propagated to authenticated client handler listen_for_requests ==="); - // Ping timeout future used to check if the client responded to our ping request let mut ping_timeout: OptionFuture<_> = None.into(); @@ -690,7 +647,7 @@ impl AuthenticatedHandler { self.handle_ping_timeout().await; }, socket_msg = self.inner.read_websocket_message() => { - let span = info_span!(parent: &child_span, "client_message_received"); + let span = info_span!(parent: &context_ext.root_span, "client_message_received"); let _enter = span.enter(); let socket_msg = match socket_msg { None => break, @@ -715,7 +672,7 @@ impl AuthenticatedHandler { } }, mix_messages = self.mix_receiver.next() => { - let span = info_span!(parent: &child_span, "mix_message_received"); + let span = info_span!(parent: &context_ext.root_span, "mix_message_received"); let _enter = span.enter(); let mix_messages = match mix_messages { None => { diff --git a/gateway/src/node/client_handling/websocket/connection_handler/fresh.rs b/gateway/src/node/client_handling/websocket/connection_handler/fresh.rs index a7ad4e1458..890131d7c5 100644 --- a/gateway/src/node/client_handling/websocket/connection_handler/fresh.rs +++ b/gateway/src/node/client_handling/websocket/connection_handler/fresh.rs @@ -13,7 +13,7 @@ use futures::{ channel::{mpsc, oneshot}, SinkExt, StreamExt, }; -use nym_bin_common::opentelemetry::context::ContextCarrier; +use nym_bin_common::opentelemetry::context::{ContextCarrier, AsyncSpanContextExt, new_span_context_with_id}; use nym_credentials_interface::AvailableBandwidth; use nym_crypto::aes::cipher::crypto_common::rand_core::RngCore; use nym_crypto::asymmetric::ed25519; @@ -40,9 +40,9 @@ use opentelemetry::Context; use opentelemetry_sdk::logs::TraceContext; use opentelemetry_sdk::trace::{IdGenerator, RandomIdGenerator}; use opentelemetry::propagation::TextMapPropagator; +use opentelemetry::TraceId; use opentelemetry_sdk::propagation::TraceContextPropagator; use rand::CryptoRng; -use std::collections::HashMap; use std::net::SocketAddr; use std::time::Duration; use thiserror::Error; @@ -875,40 +875,70 @@ impl FreshHandler { S: AsyncRead + AsyncWrite + Unpin + Send, R: CryptoRng + RngCore + Send, { - let (remote_cx_span, trace_id) = if let ClientControlRequest::AuthenticateV2(ref auth_req) = request { + let context_ext = if let ClientControlRequest::AuthenticateV2(ref auth_req) = request { if let Some(otel_context) = &auth_req.otel_context { - let carrier = ContextCarrier::from_map(otel_context.clone()); - - let extracted_trace_id = carrier.extract_trace_id(); - info!("Extracted trace id: {:?}", extracted_trace_id); - - let propagator = TraceContextPropagator::new(); - let extracted_context = propagator.extract(&carrier); + let context_ext = AsyncSpanContextExt::new() + .with_context_carrier(ContextCarrier::from_map(otel_context.clone())); + info!("Extracted trace id: {:?}", context_ext.trace_id); + + let extractor = TraceContextPropagator::new(); + let extracted_context = extractor.extract(&context_ext.context_carrier); tracing::Span::current().set_parent(extracted_context); let span = info_span!("extracted_otel_context"); + let context_ext = context_ext.set_root_span(span); warn!("==== Context propagation successful ===="); - - (Some(span), extracted_trace_id) + + context_ext } else { warn!("No OpenTelemetry context provided in the request"); - (None, None) + AsyncSpanContextExt::new() } } else { - (None, None) + warn!("No OpenTelemetry context provided in the request"); + AsyncSpanContextExt::new() }; - let child_span = if let Some(ref parent_span) = remote_cx_span { - info_span!(parent: parent_span, "handling initial client request with otel context") - } else { + let child_span = if context_ext.root_span.is_none() { info_span!("handling_initial_client_request") + } else { + info_span!(parent: &context_ext.root_span, "handling initial client request with otel context") }; let _enter = child_span.enter(); - let context_carrier = match &request { - ClientControlRequest::AuthenticateV2(ref auth_req) => auth_req.otel_context.clone(), - _ => None, - }; + // let (remote_cx_span, trace_id) = if let ClientControlRequest::AuthenticateV2(ref auth_req) = request { + // if let Some(otel_context) = &auth_req.otel_context { + // let carrier = ContextCarrier::from_map(otel_context.clone()); + + // let extracted_trace_id = carrier.extract_trace_id(); + // info!("Extracted trace id: {:?}", extracted_trace_id); + + // let propagator = TraceContextPropagator::new(); + // let extracted_context = propagator.extract(&carrier); + + // tracing::Span::current().set_parent(extracted_context); + // let span = info_span!("extracted_otel_context"); + // warn!("==== Context propagation successful ===="); + + // (Some(span), extracted_trace_id) + // } else { + // warn!("No OpenTelemetry context provided in the request"); + // (None, None) + // } + // } else { + // (None, None) + // }; + + // let child_span = if let Some(ref parent_span) = remote_cx_span { + // info_span!(parent: parent_span, "handling initial client request with otel context") + // } else { + // info_span!("handling_initial_client_request") + // }; + + // let context_carrier = match &request { + // ClientControlRequest::AuthenticateV2(ref auth_req) => auth_req.otel_context.clone(), + // _ => None, + // }; let auth_result = match request { ClientControlRequest::Authenticate { @@ -921,14 +951,14 @@ impl FreshHandler { self.handle_legacy_authenticate(protocol_version, address, enc_address, iv) .await } - ClientControlRequest::AuthenticateV2(req) => self.handle_authenticate_v2(req, context_carrier).await, + ClientControlRequest::AuthenticateV2(req) => self.handle_authenticate_v2(req, Some(context_ext.context_carrier.into_map())).await, ClientControlRequest::RegisterHandshakeInitRequest { protocol_version, data, } => self.handle_register(protocol_version, data).await, ClientControlRequest::SupportedProtocol { .. } => { self.handle_reply_supported_protocol_request().await; - return Ok((None, trace_id)); + return Ok((None, context_ext.trace_id)); } _ => { debug!("received an invalid client request"); @@ -968,7 +998,7 @@ impl FreshHandler { return Err(InitialAuthenticationError::EmptyClientDetails); }; - Ok((Some(client_details), trace_id)) + Ok((Some(client_details), context_ext.trace_id)) } #[instrument(skip_all)] @@ -1001,16 +1031,7 @@ impl FreshHandler { if let (Some(registration_details), Some(trace_id)) = (maybe_auth_res, maybe_trace_id) { let span = { - let id_gen = opentelemetry_sdk::trace::RandomIdGenerator::default(); - let span_id = id_gen.new_span_id(); - let span_context = SpanContext::new( - trace_id, - span_id, - TraceFlags::SAMPLED, - true, - Default::default(), - ); - let cx = opentelemetry::Context::current().with_remote_span_context(span_context); + let cx = new_span_context_with_id(trace_id); let _context_guard = cx.clone().attach(); let span = info_span!("authentication handler with otel", %trace_id); diff --git a/gateway/src/node/client_handling/websocket/connection_handler/mod.rs b/gateway/src/node/client_handling/websocket/connection_handler/mod.rs index a2d99ddf9c..54fcfc4d02 100644 --- a/gateway/src/node/client_handling/websocket/connection_handler/mod.rs +++ b/gateway/src/node/client_handling/websocket/connection_handler/mod.rs @@ -2,14 +2,11 @@ // SPDX-License-Identifier: GPL-3.0-only use crate::config::Config; -use nym_bin_common::opentelemetry::context::ContextCarrier; +use nym_bin_common::opentelemetry::context::new_span_context_with_id; use nym_credential_verification::BandwidthFlushingBehaviourConfig; use nym_gateway_requests::shared_key::SharedGatewayKey; use nym_gateway_requests::ServerResponse; use nym_sphinx::DestinationAddressBytes; -use opentelemetry::propagation::TextMapPropagator; -use opentelemetry_sdk::propagation::TraceContextPropagator; -use opentelemetry_sdk::trace::IdGenerator; use rand::{CryptoRng, Rng}; use std::time::Duration; use time::OffsetDateTime; @@ -17,7 +14,6 @@ use tokio::io::{AsyncRead, AsyncWrite}; use tokio_tungstenite::WebSocketStream; use tracing::{debug, instrument, trace, warn, info_span}; use tracing_opentelemetry::OpenTelemetrySpanExt; -use opentelemetry::trace::{SpanContext, TraceContextExt, TraceFlags}; pub(crate) use self::authenticated::AuthenticatedHandler; pub(crate) use self::fresh::FreshHandler; diff --git a/sdk/rust/nym-sdk/examples/surb_reply.rs b/sdk/rust/nym-sdk/examples/surb_reply.rs index 081bedbdcc..02c918a605 100644 --- a/sdk/rust/nym-sdk/examples/surb_reply.rs +++ b/sdk/rust/nym-sdk/examples/surb_reply.rs @@ -1,72 +1,30 @@ use nym_sdk::mixnet::{ AnonymousSenderTag, MixnetClientBuilder, MixnetMessageSender, ReconstructedMessage, - // StoragePaths, }; use opentelemetry::trace::{TraceContextExt, Tracer}; use opentelemetry::{global, Context}; -// use opentelemetry::trace::Tracer; -// use opentelemetry::Context; -// use opentelemetry::global; use tracing::warn; -use tracing::{instrument, info_span}; -use tracing_opentelemetry::OpenTelemetrySpanExt; -use std::path::PathBuf; -use tempfile::TempDir; -use nym_sdk::mixnet::StoragePaths; +use tracing::instrument; #[tokio::main] #[instrument(name = "sdk-example-surb-reply", skip_all)] async fn main() { nym_bin_common::opentelemetry::setup_tracing_logger("sdk-example-surb-reply".to_string()).unwrap(); - // let main_span = info_span!("surb_example_session"); - // let _main_span_enter = main_span.enter(); - - // let current_span = tracing::Span::current(); - // let otel_context = current_span.context(); - // let binding = otel_context.span(); - // let span_context = binding.span_context(); - // let trace_id = span_context.trace_id(); - - // warn!("Starting the SURB reply example - trace id: {}", trace_id); - // warn!("Otel context: {:?}", otel_context); - // warn!("trace id: {}", trace_id); let tracer = global::tracer("sdk-example-surb-reply"); - let span = tracer.start("test_span"); + let span = tracer.start("client-root-span"); let cx = Context::current_with_span(span); let _guard = cx.clone().attach(); let trace_id = cx.span().span_context().trace_id(); warn!("Main TRACE_ID: {:?}", trace_id); - let span = info_span!( - "surb_reply_example_session", - trace_id = %trace_id.to_string() - ); - let _enter = span.enter(); - - let otel_context = opentelemetry::Context::current(); - warn!("OTEL CONTEXT: {:?}", otel_context); - let span = otel_context.span(); - let context = span.span_context(); - let trace_id = context.trace_id(); - warn!("TRACE_ID: {:?}", trace_id); + // let span = info_span!( + // "surb_reply_example_session", + // trace_id = %trace_id.to_string() + // ); + // let _enter = span.enter(); - // Specify some config options - let config_dir: PathBuf = TempDir::new().unwrap().path().to_path_buf(); - let storage_paths = StoragePaths::new_from_dir(&config_dir).unwrap(); - - // Create the client with a storage backend, and enable it by giving it some paths. If keys - // exists at these paths, they will be loaded, otherwise they will be generated. - // let client = MixnetClientBuilder::new_with_default_storage(storage_paths) - // .await - // .unwrap() - // .build() - // .unwrap(); - - // Now we connect to the mixnet, using keys now stored in the paths provided. - // let mut client = client.connect_to_mixnet().await.unwrap(); - // Create a mixnet client which connect to a specific node let client_builder = MixnetClientBuilder::new_ephemeral(); @@ -75,10 +33,11 @@ async fn main() { .build() .unwrap(); - let mut client = mixnet_client.connect_to_mixnet().await.unwrap(); // Now we connect to the mixnet, using keys now stored in the paths provided. // let mut client = client.connect_to_mixnet().await.unwrap(); + let mut client = mixnet_client.connect_to_mixnet().await.unwrap(); + // Be able to get our client address let our_address = client.nym_address(); println!("\nOur client nym address is: {our_address}");