better message structures

Author: erhant

src/handlers/workflow.rs

@@ -5,7 +5,7 @@ use ollama_workflows::{Entry, Executor, ModelProvider, ProgramMemory, Workflow};
 use serde::Deserialize;
 
 use crate::node::DriaComputeNode;
-use crate::utils::payload::{TaskRequest, TaskRequestPayload};
+use crate::utils::payload::{TaskRequestPayload, TaskResponsePayload};
 use crate::utils::{get_current_time_nanos, DKNMessage};
 
 use super::ComputeHandler;
@@ -62,12 +62,6 @@ impl ComputeHandler for WorkflowHandler {
         // obtain public key from the payload
         let task_public_key = hex::decode(&task.public_key)?;
 
-        let task = TaskRequest {
-            task_id: task.task_id,
-            input: task.input,
-            public_key: task_public_key,
-        };
-
         // read model / provider from the task
         let (model_provider, model) = node
             .config
@@ -112,16 +106,16 @@ impl ComputeHandler for WorkflowHandler {
         }
 
         // prepare signed and encrypted payload
-        let payload = DKNMessage::new_signed_encrypted_payload(
+        let payload = TaskResponsePayload::new(
             result,
             &task.task_id,
-            &task.public_key,
+            &task_public_key,
             &node.config.secret_key,
         )?;
         let payload_str = payload.to_string()?;
-        let message = DKNMessage::new(payload_str, result_topic);
 
         // publish the result
+        let message = DKNMessage::new(payload_str, result_topic);
         node.publish(message)?;
 
         Ok(MessageAcceptance::Accept)

src/utils/message.rs

@@ -1,7 +1,6 @@
 use crate::utils::{
     crypto::{sha256hash, sign_bytes_recoverable},
     get_current_time_nanos,
-    payload::TaskResponsePayload,
 };
 use base64::{prelude::BASE64_STANDARD, Engine};
 use core::fmt;
@@ -10,19 +9,22 @@ use eyre::Result;
 use libsecp256k1::SecretKey;
 use serde::{Deserialize, Serialize};
 
-/// A parsed message from gossipsub. When first received, the message data is simply a vector of bytes.
-/// We treat that bytearray as a stringified JSON object, and parse it into this struct.
-///
-/// TODO: these are all available at protocol level as well
-/// - payload is the data itself
-/// - topic is available as TopicHash of Gossipsub
-/// - version is given within the Identify protocol
-/// - timestamp is available at protocol level via DataTransform
+/// A message within Dria Knowledge Network.
 #[derive(Serialize, Deserialize, Debug, Clone)]
 pub struct DKNMessage {
+    /// Base64 encoded data
     pub(crate) payload: String,
+    /// The topic of the message, derived from `TopicHash`
+    ///
+    /// NOTE: This can be obtained via TopicHash in GossipSub
     pub(crate) topic: String,
+    /// The version of the Dria Compute Node
+    ///
+    /// NOTE: This can be obtained via Identify protocol version
     pub(crate) version: String,
+    /// The timestamp of the message, in nanoseconds
+    ///
+    /// NOTE: This can be obtained via DataTransform in GossipSub
     pub(crate) timestamp: u128,
 }
 
@@ -61,39 +63,6 @@ impl DKNMessage {
         Self::new(signed_payload, topic)
     }
 
-    /// Creates the payload of a computation result.
-    ///
-    /// - Sign `task_id || payload` with node `self.secret_key`
-    /// - Encrypt `result` with `task_public_key`
-    /// TODO: this is not supposed to be here
-    pub fn new_signed_encrypted_payload(
-        payload: impl AsRef<[u8]>,
-        task_id: &str,
-        encrypting_public_key: &[u8],
-        signing_secret_key: &SecretKey,
-    ) -> Result<TaskResponsePayload> {
-        // create the message `task_id || payload`
-        let mut preimage = Vec::new();
-        preimage.extend_from_slice(task_id.as_ref());
-        preimage.extend_from_slice(payload.as_ref());
-
-        // sign the message
-        // TODO: use `sign_recoverable` here instead?
-        let digest = libsecp256k1::Message::parse(&sha256hash(preimage));
-        let (signature, recid) = libsecp256k1::sign(&digest, signing_secret_key);
-        let signature: [u8; 64] = signature.serialize();
-        let recid: [u8; 1] = [recid.serialize()];
-
-        // encrypt payload itself
-        let ciphertext = ecies::encrypt(encrypting_public_key, payload.as_ref())?;
-
-        Ok(TaskResponsePayload {
-            ciphertext: hex::encode(ciphertext),
-            signature: format!("{}{}", hex::encode(signature), hex::encode(recid)),
-            task_id: task_id.to_string(),
-        })
-    }
-
     /// Decodes the base64 payload into bytes.
     pub fn decode_payload(&self) -> Result<Vec<u8>, base64::DecodeError> {
         BASE64_STANDARD.decode(&self.payload)
@@ -161,6 +130,7 @@ impl TryFrom<libp2p::gossipsub::Message> for DKNMessage {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::utils::payload::TaskResponsePayload;
     use crate::{utils::crypto::sha256hash, DriaComputeNodeConfig};
     use ecies::decrypt;
     use libsecp256k1::SecretKey;
@@ -252,7 +222,7 @@ mod tests {
         let task_public_key = PublicKey::from_secret_key(&task_secret_key);
 
         // create payload
-        let payload = DKNMessage::new_signed_encrypted_payload(
+        let payload = TaskResponsePayload::new(
             RESULT,
             TASK_ID,
             &task_public_key.serialize(),

src/utils/payload.rs

@@ -1,10 +1,11 @@
+use super::crypto::sha256hash;
+use crate::utils::{filter::FilterPayload, get_current_time_nanos};
 use eyre::Result;
 use fastbloom_rs::BloomFilter;
+use libsecp256k1::SecretKey;
 use serde::{Deserialize, Serialize};
 use uuid::Uuid;
 
-use crate::utils::{filter::FilterPayload, get_current_time_nanos};
-
 /// A computation task is the task of computing a result from a given input. The result is encrypted with the public key of the requester.
 /// Plain result is signed by the compute node's private key, and a commitment is computed from the signature and plain result.
 ///
@@ -15,16 +16,48 @@ use crate::utils::{filter::FilterPayload, get_current_time_nanos};
 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.
+    /// Signature of the payload with task id, hexadecimally encoded.
     pub signature: String,
-    /// Computation result encrypted with the public key of the task.
+    /// Result encrypted with the public key of the task, Hexadecimally encoded.
     pub ciphertext: String,
 }
 
 impl TaskResponsePayload {
     pub fn to_string(&self) -> Result<String> {
         serde_json::to_string(&serde_json::json!(self)).map_err(Into::into)
     }
+
+    /// Creates the payload of a computation result.
+    ///
+    /// - Sign `task_id || payload` with node `self.secret_key`
+    /// - Encrypt `result` with `task_public_key`
+    pub fn new(
+        payload: impl AsRef<[u8]>,
+        task_id: &str,
+        encrypting_public_key: &[u8],
+        signing_secret_key: &SecretKey,
+    ) -> Result<Self> {
+        // create the message `task_id || payload`
+        let mut preimage = Vec::new();
+        preimage.extend_from_slice(task_id.as_ref());
+        preimage.extend_from_slice(payload.as_ref());
+
+        // sign the message
+        // TODO: use `sign_recoverable` here instead?
+        let digest = libsecp256k1::Message::parse(&sha256hash(preimage));
+        let (signature, recid) = libsecp256k1::sign(&digest, signing_secret_key);
+        let signature: [u8; 64] = signature.serialize();
+        let recid: [u8; 1] = [recid.serialize()];
+
+        // encrypt payload itself
+        let ciphertext = ecies::encrypt(encrypting_public_key, payload.as_ref())?;
+
+        Ok(TaskResponsePayload {
+            ciphertext: hex::encode(ciphertext),
+            signature: format!("{}{}", hex::encode(signature), hex::encode(recid)),
+            task_id: task_id.to_string(),
+        })
+    }
 }
 
 /// A generic task request, given by Dria.
@@ -55,11 +88,3 @@ impl<T> TaskRequestPayload<T> {
         }
     }
 }
-
-/// A parsed `TaskRequestPayload`.
-#[derive(Debug, Clone)]
-pub struct TaskRequest<T> {
-    pub task_id: String,
-    pub(crate) input: T,
-    pub(crate) public_key: Vec<u8>,
-}