Lazer solana audit fixes (#2250)

Riateche · web-flow · commit eb7f460ab8d1 · 2025-01-14T15:40:30.000Z
* fix(lazer): verify account size before applying migration

* fix(lazer): fetch message offset and data from sysvar and verify message data

* test(lazer): add test with wrong message offset
diff --git a/lazer/contracts/solana/programs/pyth-lazer-solana-contract/src/lib.rs b/lazer/contracts/solana/programs/pyth-lazer-solana-contract/src/lib.rs
@@ -103,6 +103,9 @@ pub mod pyth_lazer_solana_contract {
         if old_data[0..ANCHOR_DISCRIMINATOR_BYTES] != Storage::DISCRIMINATOR {
             return Err(ProgramError::InvalidAccountData.into());
         }
+        if old_data.len() != StorageV010::SERIALIZED_LEN + ANCHOR_DISCRIMINATOR_BYTES {
+            return Err(ProgramError::InvalidAccountData.into());
+        }
         let old_storage = StorageV010::deserialize(&mut &old_data[ANCHOR_DISCRIMINATOR_BYTES..])?;
         if old_storage.top_authority != ctx.accounts.top_authority.key() {
             return Err(ProgramError::MissingRequiredSignature.into());
@@ -194,7 +197,6 @@ pub mod pyth_lazer_solana_contract {
         message_data: Vec<u8>,
         ed25519_instruction_index: u16,
         signature_index: u8,
-        message_offset: u16,
     ) -> Result<VerifiedMessage> {
         system_program::transfer(
             CpiContext::new(
@@ -213,7 +215,6 @@ pub mod pyth_lazer_solana_contract {
             &message_data,
             ed25519_instruction_index,
             signature_index,
-            message_offset,
         )
         .map_err(|err| {
             msg!("signature verification error: {:?}", err);
diff --git a/lazer/contracts/solana/programs/pyth-lazer-solana-contract/src/signature.rs b/lazer/contracts/solana/programs/pyth-lazer-solana-contract/src/signature.rs
@@ -2,7 +2,9 @@ use {
     crate::Storage,
     anchor_lang::{
         prelude::{borsh, AccountInfo, Clock, ProgramError, Pubkey, SolanaSysvar},
-        solana_program::{ed25519_program, pubkey::PUBKEY_BYTES, sysvar},
+        solana_program::{
+            ed25519_program, program_memory::sol_memcmp, pubkey::PUBKEY_BYTES, sysvar,
+        },
         AnchorDeserialize, AnchorSerialize,
     },
     bytemuck::{cast_slice, checked::try_cast_slice, Pod, Zeroable},
@@ -127,6 +129,8 @@ pub enum SignatureVerificationError {
     InvalidStorageData,
     #[error("not a trusted signer")]
     NotTrustedSigner,
+    #[error("invalid message data")]
+    InvalidMessageData,
 }
 
 impl From<SignatureVerificationError> for ProgramError {
@@ -148,6 +152,10 @@ impl From<SignatureVerificationError> for anchor_lang::error::Error {
     }
 }
 
+fn slice_eq(a: &[u8], b: &[u8]) -> bool {
+    a.len() == b.len() && sol_memcmp(a, b, a.len()) == 0
+}
+
 /// Verifies a ed25519 signature on Solana by checking that the transaction contains
 /// a correct call to the built-in `ed25519_program`.
 ///
@@ -163,7 +171,6 @@ pub fn verify_message(
     message_data: &[u8],
     ed25519_instruction_index: u16,
     signature_index: u8,
-    message_offset: u16,
 ) -> Result<VerifiedMessage, SignatureVerificationError> {
     const SOLANA_FORMAT_MAGIC_LE: u32 = 2182742457;
 
@@ -175,25 +182,25 @@ pub fn verify_message(
         return Err(SignatureVerificationError::Ed25519InstructionMustPrecedeCurrentInstruction);
     }
 
-    let instruction = sysvar::instructions::load_instruction_at_checked(
+    let ed25519_instruction = sysvar::instructions::load_instruction_at_checked(
         ed25519_instruction_index.into(),
         instructions_sysvar,
     )
     .map_err(SignatureVerificationError::LoadInstructionAtFailed)?;
 
-    if instruction.program_id != ed25519_program::ID {
+    if ed25519_instruction.program_id != ed25519_program::ID {
         return Err(SignatureVerificationError::InvalidEd25519InstructionProgramId);
     }
-    if instruction.data.len() < ED25519_PROGRAM_INPUT_HEADER_LEN {
+    if ed25519_instruction.data.len() < ED25519_PROGRAM_INPUT_HEADER_LEN {
         return Err(SignatureVerificationError::InvalidEd25519InstructionDataLength);
     }
 
-    let num_signatures = instruction.data[0];
+    let num_signatures = ed25519_instruction.data[0];
     if signature_index >= num_signatures {
         return Err(SignatureVerificationError::InvalidSignatureIndex);
     }
     let args: &[Ed25519SignatureOffsets] =
-        try_cast_slice(&instruction.data[ED25519_PROGRAM_INPUT_HEADER_LEN..])
+        try_cast_slice(&ed25519_instruction.data[ED25519_PROGRAM_INPUT_HEADER_LEN..])
             .map_err(|_| SignatureVerificationError::InvalidEd25519InstructionDataLength)?;
 
     let args_len = args
@@ -205,19 +212,37 @@ pub fn verify_message(
     }
     let offsets = &args[usize::from(signature_index)];
 
-    let expected_signature_offset = message_offset
-        .checked_add(MAGIC_LEN)
+    let message_offset = offsets
+        .signature_offset
+        .checked_sub(MAGIC_LEN)
         .ok_or(SignatureVerificationError::MessageOffsetOverflow)?;
-    if offsets.signature_offset != expected_signature_offset {
-        return Err(SignatureVerificationError::InvalidSignatureOffset);
+
+    let self_instruction = sysvar::instructions::load_instruction_at_checked(
+        self_instruction_index.into(),
+        instructions_sysvar,
+    )
+    .map_err(SignatureVerificationError::LoadInstructionAtFailed)?;
+
+    let message_end_offset = offsets
+        .message_data_offset
+        .checked_add(offsets.message_data_size)
+        .ok_or(SignatureVerificationError::MessageOffsetOverflow)?;
+    let expected_message_data = self_instruction
+        .data
+        .get(usize::from(message_offset)..usize::from(message_end_offset))
+        .ok_or(SignatureVerificationError::InvalidMessageOffset)?;
+
+    if !slice_eq(expected_message_data, message_data) {
+        return Err(SignatureVerificationError::InvalidMessageData);
     }
 
     let magic = LE::read_u32(&message_data[..MAGIC_LEN.into()]);
     if magic != SOLANA_FORMAT_MAGIC_LE {
         return Err(SignatureVerificationError::FormatMagicMismatch);
     }
 
-    let expected_public_key_offset = expected_signature_offset
+    let expected_public_key_offset = offsets
+        .signature_offset
         .checked_add(SIGNATURE_LEN)
         .ok_or(SignatureVerificationError::MessageOffsetOverflow)?;
     if offsets.public_key_offset != expected_public_key_offset {
diff --git a/lazer/contracts/solana/programs/pyth-lazer-solana-contract/tests/test1.rs b/lazer/contracts/solana/programs/pyth-lazer-solana-contract/tests/test1.rs
@@ -1,17 +1,17 @@
 use {
     anchor_lang::{prelude::AccountMeta, InstructionData},
     pyth_lazer_solana_contract::{ed25519_program_args, ANCHOR_DISCRIMINATOR_BYTES},
-    solana_program_test::{BanksClient, ProgramTest},
+    solana_program_test::{BanksClient, BanksClientError, ProgramTest},
     solana_sdk::{
         account::Account,
         ed25519_program,
         hash::Hash,
-        instruction::Instruction,
+        instruction::{Instruction, InstructionError},
         pubkey::{Pubkey, PUBKEY_BYTES},
         signature::Keypair,
         signer::Signer,
         system_instruction, system_program, system_transaction, sysvar,
-        transaction::Transaction,
+        transaction::{Transaction, TransactionError},
     },
     std::env,
 };
@@ -103,24 +103,45 @@ impl Setup {
     }
 
     async fn verify_message(&mut self, message: &[u8], treasury: Pubkey) {
+        let treasury_starting_lamports = self
+            .banks_client
+            .get_account(treasury)
+            .await
+            .unwrap()
+            .unwrap()
+            .lamports;
+
+        // 8 bytes for Anchor header, 4 bytes for Vec length.
+        self.verify_message_with_offset(message, treasury, 12)
+            .await
+            .unwrap();
+
+        assert_eq!(
+            self.banks_client
+                .get_account(treasury)
+                .await
+                .unwrap()
+                .unwrap()
+                .lamports,
+            treasury_starting_lamports + 1,
+        );
+    }
+
+    async fn verify_message_with_offset(
+        &mut self,
+        message: &[u8],
+        treasury: Pubkey,
+        message_offset: u16,
+    ) -> Result<(), BanksClientError> {
         // Instruction #0 will be ed25519 instruction;
         // Instruction #1 will be our contract instruction.
         let instruction_index = 1;
-        // 8 bytes for Anchor header, 4 bytes for Vec length.
-        let message_offset = 12;
         let ed25519_args = dbg!(pyth_lazer_solana_contract::Ed25519SignatureOffsets::new(
             message,
             instruction_index,
             message_offset,
         ));
 
-        let treasury_starting_lamports = self
-            .banks_client
-            .get_account(treasury)
-            .await
-            .unwrap()
-            .unwrap()
-            .lamports;
         let mut transaction_verify = Transaction::new_with_payer(
             &[
                 Instruction::new_with_bytes(
@@ -134,7 +155,6 @@ impl Setup {
                         message_data: message.to_vec(),
                         ed25519_instruction_index: 0,
                         signature_index: 0,
-                        message_offset,
                     }
                     .data(),
                     vec![
@@ -152,17 +172,6 @@ impl Setup {
         self.banks_client
             .process_transaction(transaction_verify)
             .await
-            .unwrap();
-
-        assert_eq!(
-            self.banks_client
-                .get_account(treasury)
-                .await
-                .unwrap()
-                .unwrap()
-                .lamports,
-            treasury_starting_lamports + 1,
-        );
     }
 }
 
@@ -207,6 +216,84 @@ async fn test_basic() {
     setup.verify_message(&message, treasury).await;
 }
 
+#[tokio::test]
+async fn test_rejects_wrong_offset() {
+    let mut setup = Setup::new().await;
+    let treasury = setup.create_treasury().await;
+
+    let mut transaction_init_contract = Transaction::new_with_payer(
+        &[Instruction::new_with_bytes(
+            pyth_lazer_solana_contract::ID,
+            &pyth_lazer_solana_contract::instruction::Initialize {
+                top_authority: setup.payer.pubkey(),
+                treasury,
+            }
+            .data(),
+            vec![
+                AccountMeta::new(setup.payer.pubkey(), true),
+                AccountMeta::new(pyth_lazer_solana_contract::STORAGE_ID, false),
+                AccountMeta::new_readonly(system_program::ID, false),
+            ],
+        )],
+        Some(&setup.payer.pubkey()),
+    );
+    transaction_init_contract.sign(&[&setup.payer], setup.recent_blockhash);
+    setup
+        .banks_client
+        .process_transaction(transaction_init_contract)
+        .await
+        .unwrap();
+
+    let verifying_key =
+        hex::decode("74313a6525edf99936aa1477e94c72bc5cc617b21745f5f03296f3154461f214").unwrap();
+    let verifying_key_2 =
+        hex::decode("AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA").unwrap();
+    let message = hex::decode(
+        [
+            // --- First copy of the message (this data is returned by the Lazer contract)
+
+            // SOLANA_FORMAT_MAGIC_LE
+            "b9011a82",
+            // Signature
+            "e5cddee2c1bd364c8c57e1c98a6a28d194afcad410ff412226c8b2ae931ff59a57147cb47c7307afc2a0a1abec4dd7e835a5b7113cf5aeac13a745c6bed6c600",
+            // Pubkey
+            "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA",
+            // Payload length (could be adjusted)
+            "1c00",
+            // Payload
+            "BBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBB",
+
+            // --- Second copy of the message (this data is used by the ed25519 program)
+
+            // Unused, was SOLANA_FORMAT_MAGIC_LE, could be removed, left it for slightly easier offset adjustments
+            "AABBCCDD",
+            // Signature
+            "e5cddee2c1bd364c8c57e1c98a6a28d194afcad410ff412226c8b2ae931ff59a57147cb47c7307afc2a0a1abec4dd7e835a5b7113cf5aeac13a745c6bed6c600",
+            // Pubkey
+            "74313a6525edf99936aa1477e94c72bc5cc617b21745f5f03296f3154461f214",
+            // Payload length
+            "1c00",
+            // Payload
+            "75d3c7931c9773f30a240600010102000000010000e1f50500000000"
+        ].concat()
+    )
+    .unwrap();
+
+    setup.set_trusted(verifying_key.try_into().unwrap()).await;
+    setup.set_trusted(verifying_key_2.try_into().unwrap()).await;
+    let err = setup
+        .verify_message_with_offset(&message, treasury, 12 + 130)
+        .await
+        .unwrap_err();
+    assert!(matches!(
+        err,
+        BanksClientError::TransactionError(TransactionError::InstructionError(
+            1,
+            InstructionError::InvalidInstructionData
+        ))
+    ));
+}
+
 #[tokio::test]
 async fn test_migrate_from_0_1_0() {
     let mut program_test = program_test();
@@ -277,3 +364,58 @@ async fn test_migrate_from_0_1_0() {
     // because it was present in the original storage PDA data.
     setup.verify_message(&message, treasury).await;
 }
+
+#[tokio::test]
+async fn test_disallows_extra_migrate() {
+    let mut setup = Setup::new().await;
+    let treasury = setup.create_treasury().await;
+
+    let mut transaction_init_contract = Transaction::new_with_payer(
+        &[Instruction::new_with_bytes(
+            pyth_lazer_solana_contract::ID,
+            &pyth_lazer_solana_contract::instruction::Initialize {
+                top_authority: setup.payer.pubkey(),
+                treasury,
+            }
+            .data(),
+            vec![
+                AccountMeta::new(setup.payer.pubkey(), true),
+                AccountMeta::new(pyth_lazer_solana_contract::STORAGE_ID, false),
+                AccountMeta::new_readonly(system_program::ID, false),
+            ],
+        )],
+        Some(&setup.payer.pubkey()),
+    );
+    transaction_init_contract.sign(&[&setup.payer], setup.recent_blockhash);
+    setup
+        .banks_client
+        .process_transaction(transaction_init_contract)
+        .await
+        .unwrap();
+
+    let mut transaction_migrate_contract = Transaction::new_with_payer(
+        &[Instruction::new_with_bytes(
+            pyth_lazer_solana_contract::ID,
+            &pyth_lazer_solana_contract::instruction::MigrateFrom010 { treasury }.data(),
+            vec![
+                AccountMeta::new(setup.payer.pubkey(), true),
+                AccountMeta::new(pyth_lazer_solana_contract::STORAGE_ID, false),
+                AccountMeta::new_readonly(system_program::ID, false),
+            ],
+        )],
+        Some(&setup.payer.pubkey()),
+    );
+    transaction_migrate_contract.sign(&[&setup.payer], setup.recent_blockhash);
+    let err = setup
+        .banks_client
+        .process_transaction(transaction_migrate_contract)
+        .await
+        .unwrap_err();
+    assert!(matches!(
+        err,
+        BanksClientError::TransactionError(TransactionError::InstructionError(
+            0,
+            InstructionError::InvalidAccountData
+        ))
+    ));
+}
