feat(rust): use a better encoding for sending messages

etorreborre · etorreborre · commit 84da455e104e · 2025-02-10T12:51:40.000+01:00
diff --git a/implementations/rust/ockam/ockam_api/src/nodes/service/messages.rs b/implementations/rust/ockam/ockam_api/src/nodes/service/messages.rs
@@ -1,7 +1,8 @@
 use crate::error::ApiError;
 use crate::nodes::{BackgroundNodeClient, NodeManager, NodeManagerWorker};
 use miette::IntoDiagnostic;
-use minicbor::{CborLen, Decode, Encode};
+use minicbor::encode::Write;
+use minicbor::{encode, CborLen, Decode, Decoder, Encode, Encoder};
 use ockam_core::api::{Error, Request, Response};
 use ockam_core::{self, async_trait, Decodable, Encodable, Encoded, Message, Result};
 use ockam_multiaddr::MultiAddr;
@@ -13,23 +14,23 @@ const TARGET: &str = "ockam_api::message";
 
 #[async_trait]
 pub trait Messages {
-    async fn send_message<R: Message>(
+    async fn send_message<T: Message, R: Message>(
         &self,
         ctx: &Context,
         to: &MultiAddr,
-        message: Vec<u8>,
+        message: T,
         timeout: Option<Duration>,
     ) -> miette::Result<R>;
 }
 
 #[async_trait]
 impl Messages for NodeManager {
     #[instrument(skip_all)]
-    async fn send_message<R: Message>(
+    async fn send_message<T: Message, R: Message>(
         &self,
         ctx: &Context,
         to: &MultiAddr,
-        message: Vec<u8>,
+        message: T,
         timeout: Option<Duration>,
     ) -> miette::Result<R> {
         let connection = self
@@ -56,11 +57,11 @@ impl Messages for NodeManager {
 #[async_trait]
 impl Messages for BackgroundNodeClient {
     #[instrument(skip_all)]
-    async fn send_message<R: Message>(
+    async fn send_message<T: Message, R: Message>(
         &self,
         ctx: &Context,
         to: &MultiAddr,
-        message: Vec<u8>,
+        message: T,
         timeout: Option<Duration>,
     ) -> miette::Result<R> {
         let request = Request::post("v0/message").body(SendMessage::new(to, message));
@@ -69,10 +70,10 @@ impl Messages for BackgroundNodeClient {
 }
 
 impl NodeManagerWorker {
-    pub(crate) async fn send_message<R: Message>(
+    pub(crate) async fn send_message<T: Message, R: Message>(
         &self,
         ctx: &Context,
-        send_message: SendMessage,
+        send_message: SendMessage<T>,
     ) -> Result<Response<R>, Response<Error>> {
         let multiaddr = send_message.multiaddr()?;
         let msg = send_message.message;
@@ -96,25 +97,51 @@ impl NodeManagerWorker {
 #[derive(Debug, Clone, Encode, Decode, CborLen, Message)]
 #[rustfmt::skip]
 #[cbor(map)]
-pub struct SendMessage {
+pub struct SendMessage<T: Message> {
     #[n(1)] pub route: String,
-    #[n(2)] pub message: Vec<u8>,
+    #[n(2)] pub message: T,
 }
 
-impl Encodable for SendMessage {
+impl<T: Message> SendMessage<T> {
+    fn encode_send_message<W>(self, buf: W) -> Result<(), encode::Error<W::Error>>
+    where
+        W: Write,
+    {
+        let mut e = Encoder::new(buf);
+        e.encode(&self.route)?;
+        e.writer_mut()
+            .write_all(&<T as Encodable>::encode(self.message).map_err(encode::Error::message)?)
+            .map_err(|_| encode::Error::message("encoding error"))?;
+        Ok(())
+    }
+
+    fn into_vec(self) -> Result<Vec<u8>, encode::Error<<Vec<u8> as Write>::Error>> {
+        let mut buf = Vec::new();
+        self.encode_send_message(&mut buf)?;
+        Ok(buf)
+    }
+}
+
+impl<T: Message> Encodable for SendMessage<T> {
     fn encode(self) -> Result<Encoded> {
-        Ok(minicbor::to_vec(self)?)
+        Ok(self.into_vec()?)
     }
 }
 
-impl Decodable for SendMessage {
+impl<T: Message> Decodable for SendMessage<T> {
     fn decode(e: &[u8]) -> Result<Self> {
-        Ok(minicbor::decode(e)?)
+        let mut dec = Decoder::new(e);
+        let route: String = dec.decode()?;
+        let message = dec.input().get(dec.position()..e.len()).unwrap();
+        Ok(SendMessage {
+            route,
+            message: <T as Decodable>::decode(message)?,
+        })
     }
 }
 
-impl SendMessage {
-    pub fn new(route: &MultiAddr, message: Vec<u8>) -> Self {
+impl<T: Message> SendMessage<T> {
+    pub fn new(route: &MultiAddr, message: T) -> Self {
         Self {
             route: route.to_string(),
             message,
diff --git a/implementations/rust/ockam/ockam_api/src/nodes/service/worker.rs b/implementations/rust/ockam/ockam_api/src/nodes/service/worker.rs
@@ -266,10 +266,11 @@ impl NodeManagerWorker {
 
             // ==*== Messages ==*==
             (Post, ["v0", "message"]) => {
-                let send_message: SendMessage = Decodable::decode(&body)?;
+                let send_message: SendMessage<Vec<u8>> = Decodable::decode(&body)?;
                 encode_response(
                     &header,
-                    self.send_message::<Vec<u8>>(ctx, send_message).await,
+                    self.send_message::<Vec<u8>, Vec<u8>>(ctx, send_message)
+                        .await,
                 )
             }
 
diff --git a/implementations/rust/ockam/ockam_api/src/uppercase.rs b/implementations/rust/ockam/ockam_api/src/uppercase.rs
@@ -9,10 +9,8 @@ impl Worker for Uppercase {
 
     #[instrument(skip_all, name = "Uppercase::handle_message")]
     async fn handle_message(&mut self, ctx: &mut Context, msg: Routed<String>) -> Result<()> {
-        ctx.send(
-            msg.return_route().clone(),
-            msg.into_body()?.to_uppercase().as_bytes().to_vec(),
-        )
-        .await
+        let return_route = msg.return_route().clone();
+        ctx.send(return_route.clone(), msg.into_body()?.to_uppercase())
+            .await
     }
 }
diff --git a/implementations/rust/ockam/ockam_command/src/message/send.rs b/implementations/rust/ockam/ockam_command/src/message/send.rs
@@ -6,9 +6,8 @@ use tracing::info;
 use ockam::Context;
 use ockam_api::address::extract_address_value;
 use ockam_api::nodes::service::messages::Messages;
-use ockam_api::nodes::BackgroundNodeClient;
 use ockam_api::nodes::InMemoryNode;
-use ockam_core::Encodable;
+use ockam_api::nodes::{BackgroundNodeClient, NodeManager};
 use ockam_multiaddr::MultiAddr;
 
 use crate::project::util::{
@@ -74,11 +73,9 @@ impl Command for SendCommand {
 
         // Setup environment depending on whether we are sending the message from a background node
         // or an in-memory node
-        let response: Vec<u8> = if let Some(node) = &self.from {
-            BackgroundNodeClient::create_to_node(ctx, &opts.state, node.as_str())?
-                .send_message(ctx, &to, self.get_message()?, Some(self.timeout.timeout))
-                .await
-                .map_err(Error::Retry)?
+        let result = if let Some(node) = &self.from {
+            let client = BackgroundNodeClient::create_to_node(ctx, &opts.state, node.as_str())?;
+            self.send_message(&client, ctx, &to).await?
         } else {
             let identity_name = opts
                 .state
@@ -113,18 +110,8 @@ impl Command for SendCommand {
             .map_err(Error::Retry)?;
             let to = clean_projects_multiaddr(to, projects_sc)?;
             info!("sending to {to}");
-            node_manager
-                .send_message(ctx, &to, self.get_message()?, Some(self.timeout.timeout))
-                .await
-                .map_err(Error::Retry)?
-        };
-
-        let result = if self.hex {
-            hex::encode(response)
-        } else {
-            String::from_utf8(response)
-                .into_diagnostic()
-                .context("Received content is not a valid utf8 string")?
+            let n: &NodeManager = &node_manager;
+            self.send_message(n, ctx, &to).await?
         };
 
         opts.terminal.stdout().plain(result).write_line()?;
@@ -133,13 +120,28 @@ impl Command for SendCommand {
 }
 
 impl SendCommand {
-    fn get_message(&self) -> crate::Result<Vec<u8>> {
+    async fn send_message(
+        self,
+        client: &impl Messages,
+        ctx: &Context,
+        to: &MultiAddr,
+    ) -> crate::Result<String> {
         if self.hex {
-            Ok(hex::decode(self.message.clone())
+            let to_send = hex::decode(self.message.clone())
                 .into_diagnostic()
-                .context("The message is not a valid hex string")?)
+                .context("The message is not a valid hex string")?;
+            let response: Vec<u8> = client
+                .send_message(ctx, to, to_send, Some(self.timeout.timeout))
+                .await
+                .map_err(Error::Retry)?;
+            Ok(hex::encode(response))
         } else {
-            Ok(Encodable::encode(self.message.clone())?)
+            client
+                .send_message(ctx, to, self.message, Some(self.timeout.timeout))
+                .await
+                .map_err(Error::Retry)
+                .into_diagnostic()
+                .context("Received content is not a valid utf8 string")
         }
     }
 }

Original file line number	Diff line number	Diff line change
`@@ -266,10 +266,11 @@ impl NodeManagerWorker {`
`266`	`266`
`267`	`267`	`// ==== Messages ====`
`268`	`268`	`(Post, ["v0", "message"]) => {`
`269`		`- let send_message: SendMessage = Decodable::decode(&body)?;`
	`269`	`+ let send_message: SendMessage<Vec<u8>> = Decodable::decode(&body)?;`
`270`	`270`	`encode_response(`
`271`	`271`	`&header,`
`272`		`- self.send_message::<Vec<u8>>(ctx, send_message).await,`
	`272`	`+ self.send_message::<Vec<u8>, Vec<u8>>(ctx, send_message)`
	`273`	`+ .await,`
`273`	`274`	`)`
`274`	`275`	`}`
`275`	`276`
Original file line number	Diff line number	Diff line change
`@@ -9,10 +9,8 @@ impl Worker for Uppercase {`
`9`	`9`
`10`	`10`	`#[instrument(skip_all, name = "Uppercase::handle_message")]`
`11`	`11`	`async fn handle_message(&mut self, ctx: &mut Context, msg: Routed<String>) -> Result<()> {`
`12`		`- ctx.send(`
`13`		`- msg.return_route().clone(),`
`14`		`- msg.into_body()?.to_uppercase().as_bytes().to_vec(),`
`15`		`- )`
`16`		`- .await`
	`12`	`+ let return_route = msg.return_route().clone();`
	`13`	`+ ctx.send(return_route.clone(), msg.into_body()?.to_uppercase())`
	`14`	`+ .await`
`17`	`15`	`}`
`18`	`16`	`}`