more tidyups and sanitizations

erhant · erhant · commit d3ca0b817243 · 2024-10-01T19:12:38.000+03:00
diff --git a/Makefile b/Makefile
@@ -37,10 +37,6 @@ profile-mem:
 version:
 	  @cargo pkgid | cut -d@ -f2
 
-.PHONY: ollama-cpu #  | Run Ollama CPU container
-ollama-cpu:
-		docker run -p=11434:11434 -v=${HOME}/.ollama:/root/.ollama ollama/ollama
-
 ###############################################################################
 .PHONY: test #         | Run tests
 test:
@@ -58,7 +54,7 @@ format:
 ###############################################################################
 .PHONY: docs #         | Generate & open crate documentation
 docs:
-		cargo doc --open --no-deps
+		cargo doc --open --no-deps --document-private-items
 
 .PHONY: env #          | Print active environment
 env:
diff --git a/src/config/ollama.rs b/src/config/ollama.rs
@@ -1,5 +1,3 @@
-use std::time::Duration;
-
 use eyre::{eyre, Result};
 use ollama_workflows::{
     ollama_rs::{
@@ -12,6 +10,7 @@ use ollama_workflows::{
     },
     Model,
 };
+use std::time::Duration;
 
 const DEFAULT_OLLAMA_HOST: &str = "http://127.0.0.1";
 const DEFAULT_OLLAMA_PORT: u16 = 11434;
diff --git a/src/handlers/workflow.rs b/src/handlers/workflow.rs
@@ -92,6 +92,7 @@ impl ComputeHandler for WorkflowHandler {
             .map(|prompt| Entry::try_value_or_str(&prompt));
 
         // execute workflow with cancellation
+        // TODO: is there a better way to handle this?
         let result: String;
         tokio::select! {
             _ = node.cancellation.cancelled() => {
@@ -111,7 +112,17 @@ impl ComputeHandler for WorkflowHandler {
         }
 
         // publish the result
-        node.send_result(result_topic, &task.public_key, &task.task_id, result)?;
+        let payload = P2PMessage::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);
+
+        node.publish(message)?;
+
         Ok(MessageAcceptance::Accept)
     }
 }
diff --git a/src/main.rs b/src/main.rs
@@ -1,4 +1,4 @@
-use dkn_compute::{DriaComputeNode, DriaComputeNodeConfig};
+use dkn_compute::*;
 use eyre::Result;
 use tokio_util::sync::CancellationToken;
 
@@ -17,12 +17,11 @@ async fn main() -> Result<()> {
 
 ██████╗ ██████╗ ██╗ █████╗ 
 ██╔══██╗██╔══██╗██║██╔══██╗   Dria Compute Node 
-██║  ██║██████╔╝██║███████║   v{}
+██║  ██║██████╔╝██║███████║   v{DRIA_COMPUTE_NODE_VERSION}
 ██║  ██║██╔══██╗██║██╔══██║   https://dria.co
 ██████╔╝██║  ██║██║██║  ██║
 ╚═════╝ ╚═╝  ╚═╝╚═╝╚═╝  ╚═╝
 "#,
-        dkn_compute::DRIA_COMPUTE_NODE_VERSION
     );
 
     let token = CancellationToken::new();
diff --git a/src/node.rs b/src/node.rs
@@ -185,7 +185,6 @@ impl DriaComputeNode {
                         // validate the message based on the result
                         match handle_result {
                             Ok(acceptance) => {
-
                                 self.p2p.validate_message(&message_id, &peer_id, acceptance)?;
                             },
                             Err(err) => {
@@ -235,32 +234,13 @@ impl DriaComputeNode {
 
         // check dria signature
         // NOTE: when we have many public keys, we should check the signature against all of them
+        // TODO: public key here will be given dynamically
         if !message.is_signed(&self.config.admin_public_key)? {
             return Err(eyre!("Invalid signature."));
         }
 
         Ok(message)
     }
-
-    /// Given a task with `id` and respective `public_key`, sign-then-encrypt the result.
-    pub fn send_result<R: AsRef<[u8]>>(
-        &mut self,
-        response_topic: &str,
-        public_key: &[u8],
-        task_id: &str,
-        result: R,
-    ) -> Result<()> {
-        let payload = P2PMessage::new_signed_encrypted_payload(
-            result.as_ref(),
-            task_id,
-            public_key,
-            &self.config.secret_key,
-        )?;
-        let payload_str = payload.to_string()?;
-        let message = P2PMessage::new(payload_str, response_topic);
-
-        self.publish(message)
-    }
 }
 
 #[cfg(test)]
diff --git a/src/p2p/client.rs b/src/p2p/client.rs
@@ -12,15 +12,17 @@ use libp2p_identity::Keypair;
 use std::time::Duration;
 use std::time::Instant;
 
+use super::*;
 use crate::utils::AvailableNodes;
 
-use super::{DriaBehaviour, DriaBehaviourEvent, P2P_KADEMLIA_PROTOCOL, P2P_PROTOCOL_STRING};
-
-/// Underlying libp2p client.
+/// P2P client, exposes a simple interface to handle P2P communication.
 pub struct P2PClient {
     /// `Swarm` instance, everything is accesses through this one.
     swarm: Swarm<DriaBehaviour>,
-    /// Peer count for (All, Mesh).
+    /// Peer count for All and Mesh peers.
+    ///
+    /// Mesh usually contains much fewer peers than All, as they are the ones
+    /// used for actual gossipping.
     peer_count: (usize, usize),
     /// Last time the peer count was refreshed.
     peer_last_refreshed: Instant,
@@ -253,7 +255,7 @@ impl P2PClient {
         if let Some(kad_protocol) = info
             .protocols
             .iter()
-            .find(|p| p.to_string().starts_with(P2P_KADEMLIA_PREFIX!()))
+            .find(|p| p.to_string().starts_with(P2P_KADEMLIA_PREFIX))
         {
             // if it matches our protocol, add it to the Kademlia routing table
             if *kad_protocol == P2P_KADEMLIA_PROTOCOL {
diff --git a/src/p2p/data_transform.rs b/src/p2p/data_transform.rs
@@ -1,3 +1,4 @@
+///! https://docs.rs/libp2p-gossipsub/latest/libp2p_gossipsub/trait.DataTransform.html
 use libp2p::gossipsub::{DataTransform, Message, RawMessage, TopicHash};
 use std::io::{Error, ErrorKind};
 use std::time::{SystemTime, UNIX_EPOCH};
diff --git a/src/p2p/mod.rs b/src/p2p/mod.rs
@@ -1,40 +1,10 @@
-use libp2p::StreamProtocol;
-
-/// Kademlia protocol prefix, as a macro so that it can be used in const macros.
-macro_rules! P2P_KADEMLIA_PREFIX {
-    () => {
-        "/dria/kad/"
-    };
-}
-
-/// Kademlia protocol prefix, as a macro so that it can be used in const macros.
-macro_rules! P2P_IDENTITY_PREFIX {
-    () => {
-        "dria/"
-    };
-}
-
-/// Kademlia protocol version, in the form of `/dria/kad/<version>`.
-/// Notice the `/` at the start.
-pub(crate) const P2P_KADEMLIA_PROTOCOL: StreamProtocol = StreamProtocol::new(concat!(
-    P2P_KADEMLIA_PREFIX!(),
-    env!("CARGO_PKG_VERSION_MAJOR"),
-    ".",
-    env!("CARGO_PKG_VERSION_MINOR")
-));
-
-/// Protocol string, checked by Identify protocol
-pub(crate) const P2P_PROTOCOL_STRING: &str = concat!(
-    P2P_IDENTITY_PREFIX!(),
-    env!("CARGO_PKG_VERSION_MAJOR"),
-    ".",
-    env!("CARGO_PKG_VERSION_MINOR")
-);
-
 mod behaviour;
 pub use behaviour::{DriaBehaviour, DriaBehaviourEvent};
 
 mod client;
 pub use client::P2PClient;
 
+mod versioning;
+pub use versioning::*;
+
 mod data_transform;
diff --git a/src/p2p/versioning.rs b/src/p2p/versioning.rs
@@ -0,0 +1,35 @@
+use libp2p::StreamProtocol;
+
+/// Kademlia protocol prefix, as a macro so that it can be used in literal-expecting constants.
+macro_rules! P2P_KADEMLIA_PREFIX {
+    () => {
+        "/dria/kad/"
+    };
+}
+pub const P2P_KADEMLIA_PREFIX: &str = P2P_KADEMLIA_PREFIX!();
+
+/// Identity protocol name prefix, as a macro so that it can be used in literal-expecting constants.
+macro_rules! P2P_IDENTITY_PREFIX {
+    () => {
+        "dria/"
+    };
+}
+
+/// Kademlia protocol version, in the form of `/dria/kad/<version>`, **notice the `/` at the start**.
+///
+/// It is important that this protocol matches EXACTLY among the nodes, otherwise there is a protocol-level logic
+/// that will prevent peers from finding eachother within the DHT.
+pub const P2P_KADEMLIA_PROTOCOL: StreamProtocol = StreamProtocol::new(concat!(
+    P2P_KADEMLIA_PREFIX!(),
+    env!("CARGO_PKG_VERSION_MAJOR"),
+    ".",
+    env!("CARGO_PKG_VERSION_MINOR")
+));
+
+/// Protocol string, checked by Identify protocol handlers.
+pub const P2P_PROTOCOL_STRING: &str = concat!(
+    P2P_IDENTITY_PREFIX!(),
+    env!("CARGO_PKG_VERSION_MAJOR"),
+    ".",
+    env!("CARGO_PKG_VERSION_MINOR")
+);
diff --git a/src/utils/crypto.rs b/src/utils/crypto.rs
@@ -22,14 +22,12 @@ pub fn keccak256hash(data: impl AsRef<[u8]>) -> [u8; 32] {
 /// and then (x || y) is hashed using Keccak256. The last 20 bytes of this hash is taken as the address.
 #[inline]
 pub fn to_address(public_key: &PublicKey) -> [u8; 20] {
-    let public_key_serial = &public_key.serialize()[1..];
+    let public_key_xy = &public_key.serialize()[1..];
     let mut addr = [0u8; 20];
-    addr.copy_from_slice(&keccak256hash(public_key_serial)[12..32]);
+    addr.copy_from_slice(&keccak256hash(public_key_xy)[12..32]);
     addr
 }
 
-// TODO: add peerId
-
 /// Shorthand to sign a digest (bytes) with node's secret key and return signature & recovery id
 /// serialized to 65 byte hex-string.
 #[inline]
diff --git a/src/utils/filter.rs b/src/utils/filter.rs
@@ -14,14 +14,15 @@ pub struct FilterPayload {
 
 impl FilterPayload {
     /// Shorthand function to create the underlying `BloomFilter` and check if it contains the given address.
-    #[inline]
     pub fn contains(&self, address: &[u8]) -> Result<bool> {
         BloomFilter::try_from(self)
             .map(|filter| filter.contains(address))
             .wrap_err("Could not create filter.")
     }
 }
 
+// FIXME: too many TryFrom's here, simplify in a single function here!
+
 impl TryFrom<&FilterPayload> for String {
     type Error = serde_json::Error;
 
@@ -109,28 +110,4 @@ mod tests {
 
         BloomFilter::try_from(&filter_payload).expect("Should parse filter");
     }
-
-    #[test]
-    fn test_filter_read_3() {
-        const FILTER_HEX: &str = "e7799ef73dcff3bc";
-        let filter_payload = FilterPayload {
-            hex: FILTER_HEX.to_string(),
-            hashes: 45,
-        };
-
-        let bf = BloomFilter::try_from(&filter_payload).expect("Should parse filter");
-
-        let addr = "6c460f37b78a2088a71e09dfb1c6238d7a34a346";
-        let addr_hex = hex::decode(addr).unwrap();
-        println!(
-            "{} ({}): {}",
-            addr,
-            addr_hex.len(),
-            bf.contains(&hex::decode(addr).unwrap())
-        );
-
-        let addr = "90c010775ffb266bc2136a13c08aa20ef11fe5cb";
-        let addr_hex = hex::decode(addr).unwrap();
-        println!("{} ({}): {}", addr, addr_hex.len(), bf.contains(&addr_hex));
-    }
 }
diff --git a/src/utils/message.rs b/src/utils/message.rs
@@ -3,7 +3,6 @@ use crate::utils::{
     get_current_time_nanos,
     payload::TaskResponsePayload,
 };
-
 use base64::{prelude::BASE64_STANDARD, Engine};
 use core::fmt;
 use ecies::PublicKey;
@@ -62,26 +61,30 @@ impl P2PMessage {
         Self::new(signed_payload, topic)
     }
 
-    /// Creates the payload of a computation result, as per Dria Whitepaper section 5.1 algorithm 2:
+    /// 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> {
-        // sign payload
+        // 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
+        // encrypt payload itself
         let ciphertext = ecies::encrypt(encrypting_public_key, payload.as_ref())?;
 
         Ok(TaskResponsePayload {
diff --git a/src/utils/payload.rs b/src/utils/payload.rs
@@ -10,7 +10,7 @@ use crate::utils::{filter::FilterPayload, get_current_time_nanos};
 ///
 /// To check the commitment, one must decrypt the ciphertext and parse plaintext from it,
 /// and compute the digest using SHA256. That digest will then be used for the signature check.
-#[derive(Serialize, Deserialize, Debug, Clone)]
+#[derive(Debug, Clone, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub struct TaskResponsePayload {
     /// A signature on the digest of plaintext result, prepended with task id.
@@ -30,7 +30,7 @@ impl TaskResponsePayload {
 }
 
 /// A generic task request, given by Dria.
-#[derive(Serialize, Deserialize, Debug, Clone)]
+#[derive(Debug, Clone, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub struct TaskRequestPayload<T> {
     /// The unique identifier of the task.

Original file line number	Diff line number	Diff line change
`@@ -1,3 +1,4 @@`
	`1`	`+///! https://docs.rs/libp2p-gossipsub/latest/libp2p_gossipsub/trait.DataTransform.html`
`1`	`2`	`use libp2p::gossipsub::{DataTransform, Message, RawMessage, TopicHash};`
`2`	`3`	`use std::io::{Error, ErrorKind};`
`3`	`4`	`use std::time::{SystemTime, UNIX_EPOCH};`