fix: remove num_trusted_signers, remove max signers functionality, add getters for private struct fields, improve tests

tejasbadadare · tejasbadadare · commit 69f2946a4fef · 2025-02-19T11:06:09.000-08:00
diff --git a/lazer/contracts/aptos/sources/pyth_lazer.move b/lazer/contracts/aptos/sources/pyth_lazer.move
@@ -9,63 +9,103 @@ module pyth_lazer::pyth_lazer {
     /// Error codes
     const ENO_PERMISSIONS: u64 = 1;
     const EINVALID_SIGNER: u64 = 2;
-    const ENO_SPACE: u64 = 3;
     const ENO_SUCH_PUBKEY: u64 = 4;
     const EINVALID_SIGNATURE: u64 = 5;
     const EINSUFFICIENT_FEE: u64 = 6;
 
     /// Constants
-    const MAX_NUM_TRUSTED_SIGNERS: u8 = 2;
     const ED25519_PUBLIC_KEY_LENGTH: u64 = 32;
 
     /// Stores information about a trusted signer including their public key and expiration
-    struct TrustedSignerInfo has store, drop {
-        pubkey: vector<u8>,  // Ed25519 public key (32 bytes)
-        expires_at: u64,     // Unix timestamp
+    struct TrustedSignerInfo has store, drop, copy {
+        pubkey: vector<u8>, // Ed25519 public key (32 bytes)
+        expires_at: u64 // Unix timestamp
     }
 
     /// Main storage for the Lazer contract
     struct Storage has key {
         top_authority: address,
         treasury: address,
         single_update_fee: u64,
-        num_trusted_signers: u8,
-        trusted_signers: vector<TrustedSignerInfo>,
+        trusted_signers: vector<TrustedSignerInfo>
     }
 
     /// Events
     struct TrustedSignerUpdateEvent has drop, store {
         pubkey: vector<u8>,
-        expires_at: u64,
+        expires_at: u64
     }
 
-    /// Initialize the Lazer contract with top authority and treasury
+    /// Initialize the Lazer contract with top authority and treasury. One-time operation.
     public entry fun initialize(
         account: &signer,
         top_authority: address,
         treasury: address,
     ) {
+        // Initialize must be called by the contract account
+        assert!(signer::address_of(account) == @pyth_lazer, ENO_PERMISSIONS);
         let storage = Storage {
             top_authority,
             treasury,
             single_update_fee: 1, // Nominal fee
-            num_trusted_signers: 0,
-            trusted_signers: vector::empty(),
+            trusted_signers: vector::empty()
         };
+        // Can only be called once. If storage already exists in @pyth_lazer,
+        // this operation will fail (one-time initialization).
         move_to(account, storage);
     }
 
+    /// Verify a message signature and collect fee
+    public entry fun verify_message(
+        account: &signer,
+        message: vector<u8>,
+        signature: vector<u8>,
+        trusted_signer: vector<u8>
+    ) acquires Storage {
+        let storage = borrow_global<Storage>(@pyth_lazer);
+
+        // Verify fee payment
+        assert!(
+            coin::balance<AptosCoin>(signer::address_of(account))
+                >= storage.single_update_fee,
+            EINSUFFICIENT_FEE
+        );
+        coin::transfer<AptosCoin>(account, storage.treasury, storage.single_update_fee);
+
+        // Verify signer is trusted and not expired
+        let i = 0;
+        let valid = false;
+        while (i < storage.trusted_signers.length()) {
+            let signer_info = vector::borrow(&storage.trusted_signers, (i as u64));
+            if (signer_info.pubkey == trusted_signer
+                && signer_info.expires_at > timestamp::now_seconds()) {
+                valid = true;
+                break
+            };
+            i = i + 1;
+        };
+        assert!(valid, EINVALID_SIGNER);
+
+        // Verify signature
+        let sig = ed25519::new_signature_from_bytes(signature);
+        let pk = ed25519::new_unvalidated_public_key_from_bytes(trusted_signer);
+        assert!(
+            ed25519::signature_verify_strict(&sig, &pk, message),
+            EINVALID_SIGNATURE
+        );
+    }
     /// Upsert a trusted signer's information or remove them
     public entry fun update_trusted_signer(
-        account: &signer,
-        trusted_signer: vector<u8>,
-        expires_at: u64,
+        account: &signer, trusted_signer: vector<u8>, expires_at: u64
     ) acquires Storage {
         let storage = borrow_global_mut<Storage>(@pyth_lazer);
         assert!(signer::address_of(account) == storage.top_authority, ENO_PERMISSIONS);
-        assert!(vector::length(&trusted_signer) == ED25519_PUBLIC_KEY_LENGTH, EINVALID_SIGNER);
+        assert!(
+            vector::length(&trusted_signer) == ED25519_PUBLIC_KEY_LENGTH,
+            EINVALID_SIGNER
+        );
 
-        let num_signers = storage.num_trusted_signers;
+        let num_signers = storage.trusted_signers.length();
         let i = 0;
         let found = false;
 
@@ -82,51 +122,33 @@ module pyth_lazer::pyth_lazer {
             // Remove signer
             assert!(found, ENO_SUCH_PUBKEY);
             vector::remove(&mut storage.trusted_signers, (i as u64));
-            storage.num_trusted_signers = storage.num_trusted_signers - 1;
         } else if (found) {
             // Update existing signer
             let signer_info = vector::borrow_mut(&mut storage.trusted_signers, (i as u64));
             signer_info.expires_at = expires_at;
         } else {
             // Add new signer
-            assert!(storage.num_trusted_signers < MAX_NUM_TRUSTED_SIGNERS, ENO_SPACE);
-            vector::push_back(&mut storage.trusted_signers, TrustedSignerInfo {
-                pubkey: trusted_signer,
-                expires_at,
-            });
-            storage.num_trusted_signers = storage.num_trusted_signers + 1;
+            vector::push_back(
+                &mut storage.trusted_signers,
+                TrustedSignerInfo { pubkey: trusted_signer, expires_at }
+            );
         };
     }
 
-    /// Verify a message signature and collect fee
-    public entry fun verify_message(
-        account: &signer,
-        message: vector<u8>,
-        signature: vector<u8>,
-        public_key: vector<u8>,
-    ) acquires Storage {
+    /// Returns the list of trusted signers
+    public fun get_trusted_signers(): vector<TrustedSignerInfo> acquires Storage {
         let storage = borrow_global<Storage>(@pyth_lazer);
+        storage.trusted_signers
+    }
+    /// Signer pubkey getter
+    public fun get_signer_pubkey(info: &TrustedSignerInfo): vector<u8> {
+        info.pubkey
+    }
 
-        // Verify fee payment
-        assert!(coin::balance<AptosCoin>(signer::address_of(account)) >= storage.single_update_fee, EINSUFFICIENT_FEE);
-        coin::transfer<AptosCoin>(account, storage.treasury, storage.single_update_fee);
+    /// Signer expiry getter
+    public fun get_signer_expires_at(info: &TrustedSignerInfo): u64 {
+        info.expires_at
+    }
 
-        // Verify signer is trusted and not expired
-        let i = 0;
-        let valid = false;
-        while (i < storage.num_trusted_signers) {
-            let signer_info = vector::borrow(&storage.trusted_signers, (i as u64));
-            if (signer_info.pubkey == public_key && signer_info.expires_at > timestamp::now_seconds()) {
-                valid = true;
-                break
-            };
-            i = i + 1;
-        };
-        assert!(valid, EINVALID_SIGNER);
 
-        // Verify signature
-        let sig = ed25519::new_signature_from_bytes(signature);
-        let pk = ed25519::new_unvalidated_public_key_from_bytes(public_key);
-        assert!(ed25519::signature_verify_strict(&sig, &pk, message), EINVALID_SIGNATURE);
-    }
 }
diff --git a/lazer/contracts/aptos/tests/pyth_lazer_tests.move b/lazer/contracts/aptos/tests/pyth_lazer_tests.move
@@ -1,15 +1,17 @@
 #[test_only]
 module pyth_lazer::pyth_lazer_tests {
     use std::signer;
+    use std::vector;
     use aptos_framework::account;
     use aptos_framework::coin;
     use aptos_framework::timestamp;
     use aptos_framework::aptos_coin::AptosCoin;
     use aptos_std::ed25519;
-    use pyth_lazer::pyth_lazer;
+    use pyth_lazer::pyth_lazer::{Self, EINVALID_SIGNER, EINSUFFICIENT_FEE};
 
     // Test accounts
-    const TOP_AUTHORITY: address = @0x3374049c3b46a907ff2fc6b62af51975fb9dc572b7e73eb1b255ed5edcd7cee0;
+    const TOP_AUTHORITY: address =
+        @0x3374049c3b46a907ff2fc6b62af51975fb9dc572b7e73eb1b255ed5edcd7cee0;
     const TREASURY: address = @0x456;
     const USER: address = @0x789;
 
@@ -20,13 +22,14 @@ module pyth_lazer::pyth_lazer_tests {
 
     #[test_only]
     fun setup_aptos_coin(framework: &signer): coin::MintCapability<AptosCoin> {
-        let (burn_cap, freeze_cap, mint_cap) = coin::initialize<AptosCoin>(
-            framework,
-            std::string::utf8(b"Aptos Coin"),
-            std::string::utf8(b"APT"),
-            8,
-            false,
-        );
+        let (burn_cap, freeze_cap, mint_cap) =
+            coin::initialize<AptosCoin>(
+                framework,
+                std::string::utf8(b"Aptos Coin"),
+                std::string::utf8(b"APT"),
+                8,
+                false
+            );
         coin::destroy_burn_cap(burn_cap);
         coin::destroy_freeze_cap(freeze_cap);
         mint_cap
@@ -69,7 +72,7 @@ module pyth_lazer::pyth_lazer_tests {
     }
 
     #[test]
-    fun test_verify_message_success() {
+    fun test_verify_message_succeeds() {
         let (top_authority, _treasury, user) = setup();
 
         // Add a valid signer
@@ -79,47 +82,70 @@ module pyth_lazer::pyth_lazer_tests {
         // Create a valid ed25519 signature
         let signature = ed25519::new_signature_from_bytes(TEST_SIGNATURE);
         let pubkey = ed25519::new_unvalidated_public_key_from_bytes(TEST_PUBKEY);
-        assert!(ed25519::signature_verify_strict(&signature, &pubkey, TEST_MESSAGE), 0);
+        assert!(
+            ed25519::signature_verify_strict(&signature, &pubkey, TEST_MESSAGE),
+            0
+        );
 
         // This should succeed as we have a valid signer and sufficient fee
-        pyth_lazer::verify_message(&user, TEST_MESSAGE, TEST_SIGNATURE, TEST_PUBKEY);
-    }
-
-    #[test]
-    fun test_update_add_signer() {
-        let (top_authority, _treasury, _) = setup();
-
-        // Add signer
-        let expires_at = timestamp::now_seconds() + 1000;
-        pyth_lazer::update_trusted_signer(&top_authority, TEST_PUBKEY, expires_at);
-
-        // Update signer
-        let new_expires_at = timestamp::now_seconds() + 2000;
-        pyth_lazer::update_trusted_signer(&top_authority, TEST_PUBKEY, new_expires_at);
-
-        // Remove signer
-        pyth_lazer::update_trusted_signer(&top_authority, TEST_PUBKEY, 0);
+        pyth_lazer::verify_message(
+            &user,
+            TEST_MESSAGE,
+            TEST_SIGNATURE,
+            TEST_PUBKEY
+        );
     }
 
     #[test]
-    #[expected_failure(abort_code = pyth_lazer::ENO_SPACE)]
-    fun test_max_signers() {
+    fun test_add_update_remove_signers_succeeds() {
         let (top_authority, _treasury, _) = setup();
 
+        // Add multiple signers
         let expires_at = timestamp::now_seconds() + 1000;
         let pubkey1 = x"1111111111111111111111111111111111111111111111111111111111111111";
         let pubkey2 = x"2222222222222222222222222222222222222222222222222222222222222222";
         let pubkey3 = x"3333333333333333333333333333333333333333333333333333333333333333";
 
         pyth_lazer::update_trusted_signer(&top_authority, pubkey1, expires_at);
         pyth_lazer::update_trusted_signer(&top_authority, pubkey2, expires_at);
-        // This should fail as we already have 2 signers
         pyth_lazer::update_trusted_signer(&top_authority, pubkey3, expires_at);
+
+        // Verify signers were added
+        let trusted_signers = pyth_lazer::get_trusted_signers();
+        assert!(vector::length(&trusted_signers) == 3, 0);
+
+        // Verify first signer
+        let signer_info = vector::borrow(&trusted_signers, 0);
+        let signer_pubkey = pyth_lazer::get_signer_pubkey(signer_info);
+        let signer_expires_at = pyth_lazer::get_signer_expires_at(signer_info);
+        assert!(signer_pubkey == pubkey1, 0);
+        assert!(signer_expires_at == expires_at, 0);
+
+        // Update second signer
+        let new_expires_at = timestamp::now_seconds() + 2000;
+        pyth_lazer::update_trusted_signer(&top_authority, pubkey2, new_expires_at);
+
+        // Verify second signer was updated
+        let trusted_signers = pyth_lazer::get_trusted_signers();
+        let signer_info = vector::borrow(&trusted_signers, 1);
+        let signer_pubkey = pyth_lazer::get_signer_pubkey(signer_info);
+        let signer_expires_at = pyth_lazer::get_signer_expires_at(signer_info);
+        assert!(signer_pubkey == pubkey2, 0);
+        assert!(signer_expires_at == new_expires_at, 0);
+
+        // Remove all signers
+        pyth_lazer::update_trusted_signer(&top_authority, pubkey1, 0);
+        pyth_lazer::update_trusted_signer(&top_authority, pubkey2, 0);
+        pyth_lazer::update_trusted_signer(&top_authority, pubkey3, 0);
+
+        // Verify all signers were removed
+        let trusted_signers = pyth_lazer::get_trusted_signers();
+        assert!(vector::length(&trusted_signers) == 0, 0);
     }
 
     #[test]
-    #[expected_failure(abort_code = pyth_lazer::EINVALID_SIGNER)]
-    fun test_expired_signer() {
+    #[expected_failure(abort_code = EINVALID_SIGNER)]
+    fun test_expired_signer_throws_error() {
         let (top_authority, _treasury, user) = setup();
 
         // Add signer that expires in 1000 seconds
@@ -130,19 +156,29 @@ module pyth_lazer::pyth_lazer_tests {
         timestamp::fast_forward_seconds(2000);
 
         // This should fail as the signer is expired
-        pyth_lazer::verify_message(&user, TEST_MESSAGE, TEST_SIGNATURE, TEST_PUBKEY);
+        pyth_lazer::verify_message(
+            &user,
+            TEST_MESSAGE,
+            TEST_SIGNATURE,
+            TEST_PUBKEY
+        );
     }
 
     #[test]
-    #[expected_failure(abort_code = pyth_lazer::EINSUFFICIENT_FEE)]
-    fun test_insufficient_fee() {
+    #[expected_failure(abort_code = EINSUFFICIENT_FEE)]
+    fun test_insufficient_fee_throws_error() {
         let (_top_authority, _treasury, user) = setup();
 
         // Drain user's balance
         let user_balance = coin::balance<AptosCoin>(signer::address_of(&user));
         coin::transfer<AptosCoin>(&user, TREASURY, user_balance);
 
         // This should fail due to insufficient fee
-        pyth_lazer::verify_message(&user, TEST_MESSAGE, TEST_SIGNATURE, TEST_PUBKEY);
+        pyth_lazer::verify_message(
+            &user,
+            TEST_MESSAGE,
+            TEST_SIGNATURE,
+            TEST_PUBKEY
+        );
     }
 }
