Fix decrypting stored received msg (#1786)

* Fix decrypting stored received msg

* rustfmt

* Moving binary message recovery to separate function

Co-authored-by: Jędrzej Stuczyński <jedrzej.stuczynski@gmail.com>

common/client-libs/gateway-client/src/client.rs

@@ -2,19 +2,13 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use crate::bandwidth::BandwidthController;
-use crate::cleanup_socket_message;
 use crate::error::GatewayClientError;
 use crate::packet_router::PacketRouter;
 pub use crate::packet_router::{
     AcknowledgementReceiver, AcknowledgementSender, MixnetMessageReceiver, MixnetMessageSender,
 };
 use crate::socket_state::{PartiallyDelegated, SocketState};
-#[cfg(target_arch = "wasm32")]
-use crate::wasm_storage::PersistentStorage;
-#[cfg(feature = "coconut")]
-use coconut_interface::Credential;
-#[cfg(not(target_arch = "wasm32"))]
-use credential_storage::PersistentStorage;
+use crate::{cleanup_socket_message, try_decrypt_binary_message};
 use crypto::asymmetric::identity;
 use futures::{FutureExt, SinkExt, StreamExt};
 use gateway_requests::authentication::encrypted_address::EncryptedAddressBytes;
@@ -28,13 +22,20 @@ use rand::rngs::OsRng;
 use std::convert::TryFrom;
 use std::sync::Arc;
 use std::time::Duration;
-#[cfg(not(target_arch = "wasm32"))]
-use task::ShutdownListener;
 use tungstenite::protocol::Message;
 
+#[cfg(feature = "coconut")]
+use coconut_interface::Credential;
+
+#[cfg(not(target_arch = "wasm32"))]
+use credential_storage::PersistentStorage;
+#[cfg(not(target_arch = "wasm32"))]
+use task::ShutdownListener;
 #[cfg(not(target_arch = "wasm32"))]
 use tokio_tungstenite::connect_async;
 
+#[cfg(target_arch = "wasm32")]
+use crate::wasm_storage::PersistentStorage;
 #[cfg(target_arch = "wasm32")]
 use wasm_timer;
 #[cfg(target_arch = "wasm32")]
@@ -336,7 +337,15 @@ impl GatewayClient {
                     };
                     match ws_msg {
                         Message::Binary(bin_msg) => {
-                            if let Err(err) = self.packet_router.route_received(vec![bin_msg]) {
+                            // if we have established the shared key already, attempt to use it for decryption
+                            // otherwise there's not much we can do apart from just routing what we have on hand
+                            if let Some(shared_keys) = &self.shared_key {
+                                if let Some(plaintext) = try_decrypt_binary_message(bin_msg, shared_keys) {
+                                    if let Err(err) = self.packet_router.route_received(vec![plaintext]) {
+                                        log::warn!("Route received failed: {:?}", err);
+                                    }
+                                }
+                            } else if let Err(err) = self.packet_router.route_received(vec![bin_msg]) {
                                 log::warn!("Route received failed: {:?}", err);
                             }
                         }

common/client-libs/gateway-client/src/lib.rs

@@ -3,6 +3,9 @@
 
 use crate::error::GatewayClientError;
 pub use client::GatewayClient;
+use gateway_requests::registration::handshake::SharedKeys;
+use gateway_requests::BinaryResponse;
+use log::warn;
 pub use packet_router::{
     AcknowledgementReceiver, AcknowledgementSender, MixnetMessageReceiver, MixnetMessageSender,
 };
@@ -37,3 +40,21 @@ pub(crate) fn cleanup_socket_messages(
         None => Err(GatewayClientError::ConnectionAbruptlyClosed),
     }
 }
+
+pub(crate) fn try_decrypt_binary_message(
+    bin_msg: Vec<u8>,
+    shared_keys: &SharedKeys,
+) -> Option<Vec<u8>> {
+    match BinaryResponse::try_from_encrypted_tagged_bytes(bin_msg, shared_keys) {
+        Ok(bin_response) => match bin_response {
+            BinaryResponse::PushedMixMessage(plaintext) => Some(plaintext),
+        },
+        Err(err) => {
+            warn!(
+                "message received from the gateway was malformed! - {:?}",
+                err
+            );
+            None
+        }
+    }
+}

common/client-libs/gateway-client/src/socket_state.rs

@@ -1,14 +1,13 @@
 // Copyright 2021-2022 - Nym Technologies SA <contact@nymtech.net>
 // SPDX-License-Identifier: Apache-2.0
 
-use crate::cleanup_socket_messages;
 use crate::error::GatewayClientError;
 use crate::packet_router::PacketRouter;
+use crate::{cleanup_socket_messages, try_decrypt_binary_message};
 use futures::channel::oneshot;
 use futures::stream::{SplitSink, SplitStream};
 use futures::{SinkExt, StreamExt};
 use gateway_requests::registration::handshake::SharedKeys;
-use gateway_requests::BinaryResponse;
 use log::*;
 use std::sync::Arc;
 #[cfg(not(target_arch = "wasm32"))]
@@ -50,21 +49,9 @@ impl PartiallyDelegated {
             match ws_msg {
                 Message::Binary(bin_msg) => {
                     // this function decrypts the request and checks the MAC
-                    let plaintext = match BinaryResponse::try_from_encrypted_tagged_bytes(
-                        bin_msg, shared_key,
-                    ) {
-                        Ok(bin_response) => match bin_response {
-                            BinaryResponse::PushedMixMessage(plaintext) => plaintext,
-                        },
-                        Err(err) => {
-                            warn!(
-                                "message received from the gateway was malformed! - {:?}",
-                                err
-                            );
-                            continue;
-                        }
-                    };
-                    plaintexts.push(plaintext)
+                    if let Some(plaintext) = try_decrypt_binary_message(bin_msg, shared_key) {
+                        plaintexts.push(plaintext)
+                    }
                 }
                 // I think that in the future we should perhaps have some sequence number system, i.e.
                 // so each request/response pair can be easily identified, so that if messages are