From fa603be538c0330fc51e038c39a707740aeef7fa Mon Sep 17 00:00:00 2001
From: Floriane TUERNAL SABOTINOV <ftuertinov@proton.me>
Date: Wed, 27 Aug 2025 17:15:38 +0200
Subject: [PATCH] setup otel from run to forward_sphinx_packet

---
 .../websocket/connection_handler/authenticated.rs      | 10 ++++++++--
 .../websocket/connection_handler/fresh.rs              |  7 ++++++-
 .../websocket/connection_handler/mod.rs                |  2 +-
 gateway/src/node/client_handling/websocket/listener.rs |  3 +++
 4 files changed, 18 insertions(+), 4 deletions(-)

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 36b8029332..1111456fec 100644
--- a/gateway/src/node/client_handling/websocket/connection_handler/authenticated.rs
+++ b/gateway/src/node/client_handling/websocket/connection_handler/authenticated.rs
@@ -305,6 +305,7 @@ impl<R, S> AuthenticatedHandler<R, S> {
     /// # Arguments
     ///
     /// * `bin_msg`: raw message to handle.
+    #[instrument(skip_all)]
     async fn handle_binary(&mut self, bin_msg: Vec<u8>) -> Message {
         trace!("binary request");
         // this function decrypts the request and checks the MAC
@@ -582,6 +583,11 @@ impl<R, S> AuthenticatedHandler<R, S> {
     /// Simultaneously listens for incoming client requests, which realistically should only be
     /// binary requests to forward sphinx packets or increase bandwidth
     /// and for sphinx packets received from the mix network that should be sent back to the client.
+    #[instrument(level = "debug", skip_all,
+        fields(
+            client = %self.client.address.as_base58_string()
+        )
+    )]
     pub(crate) async fn listen_for_requests(mut self)
     where
         R: Rng + CryptoRng,
@@ -609,8 +615,8 @@ impl<R, S> AuthenticatedHandler<R, S> {
                 },
                 // The ping timeout expired, meaning the client didn't respond to our ping request
                 _ = &mut ping_timeout, if !ping_timeout.is_terminated() => {
-                   ping_timeout = None.into();
-                   self.handle_ping_timeout().await;
+                    ping_timeout = None.into();
+                    self.handle_ping_timeout().await;
                 },
                 socket_msg = self.inner.read_websocket_message() => {
                     let socket_msg = match socket_msg {
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 db65ac509c..e14c9a9a7c 100644
--- a/gateway/src/node/client_handling/websocket/connection_handler/fresh.rs
+++ b/gateway/src/node/client_handling/websocket/connection_handler/fresh.rs
@@ -641,6 +641,9 @@ impl<R, S> FreshHandler<R, S> {
         ))
     }
 
+    #[instrument(skip_all, fields(
+        address = %request.content.client_identity.derive_destination_address(),
+    ))]
     async fn handle_authenticate_v2(
         &mut self,
         request: Box<AuthenticateRequest>,
@@ -845,7 +848,7 @@ impl<R, S> FreshHandler<R, S> {
             debug!("failed to reply with protocol version: {err}")
         }
     }
-
+#[instrument(skip_all)]
     pub(crate) async fn handle_initial_client_request(
         &mut self,
         request: ClientControlRequest,
@@ -916,6 +919,7 @@ impl<R, S> FreshHandler<R, S> {
         Ok(Some(client_details))
     }
 
+    #[instrument(skip_all)]
     pub(crate) async fn handle_until_authenticated_or_failure(
         mut self,
     ) -> Option<AuthenticatedHandler<R, S>>
@@ -1017,6 +1021,7 @@ impl<R, S> FreshHandler<R, S> {
             .map_err(|_| InitialAuthenticationError::InvalidRequest)
     }
 
+    #[instrument(skip_all)]
     pub(crate) async fn start_handling(self)
     where
         S: AsyncRead + AsyncWrite + Unpin + Send,
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 2fd97600a7..6f98b64275 100644
--- a/gateway/src/node/client_handling/websocket/connection_handler/mod.rs
+++ b/gateway/src/node/client_handling/websocket/connection_handler/mod.rs
@@ -92,7 +92,7 @@ impl InitialAuthResult {
 }
 
 // imo there's no point in including the peer address in anything higher than debug
-#[instrument(level = "debug", skip_all, fields(peer = %handle.peer_address))]
+#[instrument(level = "debug", skip_all)]
 pub(crate) async fn handle_connection<R, S>(mut handle: FreshHandler<R, S>)
 where
     R: Rng + CryptoRng + Send,
diff --git a/gateway/src/node/client_handling/websocket/listener.rs b/gateway/src/node/client_handling/websocket/listener.rs
index b2d234d4a7..180b488bd6 100644
--- a/gateway/src/node/client_handling/websocket/listener.rs
+++ b/gateway/src/node/client_handling/websocket/listener.rs
@@ -53,6 +53,7 @@ impl Listener {
         )
     }
 
+    #[instrument(skip_all)]
     fn try_handle_accepted_connection(&self, accepted: io::Result<(TcpStream, SocketAddr)>) {
         match accepted {
             Ok((socket, remote_address)) => {
@@ -83,6 +84,7 @@ impl Listener {
                     .network
                     .new_ingress_websocket_client();
 
+                
                 // 4. spawn the task handling the client connection
                 self.shutdown.try_spawn_named(
                     async move {
@@ -104,6 +106,7 @@ impl Listener {
 
     // TODO: change the signature to pub(crate) async fn run(&self, handler: Handler)
 
+    #[instrument(skip_all)]
     pub async fn run(&mut self) {
         info!("Starting websocket listener at {}", self.address);
         let tcp_listener = match tokio::net::TcpListener::bind(self.address).await {