even more tidying up

erhant · erhant · commit 1f0e94aec574 · 2024-10-02T19:30:51.000+03:00
diff --git a/src/handlers/mod.rs b/src/handlers/mod.rs
@@ -23,6 +23,8 @@ pub trait ComputeHandler {
     /// A generic handler for DKN tasks.
     ///
     /// Returns a `MessageAcceptance` value that tells the P2P client to accept the incoming message.
+    ///
+    /// The handler has mutable reference to the compute node, and therefore can respond within the handler itself in any way it would like.
     async fn handle_compute(
         node: &mut DriaComputeNode,
         message: DKNMessage,
diff --git a/src/handlers/pingpong.rs b/src/handlers/pingpong.rs
@@ -3,7 +3,7 @@ use crate::{
     DriaComputeNode,
 };
 use async_trait::async_trait;
-use eyre::Result;
+use eyre::{Context, Result};
 use libp2p::gossipsub::MessageAcceptance;
 use ollama_workflows::{Model, ModelProvider};
 use serde::{Deserialize, Serialize};
@@ -34,7 +34,9 @@ impl ComputeHandler for PingpongHandler {
         node: &mut DriaComputeNode,
         message: DKNMessage,
     ) -> Result<MessageAcceptance> {
-        let pingpong = message.parse_payload::<PingpongPayload>(true)?;
+        let pingpong = message
+            .parse_payload::<PingpongPayload>(true)
+            .wrap_err("Could not parse ping request")?;
 
         // check deadline
         let current_time = get_current_time_nanos();
@@ -68,77 +70,3 @@ impl ComputeHandler for PingpongHandler {
         Ok(MessageAcceptance::Accept)
     }
 }
-
-#[cfg(test)]
-mod tests {
-    use crate::{
-        utils::{
-            crypto::{sha256hash, to_address},
-            filter::TaskFilter,
-            DKNMessage,
-        },
-        DriaComputeNodeConfig,
-    };
-    use fastbloom_rs::{FilterBuilder, Membership};
-    use libsecp256k1::{recover, Message, PublicKey};
-
-    use super::PingpongPayload;
-
-    #[test]
-    fn test_pingpong_payload() {
-        let pk = PublicKey::parse_compressed(&hex_literal::hex!(
-            "0208ef5e65a9c656a6f92fb2c770d5d5e2ecffe02a6aade19207f75110be6ae658"
-        ))
-        .expect("Should parse public key");
-        let message = DKNMessage {
-            payload: "Y2RmODcyNDlhY2U3YzQ2MDIzYzNkMzBhOTc4ZWY3NjViMWVhZDlmNWJhMDUyY2MxMmY0NzIzMjQyYjc0YmYyODFjMDA1MTdmMGYzM2VkNTgzMzk1YWUzMTY1ODQ3NWQyNDRlODAxYzAxZDE5MjYwMDM1MTRkNzEwMThmYTJkNjEwMXsidXVpZCI6ICI4MWE2M2EzNC05NmM2LTRlNWEtOTliNS02YjI3NGQ5ZGUxNzUiLCAiZGVhZGxpbmUiOiAxNzE0MTI4NzkyfQ==".to_string(),
-            topic: "pingpong".to_string(),
-            version: "0.0.0".to_string(),
-            timestamp: 1714129073557846272,
-        };
-
-        assert!(message.is_signed(&pk).expect("Should check signature"));
-
-        let obj = message
-            .parse_payload::<PingpongPayload>(true)
-            .expect("Should parse payload");
-        assert_eq!(obj.uuid, "81a63a34-96c6-4e5a-99b5-6b274d9de175");
-        assert_eq!(obj.deadline, 1714128792);
-    }
-
-    /// This test demonstrates the process of pingpong & task assignment.
-    ///
-    /// A heart-beat message is sent over the network by Admin Node, and compute node responds with a signature.
-    #[test]
-    fn test_pingpong_and_task_assignment() {
-        let config = DriaComputeNodeConfig::default();
-
-        // a pingpong message is signed and sent to Admin Node over the p2p network
-        let pingpong_message = Message::parse(&sha256hash(b"sign-me"));
-        let (pingpong_signature, pingpong_recid) =
-            libsecp256k1::sign(&pingpong_message, &config.secret_key);
-
-        // admin recovers the address from the signature
-        let recovered_public_key = recover(&pingpong_message, &pingpong_signature, &pingpong_recid)
-            .expect("Could not recover");
-        assert_eq!(
-            config.public_key, recovered_public_key,
-            "Public key mismatch"
-        );
-        let address = to_address(&recovered_public_key);
-        assert_eq!(address, config.address, "Address mismatch");
-
-        // admin node assigns the task to the compute node via Bloom Filter
-        let mut bloom = FilterBuilder::new(100, 0.01).build_bloom_filter();
-        bloom.add(&address);
-        let filter_payload = TaskFilter::from(bloom);
-
-        // compute node receives the filter and checks if it is tasked
-        assert!(
-            filter_payload
-                .contains(&config.address)
-                .expect("Should check filter"),
-            "Node should be tasked"
-        );
-    }
-}
diff --git a/src/handlers/workflow.rs b/src/handlers/workflow.rs
@@ -1,6 +1,7 @@
 use async_trait::async_trait;
 use eyre::{eyre, Context, Result};
 use libp2p::gossipsub::MessageAcceptance;
+use libsecp256k1::PublicKey;
 use ollama_workflows::{Entry, Executor, ModelProvider, ProgramMemory, Workflow};
 use serde::Deserialize;
 
@@ -37,7 +38,7 @@ impl ComputeHandler for WorkflowHandler {
         let config = &node.config;
         let task = message
             .parse_payload::<TaskRequestPayload<WorkflowPayload>>(true)
-            .wrap_err("Could not parse error")?;
+            .wrap_err("Could not parse workflow task")?;
 
         // check if deadline is past or not
         let current_time = get_current_time_nanos();
@@ -97,8 +98,9 @@ impl ComputeHandler for WorkflowHandler {
         match exec_result {
             Ok(result) => {
                 // obtain public key from the payload
-                let task_public_key =
+                let task_public_key_bytes =
                     hex::decode(&task.public_key).wrap_err("Could not decode public key")?;
+                let task_public_key = PublicKey::parse_slice(&task_public_key_bytes, None)?;
 
                 // prepare signed and encrypted payload
                 let payload = TaskResponsePayload::new(
diff --git a/src/payloads/response.rs b/src/payloads/response.rs
@@ -1,6 +1,6 @@
-use crate::utils::crypto::sha256hash;
+use crate::utils::crypto::{encrypt_bytes, sha256hash, sign_bytes_recoverable};
 use eyre::Result;
-use libsecp256k1::SecretKey;
+use libsecp256k1::{PublicKey, SecretKey};
 use serde::{Deserialize, Serialize};
 
 /// 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.
@@ -27,28 +27,68 @@ impl TaskResponsePayload {
     pub fn new(
         result: impl AsRef<[u8]>,
         task_id: &str,
-        encrypting_public_key: &[u8],
+        encrypting_public_key: &PublicKey,
         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(result.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, result.as_ref())?;
+        let signature = sign_bytes_recoverable(&sha256hash(preimage), signing_secret_key);
+        let ciphertext = encrypt_bytes(result, encrypting_public_key)?;
 
         Ok(TaskResponsePayload {
             ciphertext: hex::encode(ciphertext),
-            signature: format!("{}{}", hex::encode(signature), hex::encode(recid)),
+            signature,
             task_id: task_id.to_string(),
         })
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use ecies::decrypt;
+    use libsecp256k1::{recover, verify, Message, PublicKey, RecoveryId, Signature};
+    use rand::thread_rng;
+
+    #[test]
+    fn test_task_response_payload() {
+        // this is the result that we are "sending"
+        const RESULT: &[u8; 44] = b"hey im an LLM and I came up with this output";
+
+        // the signer will sign the payload, and it will be verified
+        let signer_sk = SecretKey::random(&mut thread_rng());
+        let signer_pk = PublicKey::from_secret_key(&signer_sk);
+
+        // the payload will be encrypted with this key
+        let task_sk = SecretKey::random(&mut thread_rng());
+        let task_pk = PublicKey::from_secret_key(&task_sk);
+        let task_id = uuid::Uuid::new_v4().to_string();
+
+        // creates a signed and encrypted payload
+        let payload = TaskResponsePayload::new(RESULT, &task_id, &task_pk, &signer_sk)
+            .expect("Should create payload");
+
+        // decrypt result and compare it to plaintext
+        let ciphertext_bytes = hex::decode(payload.ciphertext).unwrap();
+        let result = decrypt(&task_sk.serialize(), &ciphertext_bytes).expect("Could not decrypt");
+        assert_eq!(result, RESULT, "Result mismatch");
+
+        // verify signature
+        let signature_bytes = hex::decode(payload.signature).expect("Should decode");
+        let signature = Signature::parse_standard_slice(&signature_bytes[..64]).unwrap();
+        let recid = RecoveryId::parse(signature_bytes[64]).unwrap();
+        let mut preimage = vec![];
+        preimage.extend_from_slice(task_id.as_bytes());
+        preimage.extend_from_slice(&result);
+        let message = Message::parse(&sha256hash(preimage));
+        assert!(verify(&message, &signature, &signer_pk), "Could not verify");
+
+        // recover verifying key (public key) from signature
+        let recovered_public_key =
+            recover(&message, &signature, &recid).expect("Could not recover");
+        assert_eq!(signer_pk, recovered_public_key, "Public key mismatch");
+    }
+}
diff --git a/src/utils/crypto.rs b/src/utils/crypto.rs
@@ -1,6 +1,7 @@
 use ecies::PublicKey;
+use eyre::{Context, Result};
 use libp2p_identity::Keypair;
-use libsecp256k1::{sign, Message, SecretKey};
+use libsecp256k1::{Message, SecretKey};
 use sha2::{Digest, Sha256};
 use sha3::Keccak256;
 
@@ -33,7 +34,7 @@ pub fn to_address(public_key: &PublicKey) -> [u8; 20] {
 #[inline]
 pub fn sign_bytes_recoverable(message: &[u8; 32], secret_key: &SecretKey) -> String {
     let message = Message::parse(message);
-    let (signature, recid) = sign(&message, secret_key);
+    let (signature, recid) = libsecp256k1::sign(&message, secret_key);
 
     format!(
         "{}{}",
@@ -42,6 +43,15 @@ pub fn sign_bytes_recoverable(message: &[u8; 32], secret_key: &SecretKey) -> Str
     )
 }
 
+/// Shorthand to encrypt bytes with a given public key.
+/// Returns hexadecimal encoded ciphertext.
+#[inline]
+pub fn encrypt_bytes(data: impl AsRef<[u8]>, public_key: &PublicKey) -> Result<String> {
+    ecies::encrypt(public_key.serialize().as_slice(), data.as_ref())
+        .wrap_err("could not encrypt data")
+        .map(hex::encode)
+}
+
 /// Converts a `libsecp256k1::SecretKey` to a `libp2p_identity::secp256k1::Keypair`.
 /// To do this, we serialize the secret key and create a new keypair from it.
 #[inline]
diff --git a/src/utils/message.rs b/src/utils/message.rs
