From dbe6a5de7d2c8e892215e39adceaf6fa9b6f7e5b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jon=20H=C3=A4ggblad?= <jon.haggblad@gmail.com>
Date: Wed, 24 Jan 2024 14:54:40 +0100
Subject: [PATCH] Fix reported sizes of the received packets (#4343)

The received packet sizes as reported in the stats should include the
encryption that is decrypted by the client. Note that this does not
include the sphinx encryption, which is already removed by the exit
gateway. This is also the reason for the relatively large discrepancy of
the reported sent and received packets.
---
 .../client-core/src/client/received_buffer.rs | 19 ++++++++++++++-----
 .../client-libs/gateway-client/src/client.rs  |  5 +----
 2 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/common/client-core/src/client/received_buffer.rs b/common/client-core/src/client/received_buffer.rs
index fe6887b01f..021fd75d1a 100644
--- a/common/client-core/src/client/received_buffer.rs
+++ b/common/client-core/src/client/received_buffer.rs
@@ -50,19 +50,26 @@ struct ReceivedMessagesBufferInner<R: MessageReceiver> {
 }
 
 impl<R: MessageReceiver> ReceivedMessagesBufferInner<R> {
-    fn recover_from_fragment(&mut self, fragment_data: &[u8]) -> Option<NymMessage> {
+    fn recover_from_fragment(
+        &mut self,
+        fragment_data: &[u8],
+        fragment_data_size: usize,
+    ) -> Option<NymMessage> {
         if nym_sphinx::cover::is_cover(fragment_data) {
             trace!("The message was a loop cover message! Skipping it");
+            // NOTE: it's important to note that there is quite a bit of difference in size of
+            // received and sent packets due to the sphinx layers being removed by the exit gateway
+            // before it reaches the mixnet client.
             self.stats_tx
                 .report(PacketStatisticsEvent::CoverPacketReceived(
-                    fragment_data.len(),
+                    fragment_data_size,
                 ));
             return None;
         }
 
         self.stats_tx
             .report(PacketStatisticsEvent::RealPacketReceived(
-                fragment_data.len(),
+                fragment_data_size,
             ));
 
         let fragment = match self.message_receiver.recover_fragment(fragment_data) {
@@ -116,15 +123,17 @@ impl<R: MessageReceiver> ReceivedMessagesBufferInner<R> {
         reply_ciphertext: &mut [u8],
         reply_key: SurbEncryptionKey,
     ) -> Result<Option<NymMessage>, MessageRecoveryError> {
+        let reply_ciphertext_size = reply_ciphertext.len();
         // note: this performs decryption IN PLACE without extra allocation
         self.message_receiver
             .recover_plaintext_from_reply(reply_ciphertext, reply_key)?;
         let fragment_data = reply_ciphertext;
 
-        Ok(self.recover_from_fragment(fragment_data))
+        Ok(self.recover_from_fragment(fragment_data, reply_ciphertext_size))
     }
 
     fn process_received_regular_packet(&mut self, mut raw_fragment: Vec<u8>) -> Option<NymMessage> {
+        let raw_fragment_size = raw_fragment.len();
         let fragment_data = match self.message_receiver.recover_plaintext_from_regular_packet(
             self.local_encryption_keypair.private_key(),
             &mut raw_fragment,
@@ -136,7 +145,7 @@ impl<R: MessageReceiver> ReceivedMessagesBufferInner<R> {
             Ok(frag_data) => frag_data,
         };
 
-        self.recover_from_fragment(fragment_data)
+        self.recover_from_fragment(fragment_data, raw_fragment_size)
     }
 }
 
diff --git a/common/client-libs/gateway-client/src/client.rs b/common/client-libs/gateway-client/src/client.rs
index f8803bcc12..f65e61c29f 100644
--- a/common/client-libs/gateway-client/src/client.rs
+++ b/common/client-libs/gateway-client/src/client.rs
@@ -354,10 +354,7 @@ impl<C, St> GatewayClient<C, St> {
         msg: Message,
     ) -> Result<(), GatewayClientError> {
         match self.connection {
-            SocketState::Available(ref mut conn) => {
-                conn.send(msg).await?;
-                Ok(())
-            }
+            SocketState::Available(ref mut conn) => Ok(conn.send(msg).await?),
             SocketState::PartiallyDelegated(ref mut partially_delegated) => {
                 if let Err(err) = partially_delegated.send_without_response(msg).await {
                     error!("failed to send message without response - {err}...");