fix: buggy assigned port id and buggy websocket transport (#80)

Author: lauti7

rpc/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "dcl-rpc"
-version = "2.0.0"
+version = "2.0.1"
 edition = "2021"
 description = "Decentraland RPC Implementation"
 repository = "https://github.com/decentraland/rpc-rust"

rpc/src/client.rs

@@ -161,6 +161,11 @@ impl<T: Transport> RpcClientPort<T> {
         }
     }
 
+    /// Get the Port's ID assigned by the [`RpcServer`](`crate::server::RpcServer`)
+    pub fn port_id(&self) -> u32 {
+        self.port_id
+    }
+
     /// Sends a request to the [`RpcServer`](crate::server::RpcServer)  requesting a remote module so then the returned module is loaded into memory so that it can be used by the client.
     ///
     /// It returns a `impl ServiceClient` that should auto generated by the codegen

rpc/src/messages_handlers.rs

@@ -39,13 +39,15 @@ use tokio_util::sync::CancellationToken;
 /// It spawns a background tasks to process every request
 ///
 #[derive(Default)]
+#[cfg(feature = "server")]
 pub struct ServerMessagesHandler {
     /// Data structure in charge of handling all messages related to streams
     pub streams_handler: Arc<StreamsHandler>,
     /// Stores listeners for client streams messages
     listeners: Mutex<HashMap<u32, AsyncChannelSender<StreamPackage>>>,
 }
 
+#[cfg(feature = "server")]
 impl ServerMessagesHandler {
     pub fn new() -> Self {
         Self {
@@ -380,6 +382,7 @@ type StreamPackage = (RpcMessageTypes, u32, StreamMessage);
 /// It's the data structure that actually owns the Transport attached to a `RpcClient`. The transport is drilled down up to get to `ClientMEssagesHandler`
 ///
 ///
+#[cfg(feature = "client")]
 pub struct ClientMessagesHandler<T: Transport + ?Sized> {
     /// Transport received by a `RpcClient`
     pub transport: Arc<T>,
@@ -412,6 +415,7 @@ pub struct ClientMessagesHandler<T: Transport + ?Sized> {
     process_cancellation_token: CancellationToken,
 }
 
+#[cfg(feature = "client")]
 impl<T: Transport + ?Sized + 'static> ClientMessagesHandler<T> {
     pub fn new(transport: Arc<T>) -> Self {
         Self {
@@ -504,7 +508,7 @@ impl<T: Transport + ?Sized + 'static> ClientMessagesHandler<T> {
                     }
                 }
                 Err(error) => {
-                    error!("> ClientMessagesHandler > process > Error on receive, breaking the listening: {error:?}");
+                    error!("> ClientMessagesHandler > process > Error on receive {error:?}");
                     if matches!(error, TransportError::Closed) {
                         info!("> ClientMessagesHandler > process > closing...");
                         break;

rpc/src/server.rs

@@ -41,13 +41,13 @@ pub enum ServerError {
     /// Error on decoding bytes (`Vec<u8>`) into a given type using [`crate::rpc_protocol::parse::parse_protocol_message`] or using the [`Message::decode`]
     ProtocolError,
     /// Port was not found in the server state, possibly not created
-    PortNotFound,
+    PortNotFound(u32),
     /// Error on loading a Module, unlikely to happen
     LoadModuleError,
     /// Module was not found, not registered in the server
-    ModuleNotFound,
+    ModuleNotFound(String),
     /// Given procedure's ID was not found
-    ProcedureNotFound,
+    ProcedureNotFound(u32),
     /// Unexpexted Error while responding back or Error on sending the original procedure response
     ///
     /// This error should be use as a "re-try" when a [`Transport::send`] failed.
@@ -58,9 +58,9 @@ impl RemoteErrorResponse for ServerError {
     fn error_code(&self) -> u32 {
         match self {
             Self::ProtocolError => 1,
-            Self::PortNotFound => 2,
-            Self::ModuleNotFound => 3,
-            Self::ProcedureNotFound => 4,
+            Self::PortNotFound(_) => 2,
+            Self::ModuleNotFound(_) => 3,
+            Self::ProcedureNotFound(_) => 4,
             Self::UnexpectedErrorOnTransport => 5,
             Self::LoadModuleError => 0, // it's unlikely to happen
         }
@@ -69,10 +69,10 @@ impl RemoteErrorResponse for ServerError {
     fn error_message(&self) -> String {
         match self {
             Self::ProtocolError => "Error on parsing a message. The content seems to be corrupted and not to meet the protocol requirements".to_string(),
-            Self::PortNotFound => "The given Port's ID was not found".to_string(),
+            Self::PortNotFound(id) => format!("The given Port's ID: {id} was not found"),
             Self::LoadModuleError => "Error on loading a module".to_string(),
-            Self::ModuleNotFound => "Module wasn't found on the server, check the name".to_string(),
-            Self::ProcedureNotFound => "Procedure's ID wasn't found on the server".to_string(),
+            Self::ModuleNotFound(module_name) => format!("Module wasn't found on the server, check the name: {module_name}"),
+            Self::ProcedureNotFound(id) => format!("Procedure's ID: {id} wasn't found on the server"),
             Self::UnexpectedErrorOnTransport => "Error on the transport while sending the original procedure response".to_string()
         }
     }
@@ -194,6 +194,8 @@ pub struct RpcServer<Context, T: Transport + ?Sized> {
     server_events_receiver: Option<UnboundedReceiver<ServerEvents<T>>>,
     /// The id that will be assigned if a new transport is a attached
     next_transport_id: u32,
+    /// THe id that will be assigned to a port when it's created.
+    next_port_id: u32,
 }
 impl<Context: Send + Sync + 'static, T: Transport + ?Sized + 'static> RpcServer<Context, T> {
     pub fn create(ctx: Context) -> Self {
@@ -207,6 +209,7 @@ impl<Context: Send + Sync + 'static, T: Transport + ?Sized + 'static> RpcServer<
             context: Arc::new(ctx),
             messages_handler: Arc::new(ServerMessagesHandler::new()),
             next_transport_id: 1,
+            next_port_id: 1,
             server_events_sender: ServerEventsSender(channel.0),
             server_events_receiver: Some(channel.1),
         }
@@ -387,10 +390,10 @@ impl<Context: Send + Sync + 'static, T: Transport + ?Sized + 'static> RpcServer<
                                         }
                                     }
                                     Err(error) => {
-                                        error!(
-                                            "> From a Transport > Error on receiving: {error:?}"
-                                        );
                                         if matches!(error, TransportError::Closed) {
+                                            error!(
+                                                "> From a Transport > Transport is already closed. Breaking..."
+                                            );
                                             if tx_cloned
                                                 .send(TransportNotification::CloseTransport(id))
                                                 .is_err()
@@ -400,6 +403,7 @@ impl<Context: Send + Sync + 'static, T: Transport + ?Sized + 'static> RpcServer<
                                             }
                                             break;
                                         }
+                                        error!("> From a Transport > Error on receiving {error:?}");
                                     }
                                 }
                             }
@@ -558,7 +562,9 @@ impl<Context: Send + Sync + 'static, T: Transport + ?Sized + 'static> RpcServer<
 
                 Ok(())
             }
-            _ => Err(ServerResultError::External(ServerError::PortNotFound)),
+            _ => Err(ServerResultError::External(ServerError::PortNotFound(
+                request.port_id,
+            ))),
         }
     }
 
@@ -605,7 +611,9 @@ impl<Context: Send + Sync + 'static, T: Transport + ?Sized + 'static> RpcServer<
                 return Err(ServerResultError::External(ServerError::LoadModuleError));
             }
         } else {
-            return Err(ServerResultError::External(ServerError::PortNotFound));
+            return Err(ServerResultError::External(ServerError::PortNotFound(
+                request_module.port_id,
+            )));
         }
 
         Ok(())
@@ -627,7 +635,8 @@ impl<Context: Send + Sync + 'static, T: Transport + ?Sized + 'static> RpcServer<
         message_number: u32,
         payload: Vec<u8>,
     ) -> ServerResult<()> {
-        let port_id = (self.ports.len() + 1) as u32;
+        let port_id = self.next_port_id;
+        self.next_port_id += 1;
         let create_port = CreatePort::decode(payload.as_slice())
             .map_err(|_| ServerResultError::External(ServerError::ProtocolError))?;
         let port_name = create_port.port_name;
@@ -784,7 +793,9 @@ impl<Context> RpcServerPort<Context> {
                 .expect("Already checked."))
         } else {
             match self.registered_modules.get(&module_name) {
-                None => Err(ServerResultError::External(ServerError::ModuleNotFound)),
+                None => Err(ServerResultError::External(ServerError::ModuleNotFound(
+                    module_name,
+                ))),
                 Some(module_generator) => {
                     let mut server_module_declaration = ServerModuleDeclaration {
                         procedures: Vec::new(),
@@ -823,7 +834,9 @@ impl<Context> RpcServerPort<Context> {
     fn get_procedure(&self, procedure_id: u32) -> ServerResult<ProcedureDefinition<Context>> {
         match self.procedures.get(&procedure_id) {
             Some(procedure_definition) => Ok(procedure_definition.clone()),
-            _ => Err(ServerResultError::External(ServerError::ProcedureNotFound)),
+            _ => Err(ServerResultError::External(ServerError::ProcedureNotFound(
+                procedure_id,
+            ))),
         }
     }
 }

rpc/src/transports/web_socket.rs

@@ -5,9 +5,9 @@ use super::{Transport, TransportError, TransportMessage};
 use async_trait::async_trait;
 use futures_util::{
     stream::{SplitSink, SplitStream},
-    SinkExt, StreamExt, TryStreamExt,
+    SinkExt, StreamExt,
 };
-use log::{debug, error, info};
+use log::{debug, error};
 use std::error::Error;
 use tokio::{
     net::{TcpListener, TcpStream},
@@ -188,22 +188,21 @@ impl WebSocketTransport {
 #[async_trait]
 impl Transport for WebSocketTransport {
     async fn receive(&self) -> Result<TransportMessage, TransportError> {
-        match self.read.lock().await.try_next().await {
-            Ok(Some(message)) => {
+        match self.read.lock().await.next().await {
+            Some(Ok(message)) => {
                 if message.is_binary() {
                     let message_data = message.into_data();
                     return Ok(message_data);
                 } else {
+                    if message.is_close() {
+                        return Err(TransportError::Closed);
+                    }
                     // Ignore messages that are not binary
                     error!("> WebSocketTransport > Received message is not binary");
                     return Err(TransportError::NotBinaryMessage);
                 }
             }
-            Ok(_) => {
-                error!("> WebSocketTransport > WEIRD: Ok(None)");
-                Err(TransportError::Closed)
-            }
-            Err(err) => {
+            Some(Err(err)) => {
                 error!(
                     "> WebSocketTransport > Failed to receive message {}",
                     err.to_string()
@@ -215,6 +214,10 @@ impl Transport for WebSocketTransport {
                     error => return Err(TransportError::Internal(Box::new(error))),
                 }
             }
+            None => {
+                error!("> WebSocketTransport > None received > Closing...");
+                return Err(TransportError::Closed);
+            }
         }
     }
 
@@ -243,7 +246,7 @@ impl Transport for WebSocketTransport {
     async fn close(&self) {
         match self.write.lock().await.close().await {
             Ok(_) => {
-                info!("> WebSocketTransport > Closed successfully")
+                debug!("> WebSocketTransport > Closed successfully")
             }
             Err(err) => {
                 error!("> WebSocketTransport > Error: Couldn't close tranport: {err:?}")