small rfks

erhant · erhant · commit 304edc6b8c2e · 2024-11-28T13:13:06.000+03:00
diff --git a/compute/src/node.rs b/compute/src/node.rs
@@ -5,7 +5,7 @@ use dkn_p2p::{
     },
     DriaP2PClient, DriaP2PCommander, DriaP2PProtocol,
 };
-use eyre::{eyre, Result};
+use eyre::Result;
 use tokio::{sync::mpsc, time::Duration};
 use tokio_util::{either::Either, sync::CancellationToken};
 
@@ -165,7 +165,10 @@ impl DriaComputeNode {
                 }
 
                 // first, parse the raw gossipsub message to a prepared message
-                let message = match self.parse_message_to_prepared_message(&message) {
+                let message = match DKNMessage::try_from_gossipsub_message(
+                    &message,
+                    &self.config.admin_public_key,
+                ) {
                     Ok(message) => message,
                     Err(e) => {
                         log::error!("Error parsing message: {:?}", e);
@@ -193,7 +196,7 @@ impl DriaComputeNode {
                     PingpongHandler::LISTEN_TOPIC => {
                         PingpongHandler::handle_ping(self, &message).await
                     }
-                    _ => unreachable!(), // unreachable because of the `match` above
+                    _ => unreachable!("unreachable due to match expression"),
                 };
 
                 // validate the message based on the result
@@ -216,46 +219,20 @@ impl DriaComputeNode {
         }
     }
 
-    /// Peer refresh simply reports the peer count to the user.
-    async fn handle_peer_refresh(&self) {
-        match self.p2p.peer_counts().await {
-            Ok((mesh, all)) => log::info!("Peer Count (mesh/all): {} / {}", mesh, all),
-            Err(e) => log::error!("Error getting peer counts: {:?}", e),
-        }
-    }
-
-    /// Updates the local list of available nodes by refreshing it.
-    /// Dials the RPC nodes again for better connectivity.
-    async fn handle_available_nodes_refresh(&mut self) {
-        log::info!("Refreshing available nodes.");
-
-        // refresh available nodes
-        if let Err(e) = self.available_nodes.populate_with_api().await {
-            log::error!("Error refreshing available nodes: {:?}", e);
-        };
-
-        // dial all rpc nodes
-        for rpc_addr in self.available_nodes.rpc_addrs.iter() {
-            log::debug!("Dialling RPC node: {}", rpc_addr);
-            if let Err(e) = self.p2p.dial(rpc_addr.clone()).await {
-                log::warn!("Error dialling RPC node: {:?}", e);
-            };
-        }
-    }
-
     /// Runs the main loop of the compute node.
     /// This method is not expected to return until cancellation occurs.
     pub async fn run(&mut self) -> Result<()> {
+        // prepare durations for sleeps
+        let peer_refresh_duration = Duration::from_secs(PEER_REFRESH_INTERVAL_SECS);
+        let available_node_refresh_duration =
+            Duration::from_secs(AVAILABLE_NODES_REFRESH_INTERVAL_SECS);
+
         // subscribe to topics
         self.subscribe(PingpongHandler::LISTEN_TOPIC).await?;
         self.subscribe(PingpongHandler::RESPONSE_TOPIC).await?;
         self.subscribe(WorkflowHandler::LISTEN_TOPIC).await?;
         self.subscribe(WorkflowHandler::RESPONSE_TOPIC).await?;
 
-        let peer_refresh_duration = Duration::from_secs(PEER_REFRESH_INTERVAL_SECS);
-        let available_node_refresh_duration =
-            Duration::from_secs(AVAILABLE_NODES_REFRESH_INTERVAL_SECS);
-
         loop {
             tokio::select! {
                 // check peer count every now and then
@@ -321,25 +298,31 @@ impl DriaComputeNode {
         Ok(())
     }
 
-    /// Parses a given raw Gossipsub message to a prepared P2PMessage object.
-    /// 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: &Message) -> 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 = DKNMessage::try_from(message)?;
-        log::debug!("Parsed: {}", message);
-
-        // 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."));
+    /// Peer refresh simply reports the peer count to the user.
+    async fn handle_peer_refresh(&self) {
+        match self.p2p.peer_counts().await {
+            Ok((mesh, all)) => log::info!("Peer Count (mesh/all): {} / {}", mesh, all),
+            Err(e) => log::error!("Error getting peer counts: {:?}", e),
         }
+    }
 
-        Ok(message)
+    /// Updates the local list of available nodes by refreshing it.
+    /// Dials the RPC nodes again for better connectivity.
+    async fn handle_available_nodes_refresh(&mut self) {
+        log::info!("Refreshing available nodes.");
+
+        // refresh available nodes
+        if let Err(e) = self.available_nodes.populate_with_api().await {
+            log::error!("Error refreshing available nodes: {:?}", e);
+        };
+
+        // dial all rpc nodes
+        for rpc_addr in self.available_nodes.rpc_addrs.iter() {
+            log::debug!("Dialling RPC node: {}", rpc_addr);
+            if let Err(e) = self.p2p.dial(rpc_addr.clone()).await {
+                log::warn!("Error dialling RPC node: {:?}", e);
+            };
+        }
     }
 }
 
diff --git a/compute/src/utils/message.rs b/compute/src/utils/message.rs
@@ -6,7 +6,7 @@ use crate::DRIA_COMPUTE_NODE_VERSION;
 use base64::{prelude::BASE64_STANDARD, Engine};
 use core::fmt;
 use ecies::PublicKey;
-use eyre::{Context, Result};
+use eyre::{eyre, Context, Result};
 use libsecp256k1::{verify, Message, SecretKey, Signature};
 use serde::{Deserialize, Serialize};
 
@@ -44,7 +44,7 @@ impl DKNMessage {
     ///
     /// - `data` is given as bytes, it is encoded into base64 to make up the `payload` within.
     /// - `topic` is the name of the [gossipsub topic](https://docs.libp2p.io/concepts/pubsub/overview/).
-    pub fn new(data: impl AsRef<[u8]>, topic: &str) -> Self {
+    pub(crate) fn new(data: impl AsRef<[u8]>, topic: &str) -> Self {
         Self {
             payload: BASE64_STANDARD.encode(data),
             topic: topic.to_string(),
@@ -55,7 +55,7 @@ impl DKNMessage {
     }
 
     /// Creates a new Message by signing the SHA256 of the payload, and prepending the signature.
-    pub fn new_signed(data: impl AsRef<[u8]>, topic: &str, signing_key: &SecretKey) -> Self {
+    pub(crate) fn new_signed(data: impl AsRef<[u8]>, topic: &str, signing_key: &SecretKey) -> Self {
         // sign the SHA256 hash of the data
         let signature_bytes = sign_bytes_recoverable(&sha256hash(data.as_ref()), signing_key);
 
@@ -69,19 +69,19 @@ impl DKNMessage {
     }
 
     /// Sets the identity of the message.
-    pub fn with_identity(mut self, identity: String) -> Self {
+    pub(crate) fn with_identity(mut self, identity: String) -> Self {
         self.identity = identity;
         self
     }
 
     /// Decodes the base64 payload into bytes.
     #[inline(always)]
-    pub fn decode_payload(&self) -> Result<Vec<u8>, base64::DecodeError> {
+    pub(crate) fn decode_payload(&self) -> Result<Vec<u8>, base64::DecodeError> {
         BASE64_STANDARD.decode(&self.payload)
     }
 
     /// Decodes and parses the base64 payload into JSON for the provided type `T`.
-    pub fn parse_payload<T: for<'a> Deserialize<'a>>(&self, signed: bool) -> Result<T> {
+    pub(crate) fn parse_payload<T: for<'a> Deserialize<'a>>(&self, signed: bool) -> Result<T> {
         let payload = self.decode_payload()?;
 
         let body = if signed {
@@ -96,7 +96,7 @@ impl DKNMessage {
     }
 
     /// Checks if the payload is signed by the given public key.
-    pub fn is_signed(&self, public_key: &PublicKey) -> Result<bool> {
+    pub(crate) fn is_signed(&self, public_key: &PublicKey) -> Result<bool> {
         // decode base64 payload
         let data = self.decode_payload()?;
 
@@ -116,6 +116,29 @@ impl DKNMessage {
         let digest = Message::parse(&sha256hash(body));
         Ok(verify(&digest, &signature, public_key))
     }
+
+    /// Tries to parse the given gossipsub message into a DKNMessage.
+    ///
+    /// This prepared message includes the topic, payload, version and timestamp.
+    /// It also checks the signature of the message, expecting a valid signature from admin node.
+    pub(crate) fn try_from_gossipsub_message(
+        gossipsub_message: &dkn_p2p::libp2p::gossipsub::Message,
+        public_key: &libsecp256k1::PublicKey,
+    ) -> Result<Self> {
+        // 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.", gossipsub_message.topic.as_str());
+        let message = serde_json::from_slice::<DKNMessage>(&gossipsub_message.data)
+            .wrap_err("could not parse message")?;
+        log::debug!("Parsed: {}", message);
+
+        // check dria signature
+        // NOTE: when we have many public keys, we should check the signature against all of them
+        if !message.is_signed(&public_key)? {
+            return Err(eyre!("Invalid signature."));
+        }
+
+        Ok(message)
+    }
 }
 
 impl fmt::Display for DKNMessage {
diff --git a/workflows/src/providers/ollama.rs b/workflows/src/providers/ollama.rs
@@ -4,7 +4,6 @@ use ollama_workflows::{
         generation::{
             completion::request::GenerationRequest,
             embeddings::request::{EmbeddingsInput, GenerateEmbeddingsRequest},
-            options::GenerationOptions,
         },
         Ollama,
     },
@@ -171,7 +170,7 @@ impl OllamaConfig {
             // otherwise, give error
             log::error!("Please download missing model with: ollama pull {}", model);
             log::error!("Or, set OLLAMA_AUTO_PULL=true to pull automatically.");
-            Err(eyre!("Required model not pulled in Ollama."))
+            Err(eyre!("required model not pulled in Ollama"))
         }
     }
 
