More tests and add signature verification

src/core/mod.rs

@@ -3,3 +3,6 @@ mod message;
 pub mod types;
 pub mod util;
 pub mod validations;
+
+#[cfg(test)]
+mod validations_test;

src/core/util.rs

@@ -37,6 +37,10 @@ pub fn get_farcaster_time() -> Result<u64, HubError> {
     Ok(to_farcaster_time(now.as_millis() as u64)?)
 }
 
+pub fn calculate_message_hash(data_bytes: &[u8]) -> Vec<u8> {
+    blake3::hash(data_bytes).as_bytes()[0..20].to_vec()
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

src/core/validations.rs

@@ -1,11 +1,12 @@
 use crate::proto;
 use crate::storage::util::{blake3_20, bytes_compare};
+use ed25519_dalek::{Signature, VerifyingKey};
 use prost::Message;
 use thiserror::Error;
 
 const MAX_DATA_BYTES: usize = 2048;
 
-#[derive(Error, Debug, Clone)]
+#[derive(Error, Debug, Clone, PartialEq)]
 pub enum ValidationError {
     #[error("Message data is missing")]
     MissingData,
@@ -15,23 +16,63 @@ pub enum ValidationError {
     InvalidHashScheme,
     #[error("Message data too large")]
     InvalidDataLength,
+    #[error("Unrecognized signature scheme")]
+    InvalidSignatureScheme,
+    #[error("Signer is empty or invalid")]
+    MissingOrInvalidSigner,
+    #[error("Signature is empty")]
+    MissingSignature,
+    #[error("Invalid message signature")]
+    InvalidSignature,
 }
 
 pub fn validate_message(message: &proto::Message) -> Result<(), ValidationError> {
+    let data_bytes;
     if message.data_bytes.is_some() {
-        let data_bytes = message.data_bytes.as_ref().unwrap();
+        data_bytes = message.data_bytes.as_ref().unwrap().clone();
         if data_bytes.len() > MAX_DATA_BYTES {
             return Err(ValidationError::InvalidDataLength);
         }
-        validate_message_hash(message.hash_scheme, data_bytes, &message.hash)?;
     } else {
         if message.data.is_none() {
             return Err(ValidationError::MissingData);
         }
-        let data_bytes = message.data.as_ref().unwrap().encode_to_vec();
-        validate_message_hash(message.hash_scheme, &data_bytes, &message.hash)?;
+        data_bytes = message.data.as_ref().unwrap().encode_to_vec();
     }
 
+    validate_message_hash(message.hash_scheme, &data_bytes, &message.hash)?;
+    validate_signature(
+        message.signature_scheme,
+        &message.hash,
+        &message.signature,
+        &message.signer,
+    )?;
+
+    Ok(())
+}
+
+fn validate_signature(
+    signature_scheme: i32,
+    data_bytes: &Vec<u8>,
+    signature: &Vec<u8>,
+    signer: &Vec<u8>,
+) -> Result<(), ValidationError> {
+    if signature_scheme != proto::SignatureScheme::Ed25519 as i32 {
+        return Err(ValidationError::InvalidSignatureScheme);
+    }
+
+    if signature.len() == 0 {
+        return Err(ValidationError::MissingSignature);
+    }
+
+    let sig = Signature::from_slice(signature).map_err(|_| ValidationError::InvalidSignature)?;
+    let public_key = VerifyingKey::try_from(signer.as_slice())
+        .map_err(|_| ValidationError::MissingOrInvalidSigner)?;
+
+    public_key
+        .verify_strict(data_bytes.as_slice(), &sig)
+        .map_err(|_| ValidationError::InvalidSignature)?;
+
     Ok(())
 }
 

src/core/validations_test.rs

@@ -0,0 +1,122 @@
+mod tests {
+    use crate::core::util::calculate_message_hash;
+    use crate::core::validations::{validate_message, ValidationError};
+    use crate::proto;
+    use crate::storage::store::test_helper;
+    use crate::utils::factory::{messages_factory, time};
+    use prost::Message;
+
+    fn assert_validation_error(msg: &proto::Message, expected_error: ValidationError) {
+        let result = validate_message(msg);
+        assert!(result.is_err());
+        assert_eq!(result.err().unwrap(), expected_error);
+    }
+
+    fn assert_valid(msg: &proto::Message) {
+        let result = validate_message(msg);
+        assert!(result.is_ok());
+    }
+
+    #[test]
+    fn test_validates_data_bytes() {
+        let mut msg = messages_factory::casts::create_cast_add(1234, "test", None, None);
+        assert_valid(&msg);
+
+        // Set data and data_bytes to None
+        msg.data = None;
+        msg.data_bytes = None;
+
+        assert_validation_error(&msg, ValidationError::MissingData);
+
+        msg.data_bytes = Some(vec![]);
+        assert_validation_error(&msg, ValidationError::MissingData);
+
+        // when data bytes is too large
+        msg.data_bytes = Some(vec![0; 2049]);
+        assert_validation_error(&msg, ValidationError::InvalidDataLength);
+
+        // When valid
+        let mut msg = messages_factory::casts::create_cast_add(1234, "test", None, None);
+        // Valid data, but empty data_bytes
+        msg.data_bytes = None;
+        assert_valid(&msg);
+
+        // Valid data_bytes, but empty data
+        msg.data_bytes = Some(msg.data.as_ref().unwrap().encode_to_vec());
+        msg.data = None;
+        assert_valid(&msg);
+    }
+
+    fn valid_message() -> proto::Message {
+        messages_factory::casts::create_cast_add(1234, "test", None, None)
+    }
+
+    #[test]
+    fn test_validates_hash_scheme() {
+        let mut msg = valid_message();
+        assert_valid(&msg);
+
+        msg.hash_scheme = 0;
+        assert_validation_error(&msg, ValidationError::InvalidHashScheme);
+
+        msg.hash_scheme = 2;
+        assert_validation_error(&msg, ValidationError::InvalidHashScheme);
+    }
+
+    #[test]
+    fn test_validates_hash() {
+        let timestamp = time::farcaster_time();
+        let mut msg = valid_message();
+        assert_valid(&msg);
+
+        msg.data.as_mut().unwrap().timestamp = timestamp + 10;
+        assert_validation_error(&msg, ValidationError::InvalidHash);
+
+        msg.hash = vec![];
+        assert_validation_error(&msg, ValidationError::InvalidHash);
+
+        msg.hash = vec![0; 20];
+        assert_validation_error(&msg, ValidationError::InvalidHash);
+    }
+
+    #[test]
+    fn validates_signature_scheme() {
+        let mut msg = valid_message();
+        assert_valid(&msg);
+
+        msg.signature_scheme = 0;
+        assert_validation_error(&msg, ValidationError::InvalidSignatureScheme);
+
+        msg.signature_scheme = 2;
+        assert_validation_error(&msg, ValidationError::InvalidSignatureScheme);
+    }
+
+    #[test]
+    fn validates_signature() {
+        let timestamp = time::farcaster_time();
+        let mut msg = valid_message();
+        assert_valid(&msg);
+
+        // Change the data so the signature becomes invalid
+        msg.data.as_mut().unwrap().timestamp = timestamp + 10;
+        msg.hash = calculate_message_hash(&msg.data.as_ref().unwrap().encode_to_vec()); // Ensure hash is valid
+        assert_validation_error(&msg, ValidationError::InvalidSignature);
+
+        msg.signature = vec![];
+        assert_validation_error(&msg, ValidationError::MissingSignature);
+
+        msg.signature = vec![0; 64];
+        assert_validation_error(&msg, ValidationError::InvalidSignature);
+
+        msg = valid_message();
+        msg.signer = vec![];
+
+        assert_validation_error(&msg, ValidationError::MissingOrInvalidSigner);
+
+        msg.signer = test_helper::generate_signer()
+            .verifying_key()
+            .to_bytes()
+            .to_vec();
+        assert_validation_error(&msg, ValidationError::InvalidSignature);
+    }
+}

src/storage/store/engine_tests.rs

@@ -1,5 +1,6 @@
 #[cfg(test)]
 mod tests {
+    use crate::core::util::calculate_message_hash;
     use crate::proto::ShardChunk;
     use crate::proto::{self, ReactionType};
     use crate::proto::{HubEvent, ValidatorMessage};
@@ -13,6 +14,7 @@ mod tests {
     use crate::storage::trie::merkle_trie::TrieKey;
     use crate::utils::factory::{self, events_factory, messages_factory, time, username_factory};
     use ed25519_dalek::SigningKey;
+    use prost::Message;
 
     fn from_hex(s: &str) -> Vec<u8> {
         hex::decode(s).unwrap()
@@ -194,6 +196,27 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn test_engine_rejects_message_with_invalid_signature() {
+        let (mut engine, _tmpdir) = test_helper::new_engine();
+        register_user(FID_FOR_TEST, test_helper::default_signer(), &mut engine).await;
+        let mut message = default_message("msg1");
+        let current_timestamp = message.data.as_ref().unwrap().timestamp;
+        // Modify the message so the signatures is no longer correct
+        message.data.as_mut().unwrap().timestamp = current_timestamp + 1;
+        message.hash = calculate_message_hash(&message.data.as_ref().unwrap().encode_to_vec());
+
+        assert_commit_fails(&mut engine, &message).await;
+
+        assert_eq!(
+            engine
+                .validate_user_message(&message, &mut RocksDbTransactionBatch::new())
+                .unwrap_err()
+                .to_string(),
+            "Invalid message signature"
+        );
+    }
+
     #[tokio::test]
     async fn test_engine_commit_no_messages_happy_path() {
         let (mut engine, _tmpdir) = test_helper::new_engine();

src/storage/store/test_helper.rs

@@ -260,6 +260,11 @@ pub fn default_signer() -> SigningKey {
     )
 }
 
+#[allow(dead_code)]
+pub fn generate_signer() -> SigningKey {
+    SigningKey::generate(&mut rand::thread_rng())
+}
+
 #[allow(dead_code)]
 pub fn enable_logging() {
     let env_filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info"));

src/utils/factory.rs

@@ -164,6 +164,7 @@ pub mod events_factory {
 
 pub mod messages_factory {
     use super::*;
+    use crate::core::util::calculate_message_hash;
 
     pub fn farcaster_time() -> u32 {
         (std::time::SystemTime::now()
@@ -202,7 +203,7 @@ pub mod messages_factory {
         };
 
         let msg_data_bytes = msg_data.encode_to_vec();
-        let hash = blake3::hash(&msg_data_bytes).as_bytes()[0..20].to_vec();
+        let hash = calculate_message_hash(&msg_data_bytes);
 
         let signature = key.sign(&hash).to_bytes();
         message::Message {