Merge pull request #4 from AIBlockOfficial/v0.1.3

Merge v0.1.3 to Main

Author: BHouwens

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "valence_core"
-version = "0.1.2"
+version = "0.1.3"
 edition = "2021"
 license = "MIT"
 keywords = ["blockchain", "L2", "peer-to-peer", "P2P"]
@@ -24,4 +24,6 @@ serde_json = "1.0.103"
 hex = "0.4.3"
 cuckoofilter = "0.5.0"
 tracing = "0.1.37"
+tracing-subscriber = "0.3.17"
+tracing-futures = "0.2.3"
 sha3 = "0.10.8"

src/api/utils.rs

@@ -5,6 +5,7 @@ use std::convert::Infallible;
 use warp::{Filter, Future, Rejection, Reply};
 
 impl warp::reject::Reject for ApiErrorType {}
+use tracing::{debug, warn};
 
 /// Clone component/struct to use in route
 ///
@@ -61,27 +62,38 @@ pub fn sig_verify_middleware() -> impl Filter<Extract = ((),), Error = Rejection
         .and(warp::header::headers_cloned())
         .and_then(
             move |_: warp::path::FullPath, headers: warp::hyper::HeaderMap| {
+                debug!("Validating signature");
+
                 async move {
                     let public_key = headers
                         .get("public_key")
                         .and_then(|n| n.to_str().ok())
                         .unwrap_or_default();
 
+                    debug!("public_key: {:?}", public_key);
+
                     let address = headers
                         .get("address")
                         .and_then(|n| n.to_str().ok())
                         .unwrap_or_default();
 
+                    debug!("address: {:?}", address);
+
                     let signature = headers
                         .get("signature")
                         .and_then(|n| n.to_str().ok())
                         .unwrap_or_default();
 
+                    debug!("signature: {:?}", signature);
+
                     if validate_signature(public_key, address, signature) {
+                        debug!("Signature is valid");
+
                         // Proceed to the next filter/handler
                         return Ok(());
                     }
 
+                    warn!("Invalid signature");
                     Err(warp::reject::custom(ApiErrorType::InvalidSignature))
                 }
             },

src/db/handler.rs

@@ -19,12 +19,36 @@ pub trait KvStoreConnection {
     ///
     /// * `key` - Key of the data entry to set
     /// * `value` - Value of the data entry to set
-    async fn set_data<T: Serialize + Send>(
+    async fn set_data<T: Serialize + Send + DeserializeOwned>(
         &mut self,
         key: &str,
         value: T,
     ) -> Result<(), Box<dyn std::error::Error + Send + Sync>>;
 
+    /// Sets a data entry in the cache with an expiration time
+    ///
+    /// ### Arguments
+    ///
+    /// * `key` - Key of the data entry to set
+    /// * `value` - Value of the data entry to set
+    /// * `seconds` - Number of seconds to expire the data entry in
+    async fn set_data_with_expiry<T: Serialize + DeserializeOwned + Send>(
+        &mut self,
+        key: &str,
+        value: T,
+        seconds: usize,
+    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>>;
+
+    /// Deletes a data entry from the cache
+    ///
+    /// ### Arguments
+    ///
+    /// * `key` - Key of the data entry to delete
+    async fn delete_data(
+        &mut self,
+        key: &str,
+    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>>;
+
     /// Gets a data entry from the cache
     ///
     /// ### Arguments
@@ -33,7 +57,7 @@ pub trait KvStoreConnection {
     async fn get_data<T: DeserializeOwned>(
         &mut self,
         key: &str,
-    ) -> Result<Option<T>, Box<dyn std::error::Error + Send + Sync>>;
+    ) -> Result<Option<Vec<T>>, Box<dyn std::error::Error + Send + Sync>>;
 }
 
 #[async_trait]

src/db/mongo_db.rs

@@ -1,5 +1,5 @@
 use async_trait::async_trait;
-use mongodb::bson::{doc, Document};
+use mongodb::bson::{doc, DateTime, Document};
 use mongodb::{options::ClientOptions, Client};
 use serde::{de::DeserializeOwned, Serialize};
 use tracing::{event, span, trace, warn, Level};
@@ -18,6 +18,31 @@ pub struct MongoDbConn {
     pub index: MongoDbIndex,
 }
 
+impl MongoDbConn {
+    /// Creates a TTL index on the expiry field.
+    ///
+    /// NOTE: This function will need to be called in the main function when initialising a MongoDB connection.
+    pub async fn create_ttl_index(&self) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
+        let collection = self
+            .client
+            .database(&self.index.db_name)
+            .collection::<Document>(&self.index.coll_name);
+
+        // Create TTL index on the 'expiry' field
+        let index_model = mongodb::IndexModel::builder()
+            .keys(doc! { "expiry": 1 })
+            .options(Some(
+                mongodb::options::IndexOptions::builder()
+                    .expire_after(Some(std::time::Duration::from_secs(0)))
+                    .build(),
+            ))
+            .build();
+
+        collection.create_index(index_model, None).await?;
+        Ok(())
+    }
+}
+
 #[async_trait]
 impl KvStoreConnection for MongoDbConn {
     async fn init(url: &str) -> Result<Self, Box<dyn std::error::Error + Send + Sync>> {
@@ -47,7 +72,7 @@ impl KvStoreConnection for MongoDbConn {
         Ok(MongoDbConn { client, index })
     }
 
-    async fn set_data<T: Serialize + std::marker::Send>(
+    async fn set_data<T: Serialize + std::marker::Send + DeserializeOwned>(
         &mut self,
         key: &str,
         value: T,
@@ -61,20 +86,32 @@ impl KvStoreConnection for MongoDbConn {
             .database(&self.index.db_name)
             .collection::<Document>(&self.index.coll_name);
 
-        let document = match mongodb::bson::to_document(&value) {
-            Ok(document) => document,
-            Err(e) => {
-                event!(Level::ERROR, "Failed to serialize data with error: {e}");
-                Document::new()
-            }
+        // Check if the document with the given key exists
+        let filter = doc! { "_id": key };
+        let existing_doc = collection.find_one(filter.clone(), None).await?;
+
+        let mut vec: Vec<T> = if let Some(doc) = existing_doc {
+            // Deserialize the existing data
+            mongodb::bson::from_bson(doc.get("data").unwrap().clone())?
+        } else {
+            Vec::new()
         };
 
-        let filter = doc! { "_id": key };
+        // Append the new data to the vec
+        vec.push(value);
+
+        // Serialize the vec back to a BSON array
+        let serialized_vec = mongodb::bson::to_bson(&vec)?;
+
+        // Create or update the document
+        let update = doc! {
+            "$set": { "data": serialized_vec }
+        };
         match collection
-            .replace_one(
+            .update_one(
                 filter,
-                document.clone(),
-                mongodb::options::ReplaceOptions::builder()
+                update,
+                mongodb::options::UpdateOptions::builder()
                     .upsert(true)
                     .build(),
             )
@@ -86,32 +123,122 @@ impl KvStoreConnection for MongoDbConn {
             }
         };
 
-        trace!("Data set successfully");
+        trace!("Data set successfully with expiry");
+
+        Ok(())
+    }
+
+    async fn set_data_with_expiry<T: Serialize + std::marker::Send + DeserializeOwned>(
+        &mut self,
+        key: &str,
+        value: T,
+        seconds: usize,
+    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
+        // Tracing
+        let span = span!(Level::TRACE, "MongoDbConn::set_data_with_expiry");
+        let _enter = span.enter();
+
+        let collection = self
+            .client
+            .database(&self.index.db_name)
+            .collection::<Document>(&self.index.coll_name);
+
+        // Check if the document with the given key exists
+        let filter = doc! { "_id": key };
+        let existing_doc = collection.find_one(filter.clone(), None).await?;
+
+        let mut vec: Vec<T> = if let Some(doc) = existing_doc {
+            // Deserialize the existing data
+            mongodb::bson::from_bson(doc.get("data").unwrap().clone())?
+        } else {
+            Vec::new()
+        };
+
+        // Append the new data to the vec
+        vec.push(value);
+
+        // Serialize the vec back to a BSON array
+        let serialized_vec = mongodb::bson::to_bson(&vec)?;
+
+        // Calculate the expiry time
+        let expiry_time = (seconds * 1000) as i64;
+        let expiry_bson_datetime = DateTime::from_millis(expiry_time);
+
+        // Create or update the document with the new expiry time
+        let update = doc! {
+            "$set": {
+                "data": serialized_vec,
+                "expiry": expiry_bson_datetime,
+            }
+        };
+        collection
+            .update_one(
+                filter,
+                update,
+                mongodb::options::UpdateOptions::builder()
+                    .upsert(true)
+                    .build(),
+            )
+            .await?;
+
+        trace!("Data set successfully with expiry");
+
+        Ok(())
+    }
+
+    async fn delete_data(
+        &mut self,
+        key: &str,
+    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
+        // Tracing
+        let span = span!(Level::TRACE, "MongoDbConn::delete_data");
+        let _enter = span.enter();
+
+        let collection = self
+            .client
+            .database(&self.index.db_name)
+            .collection::<Document>(&self.index.coll_name);
+
+        let filter = doc! { "_id": key };
+        match collection.delete_one(filter, None).await {
+            Ok(_) => (),
+            Err(e) => {
+                event!(Level::ERROR, "Failed to delete data with error: {e}");
+            }
+        };
+
+        trace!("Data deleted successfully");
 
         Ok(())
     }
 
     async fn get_data<T: DeserializeOwned>(
         &mut self,
         key: &str,
-    ) -> Result<Option<T>, Box<dyn std::error::Error + Send + Sync>> {
+    ) -> Result<Option<Vec<T>>, Box<dyn std::error::Error + Send + Sync>> {
         // Tracing
         let span = span!(Level::TRACE, "MongoDbConn::get_data");
         let _enter = span.enter();
 
         let collection = self
             .client
             .database(&self.index.db_name)
-            .collection::<Document>(&self.index.coll_name); // Change to your actual collection name
+            .collection::<Document>(&self.index.coll_name);
 
+        // Check if the document with the given key exists
         let filter = doc! { "_id": key };
-        let result = collection.find_one(filter, None).await?;
-
-        trace!("Data retrieved successfully");
+        let doc_find = match collection.find_one(filter.clone(), None).await {
+            Ok(doc) => doc,
+            Err(e) => {
+                event!(Level::ERROR, "Failed to get data with error: {e}");
+                return Ok(None);
+            }
+        };
 
-        if let Some(document) = result {
-            let deserialized: T = mongodb::bson::from_document(document)?;
-            return Ok(Some(deserialized));
+        if let Some(doc) = doc_find {
+            // Deserialize the existing data
+            let vec: Vec<T> = mongodb::bson::from_bson(doc.get("data").unwrap().clone())?;
+            return Ok(Some(vec));
         }
 
         warn!("Data unsuccessfully deserialized");