Address review comments (#14)

Author: danimhr

Cargo.lock

@@ -33,3 +33,7 @@ wormhole-vaas-serde = "0.1.0"
 [dev-dependencies]
 serde_json = "1.0.140"
 base64 = "0.22.1"
+mockall = "0.13.1"
+
+[profile.release]
+overflow-checks = true

Cargo.toml

@@ -61,6 +61,7 @@ const INVALID_UNRELIABLE_DATA_FORMAT: &str = "Invalid unreliable data format";
 const INVALID_PDA_MESSAGE: &str = "Invalid PDA message";
 const INVALID_EMITTER_CHAIN: &str = "Invalid emitter chain";
 const INVALID_ACCUMULATOR_ADDRESS: &str = "Invalid accumulator address";
+const INVALID_VAA_VERSION: &str = "Invalid VAA version";
 
 fn decode_and_verify_update(
     wormhole_pid: &Pubkey,
@@ -87,6 +88,14 @@ fn decode_and_verify_update(
             anyhow::anyhow!(format!("{}: {}", INVALID_UNRELIABLE_DATA_FORMAT, e))
         })?;
 
+    if unreliable_data.message.vaa_version != 1 {
+        tracing::error!(
+            vaa_version = unreliable_data.message.vaa_version,
+            "Unsupported VAA version"
+        );
+        return Err(anyhow::anyhow!(INVALID_VAA_VERSION));
+    }
+
     if Chain::Pythnet != unreliable_data.emitter_chain.into() {
         tracing::error!(
             emitter_chain = unreliable_data.emitter_chain,
@@ -320,7 +329,7 @@ mod tests {
     use secp256k1::SecretKey;
     use solana_account_decoder::{UiAccount, UiAccountData};
 
-    use crate::posted_message::MessageData;
+    use crate::{posted_message::MessageData, signer::MockSigner};
 
     fn get_wormhole_pid() -> Pubkey {
         Pubkey::from_str("H3fxXJ86ADW2PNuDDmZJg6mzTtPxkYCpNuQUTgmJ7AjU").unwrap()
@@ -488,6 +497,18 @@ mod tests {
         assert_eq!(result.unwrap_err().to_string(), INVALID_EMITTER_CHAIN);
     }
 
+    #[test]
+    fn test_decode_and_verify_update_invalid_vaa_version() {
+        let mut unreliable_data = get_unreliable_data();
+        unreliable_data.message.vaa_version = 2;
+        let result = decode_and_verify_update(
+            &get_wormhole_pid(),
+            &get_accumulator_address(),
+            get_update(unreliable_data),
+        );
+        assert_eq!(result.unwrap_err().to_string(), INVALID_VAA_VERSION);
+    }
+
     #[test]
     fn test_decode_and_verify_update_invalid_emitter_address() {
         let mut unreliable_data = get_unreliable_data();
@@ -530,4 +551,21 @@ mod tests {
             "30e41be3f10d3ac813f91e49e189bbb948d030be",
         );
     }
+
+    #[tokio::test]
+    async fn test_sign_error() {
+        let mut mock_signer = MockSigner::new();
+        mock_signer
+            .expect_sign()
+            .return_once(|_| Err(anyhow::anyhow!("Mock signing error")));
+        let unreliable_data = get_unreliable_data();
+        let body = message_data_to_body(&unreliable_data);
+        let err = Observation::try_new(body, Arc::new(mock_signer))
+            .await
+            .unwrap_err();
+        assert_eq!(
+            err.to_string(),
+            "Failed to sign observation: Mock signing error"
+        );
+    }
 }

src/main.rs

@@ -127,4 +127,55 @@ mod tests {
         let decoded = PostedMessageUnreliableData::try_from_slice(encoded.as_slice()).unwrap();
         assert_eq!(decoded, post_message_unreliable_data);
     }
+
+    #[test]
+    fn test_invalid_magic() {
+        let post_message_unreliable_data = PostedMessageUnreliableData {
+            message: MessageData {
+                vaa_version: 1,
+                consistency_level: 2,
+                vaa_time: 3,
+                vaa_signature_account: [4u8; 32],
+                submission_time: 5,
+                nonce: 6,
+                sequence: 7,
+                emitter_chain: 8,
+                emitter_address: [9u8; 32],
+                payload: vec![10u8; 32],
+            },
+        };
+
+        let mut encoded = borsh::to_vec(&post_message_unreliable_data).unwrap();
+        encoded[0..3].copy_from_slice(b"foo"); // Invalid magic
+
+        let err = PostedMessageUnreliableData::try_from_slice(encoded.as_slice()).unwrap_err();
+        assert_eq!(
+            err.to_string(),
+            "Magic mismatch. Expected [109, 115, 117] but got [102, 111, 111]"
+        );
+    }
+
+    #[test]
+    fn test_invalid_data_length() {
+        let post_message_unreliable_data = PostedMessageUnreliableData {
+            message: MessageData {
+                vaa_version: 1,
+                consistency_level: 2,
+                vaa_time: 3,
+                vaa_signature_account: [4u8; 32],
+                submission_time: 5,
+                nonce: 6,
+                sequence: 7,
+                emitter_chain: 8,
+                emitter_address: [9u8; 32],
+                payload: vec![10u8; 32],
+            },
+        };
+
+        let mut encoded = borsh::to_vec(&post_message_unreliable_data).unwrap();
+        encoded = encoded[0..encoded.len() - 1].to_vec();
+
+        let err = PostedMessageUnreliableData::try_from_slice(encoded.as_slice()).unwrap_err();
+        assert_eq!(err.to_string(), "Unexpected length of input");
+    }
 }

src/posted_message.rs

@@ -1,5 +1,5 @@
 use der::{
-    asn1::{AnyRef, BitStringRef, UintRef},
+    asn1::{AnyRef, BitStringRef},
     oid::ObjectIdentifier,
     Decode, Sequence,
 };
@@ -13,12 +13,13 @@ use std::{
 use async_trait::async_trait;
 use prost::Message as ProstMessage;
 use secp256k1::{
-    ecdsa::{RecoverableSignature, RecoveryId},
+    ecdsa::{RecoverableSignature, RecoveryId, Signature},
     Message, PublicKey, Secp256k1, SecretKey,
 };
 use sequoia_openpgp::armor::{Kind, Reader, ReaderMode};
 use sha3::{Digest, Keccak256};
 
+#[cfg_attr(test, mockall::automock)]
 #[async_trait]
 pub trait Signer: Send + Sync {
     async fn sign(&self, data: [u8; 32]) -> anyhow::Result<[u8; 65]>;
@@ -181,12 +182,6 @@ pub struct SubjectPublicKeyInfo<'a> {
     pub subject_public_key: BitStringRef<'a>,
 }
 
-#[derive(Sequence)]
-struct EcdsaSignature<'a> {
-    r: UintRef<'a>,
-    s: UintRef<'a>,
-}
-
 #[async_trait]
 impl Signer for KMSSigner {
     async fn sign(&self, data: [u8; 32]) -> anyhow::Result<[u8; 65]> {
@@ -204,14 +199,13 @@ impl Signer for KMSSigner {
             .signature
             .ok_or_else(|| anyhow::anyhow!("KMS did not return a signature"))?;
 
-        let decoded_signature = EcdsaSignature::from_der(kms_signature.as_ref())
-            .map_err(|e| anyhow::anyhow!("Failed to decode SubjectPublicKeyInfo: {}", e))?;
-
-        let r_bytes = decoded_signature.r.as_bytes();
-        let s_bytes = decoded_signature.s.as_bytes();
-        let mut signature = [0u8; 65];
-        signature[(32 - r_bytes.len())..32].copy_from_slice(r_bytes);
-        signature[(64 - s_bytes.len())..64].copy_from_slice(decoded_signature.s.as_bytes());
+        let mut signature = Signature::from_der(kms_signature.as_ref())
+            .map_err(|e| anyhow::anyhow!("Failed to decode signature from KMS: {}", e))?;
+        // NOTE: AWS KMS does not guarantee that the ECDSA signature is normalized.
+        // Therefore, we must normalize it ourselves to prevent malleability,
+        // so that it can be successfully verified later using the secp256k1 standard libraries.
+        signature.normalize_s();
+        let signature_bytes = signature.serialize_compact();
 
         let public_key = self.get_public_key()?;
         for raw_id in 0..4 {
@@ -220,10 +214,12 @@ impl Signer for KMSSigner {
                 .map_err(|e| anyhow::anyhow!("Failed to create RecoveryId: {}", e))?;
             if let Ok(recovered_public_key) = secp.recover_ecdsa(
                 &Message::from_digest(data),
-                &RecoverableSignature::from_compact(&signature[..64], recid)
+                &RecoverableSignature::from_compact(&signature_bytes, recid)
                     .map_err(|e| anyhow::anyhow!("Failed to create RecoverableSignature: {}", e))?,
             ) {
                 if recovered_public_key == public_key.0 {
+                    let mut signature = [0u8; 65];
+                    signature[..64].copy_from_slice(&signature_bytes);
                     signature[64] = raw_id as u8;
                     return Ok(signature);
                 }