now entire message is published, not just payload

Author: erhant

src/handlers/mod.rs

@@ -8,13 +8,13 @@ pub use pingpong::PingpongHandler;
 mod workflow;
 pub use workflow::WorkflowHandler;
 
-use crate::{utils::P2PMessage, DriaComputeNode};
+use crate::{utils::DKNMessage, DriaComputeNode};
 
 #[async_trait]
 pub trait ComputeHandler {
     async fn handle_compute(
         node: &mut DriaComputeNode,
-        message: P2PMessage,
+        message: DKNMessage,
         result_topic: &str,
     ) -> Result<MessageAcceptance>;
 }

src/handlers/pingpong.rs

@@ -1,6 +1,6 @@
 use crate::{
     node::DriaComputeNode,
-    utils::{get_current_time_nanos, P2PMessage},
+    utils::{get_current_time_nanos, DKNMessage},
 };
 use async_trait::async_trait;
 use eyre::Result;
@@ -29,7 +29,7 @@ struct PingpongResponse {
 impl ComputeHandler for PingpongHandler {
     async fn handle_compute(
         node: &mut DriaComputeNode,
-        message: P2PMessage,
+        message: DKNMessage,
         result_topic: &str,
     ) -> Result<MessageAcceptance> {
         let pingpong = message.parse_payload::<PingpongPayload>(true)?;
@@ -54,12 +54,13 @@ impl ComputeHandler for PingpongHandler {
             models: node.config.model_config.models.clone(),
             timestamp: get_current_time_nanos(),
         };
-        let response = P2PMessage::new_signed(
+
+        let message = DKNMessage::new_signed(
             serde_json::json!(response_body).to_string(),
             result_topic,
             &node.config.secret_key,
         );
-        node.publish(response)?;
+        node.publish(message)?;
 
         // accept message, someone else may be included in the filter
         Ok(MessageAcceptance::Accept)
@@ -72,7 +73,7 @@ mod tests {
         utils::{
             crypto::{sha256hash, to_address},
             filter::FilterPayload,
-            P2PMessage,
+            DKNMessage,
         },
         DriaComputeNodeConfig,
     };
@@ -87,7 +88,7 @@ mod tests {
             "0208ef5e65a9c656a6f92fb2c770d5d5e2ecffe02a6aade19207f75110be6ae658"
         ))
         .expect("Should parse public key");
-        let message = P2PMessage {
+        let message = DKNMessage {
             payload: "Y2RmODcyNDlhY2U3YzQ2MDIzYzNkMzBhOTc4ZWY3NjViMWVhZDlmNWJhMDUyY2MxMmY0NzIzMjQyYjc0YmYyODFjMDA1MTdmMGYzM2VkNTgzMzk1YWUzMTY1ODQ3NWQyNDRlODAxYzAxZDE5MjYwMDM1MTRkNzEwMThmYTJkNjEwMXsidXVpZCI6ICI4MWE2M2EzNC05NmM2LTRlNWEtOTliNS02YjI3NGQ5ZGUxNzUiLCAiZGVhZGxpbmUiOiAxNzE0MTI4NzkyfQ==".to_string(),
             topic: "heartbeat".to_string(),
             version: "0.0.0".to_string(),

src/handlers/workflow.rs

@@ -6,7 +6,7 @@ use serde::Deserialize;
 
 use crate::node::DriaComputeNode;
 use crate::utils::payload::{TaskRequest, TaskRequestPayload};
-use crate::utils::{get_current_time_nanos, P2PMessage};
+use crate::utils::{get_current_time_nanos, DKNMessage};
 
 use super::ComputeHandler;
 
@@ -29,7 +29,7 @@ struct WorkflowPayload {
 impl ComputeHandler for WorkflowHandler {
     async fn handle_compute(
         node: &mut DriaComputeNode,
-        message: P2PMessage,
+        message: DKNMessage,
         result_topic: &str,
     ) -> Result<MessageAcceptance> {
         let task = message.parse_payload::<TaskRequestPayload<WorkflowPayload>>(true)?;
@@ -111,16 +111,17 @@ impl ComputeHandler for WorkflowHandler {
             }
         }
 
-        // publish the result
-        let payload = P2PMessage::new_signed_encrypted_payload(
+        // prepare signed and encrypted payload
+        let payload = DKNMessage::new_signed_encrypted_payload(
             result,
             &task.task_id,
             &task.public_key,
             &node.config.secret_key,
         )?;
         let payload_str = payload.to_string()?;
-        let message = P2PMessage::new(payload_str, result_topic);
+        let message = DKNMessage::new(payload_str, result_topic);
 
+        // publish the result
         node.publish(message)?;
 
         Ok(MessageAcceptance::Accept)

src/node.rs

@@ -7,7 +7,7 @@ use crate::{
     config::DriaComputeNodeConfig,
     handlers::{ComputeHandler, PingpongHandler, WorkflowHandler},
     p2p::P2PClient,
-    utils::{crypto::secret_to_keypair, AvailableNodes, P2PMessage},
+    utils::{crypto::secret_to_keypair, AvailableNodes, DKNMessage},
 };
 
 /// Number of seconds between refreshing the Admin RPC PeerIDs from Dria server.
@@ -56,8 +56,8 @@ impl DriaComputeNode {
         let p2p = P2PClient::new(keypair, listen_addr, &available_nodes)?;
 
         Ok(DriaComputeNode {
-            config,
             p2p,
+            config,
             cancellation,
             available_nodes,
             available_nodes_last_refreshed: tokio::time::Instant::now(),
@@ -86,16 +86,18 @@ impl DriaComputeNode {
         Ok(())
     }
 
-    /// Publishes a given message to the network.
-    /// The topic is expected to be provided within the message struct.
-    pub fn publish(&mut self, message: P2PMessage) -> Result<()> {
-        let message_bytes = message.payload.as_bytes().to_vec();
+    /// Publishes a given message to the network w.r.t the topic of it.
+    ///
+    /// Internally, the message is JSON serialized to bytes and then published to the network as is.
+    pub fn publish(&mut self, message: DKNMessage) -> Result<()> {
+        let message_bytes = serde_json::to_vec(&message)?;
         let message_id = self.p2p.publish(&message.topic, message_bytes)?;
         log::info!("Published message ({}) to {}", message_id, message.topic);
         Ok(())
     }
 
     /// Returns the list of connected peers.
+    #[inline(always)]
     pub fn peers(&self) -> Vec<(&libp2p_identity::PeerId, Vec<&gossipsub::TopicHash>)> {
         self.p2p.peers()
     }
@@ -223,13 +225,14 @@ impl DriaComputeNode {
     /// This prepared message includes the topic, payload, version and timestamp.
     ///
     /// This also checks the signature of the message, expecting a valid signature from admin node.
+    // TODO: move this somewhere?
     pub fn parse_message_to_prepared_message(
         &self,
         message: gossipsub::Message,
-    ) -> Result<P2PMessage> {
+    ) -> Result<DKNMessage> {
         // the received message is expected to use IdentHash for the topic, so we can see the name of the topic immediately.
         log::debug!("Parsing {} message.", message.topic.as_str());
-        let message = P2PMessage::try_from(message)?;
+        let message = DKNMessage::try_from(message)?;
         log::debug!("Parsed: {}", message);
 
         // check dria signature
@@ -270,7 +273,7 @@ mod tests {
 
         // publish a dummy message
         let topic = "foo";
-        let message = P2PMessage::new("hello from the other side", topic);
+        let message = DKNMessage::new("hello from the other side", topic);
         node.subscribe(topic).expect("should subscribe");
         node.publish(message).expect("should publish");
         node.unsubscribe(topic).expect("should unsubscribe");

src/utils/crypto.rs

@@ -60,8 +60,8 @@ mod tests {
     use hex::decode;
     use libsecp256k1::{recover, sign, verify, Message, PublicKey, SecretKey};
 
-    const DUMMY_KEY: &[u8; 32] = b"driadriadriadriadriadriadriadria";
-    const MESSAGE: &[u8] = "hello world".as_bytes();
+    const DUMMY_SECRET_KEY: &[u8; 32] = b"driadriadriadriadriadriadriadria";
+    const MESSAGE: &[u8] = b"hello world";
 
     #[test]
     fn test_hash() {
@@ -73,7 +73,7 @@ mod tests {
 
     #[test]
     fn test_address() {
-        let sk = SecretKey::parse_slice(DUMMY_KEY).expect("Should parse key.");
+        let sk = SecretKey::parse_slice(DUMMY_SECRET_KEY).expect("Should parse key.");
         let pk = PublicKey::from_secret_key(&sk);
         let addr = to_address(&pk);
         assert_eq!(
@@ -84,7 +84,7 @@ mod tests {
 
     #[test]
     fn test_encrypt_decrypt() {
-        let sk = SecretKey::parse_slice(DUMMY_KEY).expect("Should parse private key slice.");
+        let sk = SecretKey::parse_slice(DUMMY_SECRET_KEY).expect("Should parse private key slice.");
         let pk = PublicKey::from_secret_key(&sk);
         let (sk, pk) = (&sk.serialize(), &pk.serialize());
 
@@ -96,7 +96,7 @@ mod tests {
     #[test]
     fn test_sign_verify() {
         let secret_key =
-            SecretKey::parse_slice(DUMMY_KEY).expect("Should parse private key slice.");
+            SecretKey::parse_slice(DUMMY_SECRET_KEY).expect("Should parse private key slice.");
 
         // sign the message using the secret key
         let digest = sha256hash(MESSAGE);
@@ -121,7 +121,7 @@ mod tests {
     #[ignore = "run only with profiler if wanted"]
     fn test_memory_usage() {
         let secret_key =
-            SecretKey::parse_slice(DUMMY_KEY).expect("Should parse private key slice.");
+            SecretKey::parse_slice(DUMMY_SECRET_KEY).expect("Should parse private key slice.");
         let public_key = PublicKey::from_secret_key(&secret_key);
 
         // sign the message using the secret key

src/utils/filter.rs

@@ -3,7 +3,7 @@ use fastbloom_rs::{BloomFilter, Hashes, Membership};
 use serde::{Deserialize, Serialize};
 use serde_json::{json, to_string};
 
-/// A task filter is used to determine if a node is selected.
+/// A task Blfilter is used to determine if a node is selected.
 ///
 /// The filter is a Bloom Filter with a set of items and a false positive rate, it is serialized as a hex string.
 #[derive(Serialize, Deserialize, Debug, Clone)]

src/utils/message.rs

@@ -19,7 +19,7 @@ use serde::{Deserialize, Serialize};
 /// - version is given within the Identify protocol
 /// - timestamp is available at protocol level via DataTransform
 #[derive(Serialize, Deserialize, Debug, Clone)]
-pub struct P2PMessage {
+pub struct DKNMessage {
     pub(crate) payload: String,
     pub(crate) topic: String,
     pub(crate) version: String,
@@ -33,7 +33,7 @@ pub struct P2PMessage {
 /// and therefore use 128 characters: SIGNATURE_SIZE - 2.
 const SIGNATURE_SIZE_HEX: usize = 130;
 
-impl P2PMessage {
+impl DKNMessage {
     /// Creates a new message with current timestamp and version equal to the crate version.
     ///
     /// - `payload` is gives as bytes. It is to be `base64` encoded internally.
@@ -91,7 +91,6 @@ impl P2PMessage {
             ciphertext: hex::encode(ciphertext),
             signature: format!("{}{}", hex::encode(signature), hex::encode(recid)),
             task_id: task_id.to_string(),
-            timestamp: get_current_time_nanos(),
         })
     }
 
@@ -136,7 +135,7 @@ impl P2PMessage {
     }
 }
 
-impl fmt::Display for P2PMessage {
+impl fmt::Display for DKNMessage {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         let payload_decoded = self
             .decode_payload()
@@ -151,7 +150,7 @@ impl fmt::Display for P2PMessage {
     }
 }
 
-impl TryFrom<libp2p::gossipsub::Message> for P2PMessage {
+impl TryFrom<libp2p::gossipsub::Message> for DKNMessage {
     type Error = serde_json::Error;
 
     fn try_from(value: libp2p::gossipsub::Message) -> Result<Self, Self::Error> {
@@ -186,7 +185,7 @@ mod tests {
 
     #[test]
     fn test_display_message() {
-        let message = P2PMessage::new(b"hello world", "test-topic");
+        let message = DKNMessage::new(b"hello world", "test-topic");
         println!("{}", message);
     }
 
@@ -195,7 +194,7 @@ mod tests {
         // create payload & message
         let body = TestStruct::default();
         let payload = serde_json::to_vec(&json!(body)).expect("Should serialize");
-        let message = P2PMessage::new(payload, TOPIC);
+        let message = DKNMessage::new(payload, TOPIC);
 
         // decode message
         let message_body = message.decode_payload().expect("Should decode");
@@ -221,7 +220,7 @@ mod tests {
         // create payload & message with signature & body
         let body = TestStruct::default();
         let body_str = serde_json::to_string(&body).unwrap();
-        let message = P2PMessage::new_signed(body_str, TOPIC, &sk);
+        let message = DKNMessage::new_signed(body_str, TOPIC, &sk);
 
         // decode message
         let message_body = message.decode_payload().expect("Should decode");
@@ -253,7 +252,7 @@ mod tests {
         let task_public_key = PublicKey::from_secret_key(&task_secret_key);
 
         // create payload
-        let payload = P2PMessage::new_signed_encrypted_payload(
+        let payload = DKNMessage::new_signed_encrypted_payload(
             RESULT,
             TASK_ID,
             &task_public_key.serialize(),

src/utils/mod.rs

@@ -3,7 +3,7 @@ pub mod filter;
 pub mod payload;
 
 mod message;
-pub use message::P2PMessage;
+pub use message::DKNMessage;
 
 mod available_nodes;
 pub use available_nodes::AvailableNodes;

src/utils/payload.rs

@@ -13,14 +13,12 @@ use crate::utils::{filter::FilterPayload, get_current_time_nanos};
 #[derive(Debug, Clone, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub struct TaskResponsePayload {
+    /// The unique identifier of the task.
+    pub task_id: String,
     /// A signature on the digest of plaintext result, prepended with task id.
     pub signature: String,
     /// Computation result encrypted with the public key of the task.
     pub ciphertext: String,
-    /// The unique identifier of the task.
-    pub task_id: String,
-    /// Timestamp of the response.
-    pub timestamp: u128,
 }
 
 impl TaskResponsePayload {