add context to message handling

gateway/src/node/client_handling/websocket/connection_handler/authenticated.rs

@@ -14,6 +14,7 @@ use futures::{
     future::{FusedFuture, OptionFuture},
     FutureExt, StreamExt,
 };
+use nym_bin_common::opentelemetry::context::ContextCarrier;
 use nym_credential_verification::CredentialVerifier;
 use nym_credential_verification::{
     bandwidth_storage_manager::BandwidthStorageManager, ClientBandwidth,
@@ -30,10 +31,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;
+use opentelemetry::propagation::TextMapPropagator;
 use rand::{random, CryptoRng, Rng};
+use tracing_opentelemetry::OpenTelemetrySpanExt;
 use std::{process, time::Duration};
 use thiserror::Error;
-use tokio::io::{AsyncRead, AsyncWrite};
+use tokio::{io::{AsyncRead, AsyncWrite}, time::error::Elapsed};
 use tokio_tungstenite::tungstenite::{protocol::Message, Error as WsError};
 use tracing::*;
 
@@ -594,7 +598,29 @@ impl<R, S> AuthenticatedHandler<R, S> {
         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 propagator = TraceContextPropagator::new();
+            let extracted_cx = propagator.extract(&carrier);
+
+            tracing::Span::current().set_parent(extracted_cx);
+            let span = info_span!("starting point with otel");
+            warn!("=== Context propagated to authenticated client handler listen_for_requests ===");
+            Some(span)
+        } else {
+            None
+        };
+
+        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();
+        
         // Ping timeout future used to check if the client responded to our ping request
         let mut ping_timeout: OptionFuture<_> = None.into();
 

gateway/src/node/client_handling/websocket/connection_handler/fresh.rs

@@ -644,6 +644,7 @@ impl<R, S> FreshHandler<R, S> {
                 address,
                 shared_keys.key,
                 session_request_start,
+                None,
             )),
             ServerResponse::Authenticate {
                 protocol_version: Some(negotiated_protocol),
@@ -659,6 +660,7 @@ impl<R, S> FreshHandler<R, S> {
     async fn handle_authenticate_v2(
         &mut self,
         request: Box<AuthenticateRequest>,
+        otel_context: Option<std::collections::HashMap<String, String>>,
     ) -> Result<InitialAuthResult, InitialAuthenticationError>
     where
         S: AsyncRead + AsyncWrite + Unpin,
@@ -734,6 +736,7 @@ impl<R, S> FreshHandler<R, S> {
                 address,
                 shared_key.key,
                 session_request_start,
+                otel_context,
             )),
             ServerResponse::Authenticate {
                 protocol_version: Some(negotiated_protocol),
@@ -832,6 +835,7 @@ impl<R, S> FreshHandler<R, S> {
             remote_address,
             shared_keys,
             OffsetDateTime::now_utc(),
+            None
         );
 
         Ok(InitialAuthResult::new(
@@ -901,6 +905,11 @@ impl<R, S> FreshHandler<R, S> {
         };
         let _enter = child_span.enter();
 
+        let context_carrier = match &request {
+            ClientControlRequest::AuthenticateV2(ref auth_req) => auth_req.otel_context.clone(),
+            _ => None,
+        };
+
         let auth_result = match request {
             ClientControlRequest::Authenticate {
                 protocol_version,
@@ -912,7 +921,7 @@ impl<R, S> FreshHandler<R, S> {
                 self.handle_legacy_authenticate(protocol_version, address, enc_address, iv)
                     .await
             }
-            ClientControlRequest::AuthenticateV2(req) => self.handle_authenticate_v2(req).await,
+            ClientControlRequest::AuthenticateV2(req) => self.handle_authenticate_v2(req, context_carrier).await,
                 ClientControlRequest::RegisterHandshakeInitRequest {
                 protocol_version,
                 data,
@@ -958,6 +967,7 @@ impl<R, S> FreshHandler<R, S> {
             warn!("could not establish client details");
             return Err(InitialAuthenticationError::EmptyClientDetails);
         };
+
         Ok((Some(client_details), trace_id))
     }
 

gateway/src/node/client_handling/websocket/connection_handler/mod.rs

@@ -54,6 +54,7 @@ pub(crate) struct ClientDetails {
     // note, this does **NOT ALWAYS** indicate timestamp of when client connected
     // it is (for v2 auth) timestamp the client **signed** when it created the request
     pub(crate) session_request_timestamp: OffsetDateTime,
+    pub(crate) otel_context: Option<std::collections::HashMap<String, String>>,
 }
 
 impl ClientDetails {
@@ -62,12 +63,14 @@ impl ClientDetails {
         address: DestinationAddressBytes,
         shared_keys: SharedGatewayKey,
         session_request_timestamp: OffsetDateTime,
+        otel_context: Option<std::collections::HashMap<String, String>>,
     ) -> Self {
         ClientDetails {
             address,
             id,
             shared_keys,
             session_request_timestamp,
+            otel_context,
         }
     }
 }