|
| 1 | +module pyth_lazer::pyth_lazer { |
| 2 | + use std::vector; |
| 3 | + use std::signer; |
| 4 | + use aptos_framework::timestamp; |
| 5 | + use aptos_framework::coin; |
| 6 | + use aptos_framework::aptos_coin::AptosCoin; |
| 7 | + use aptos_std::ed25519; |
| 8 | + |
| 9 | + /// Error codes |
| 10 | + const ENO_PERMISSIONS: u64 = 1; |
| 11 | + const EINVALID_SIGNER: u64 = 2; |
| 12 | + const ENO_SUCH_PUBKEY: u64 = 4; |
| 13 | + const EINVALID_SIGNATURE: u64 = 5; |
| 14 | + const EINSUFFICIENT_FEE: u64 = 6; |
| 15 | + |
| 16 | + /// Constants |
| 17 | + const ED25519_PUBLIC_KEY_LENGTH: u64 = 32; |
| 18 | + |
| 19 | + /// Admin capability - holder of this resource can perform admin actions, such as key rotations |
| 20 | + struct AdminCapability has key, store {} |
| 21 | + |
| 22 | + /// Stores the admin capability until it's claimed |
| 23 | + struct PendingAdminCapability has key, drop { |
| 24 | + admin: address |
| 25 | + } |
| 26 | + |
| 27 | + /// Stores information about a trusted signer including their public key and expiration |
| 28 | + struct TrustedSignerInfo has store, drop, copy { |
| 29 | + pubkey: vector<u8>, // Ed25519 public key (32 bytes) |
| 30 | + expires_at: u64 // Unix timestamp |
| 31 | + } |
| 32 | + |
| 33 | + /// Main storage for the Lazer contract |
| 34 | + struct Storage has key { |
| 35 | + treasury: address, |
| 36 | + single_update_fee: u64, |
| 37 | + trusted_signers: vector<TrustedSignerInfo> |
| 38 | + } |
| 39 | + |
| 40 | + /// Events |
| 41 | + struct TrustedSignerUpdateEvent has drop, store { |
| 42 | + pubkey: vector<u8>, |
| 43 | + expires_at: u64 |
| 44 | + } |
| 45 | + |
| 46 | + /// Initialize the Lazer contract with top authority and treasury. One-time operation. |
| 47 | + public entry fun initialize( |
| 48 | + account: &signer, admin: address, treasury: address |
| 49 | + ) { |
| 50 | + // Initialize must be called by the contract account |
| 51 | + assert!(signer::address_of(account) == @pyth_lazer, ENO_PERMISSIONS); |
| 52 | + let storage = Storage { |
| 53 | + treasury, |
| 54 | + single_update_fee: 1, // Nominal fee |
| 55 | + trusted_signers: vector::empty() |
| 56 | + }; |
| 57 | + |
| 58 | + // Store the pending admin capability |
| 59 | + move_to(account, PendingAdminCapability { admin }); |
| 60 | + |
| 61 | + // Can only be called once. If storage already exists in @pyth_lazer, |
| 62 | + // this operation will fail (one-time initialization). |
| 63 | + move_to(account, storage); |
| 64 | + } |
| 65 | + |
| 66 | + /// Allows the designated admin to claim their capability |
| 67 | + public entry fun claim_admin_capability(account: &signer) acquires PendingAdminCapability { |
| 68 | + let pending = borrow_global<PendingAdminCapability>(@pyth_lazer); |
| 69 | + assert!(signer::address_of(account) == pending.admin, ENO_PERMISSIONS); |
| 70 | + |
| 71 | + // Create and move the admin capability to the claiming account |
| 72 | + move_to(account, AdminCapability {}); |
| 73 | + |
| 74 | + // Clean up the pending admin capability |
| 75 | + let PendingAdminCapability { admin: _ } = |
| 76 | + move_from<PendingAdminCapability>(@pyth_lazer); |
| 77 | + } |
| 78 | + |
| 79 | + /// Verify a message signature and collect fee. |
| 80 | + /// |
| 81 | + /// This is a convenience wrapper around verify_message(), which allows you to verify an update |
| 82 | + /// using an entry function. If possible, it is recommended to use update_price_feeds() instead, |
| 83 | + /// which avoids the need to pass a signer account. update_price_feeds_with_funder() should only |
| 84 | + /// be used when you need to call an entry function. |
| 85 | + public entry fun verify_message_with_funder( |
| 86 | + account: &signer, |
| 87 | + message: vector<u8>, |
| 88 | + signature: vector<u8>, |
| 89 | + trusted_signer: vector<u8> |
| 90 | + ) acquires Storage { |
| 91 | + let storage = borrow_global<Storage>(@pyth_lazer); |
| 92 | + |
| 93 | + // Verify fee payment |
| 94 | + assert!( |
| 95 | + coin::balance<AptosCoin>(signer::address_of(account)) |
| 96 | + >= storage.single_update_fee, |
| 97 | + EINSUFFICIENT_FEE |
| 98 | + ); |
| 99 | + let fee = coin::withdraw<AptosCoin>(account, storage.single_update_fee); |
| 100 | + verify_message(message, signature, trusted_signer, fee); |
| 101 | + } |
| 102 | + |
| 103 | + /// Verify a message signature with provided fee |
| 104 | + /// The provided `fee` must contain enough coins to pay a single update fee, which |
| 105 | + /// can be queried by calling calling get_update_fee(). |
| 106 | + public fun verify_message( |
| 107 | + message: vector<u8>, |
| 108 | + signature: vector<u8>, |
| 109 | + trusted_signer: vector<u8>, |
| 110 | + fee: coin::Coin<AptosCoin> |
| 111 | + ) acquires Storage { |
| 112 | + let storage = borrow_global<Storage>(@pyth_lazer); |
| 113 | + |
| 114 | + // Verify fee amount |
| 115 | + assert!(coin::value(&fee) >= storage.single_update_fee, EINSUFFICIENT_FEE); |
| 116 | + |
| 117 | + // Transfer fee to treasury |
| 118 | + coin::deposit(storage.treasury, fee); |
| 119 | + |
| 120 | + // Verify signer is trusted and not expired |
| 121 | + let i = 0; |
| 122 | + let valid = false; |
| 123 | + while (i < storage.trusted_signers.length()) { |
| 124 | + let signer_info = vector::borrow(&storage.trusted_signers, (i as u64)); |
| 125 | + if (&signer_info.pubkey == &trusted_signer |
| 126 | + && signer_info.expires_at > timestamp::now_seconds()) { |
| 127 | + valid = true; |
| 128 | + break |
| 129 | + }; |
| 130 | + i = i + 1; |
| 131 | + }; |
| 132 | + assert!(valid, EINVALID_SIGNER); |
| 133 | + |
| 134 | + // Verify signature |
| 135 | + let sig = ed25519::new_signature_from_bytes(signature); |
| 136 | + let pk = ed25519::new_unvalidated_public_key_from_bytes(trusted_signer); |
| 137 | + assert!( |
| 138 | + ed25519::signature_verify_strict(&sig, &pk, message), |
| 139 | + EINVALID_SIGNATURE |
| 140 | + ); |
| 141 | + } |
| 142 | + |
| 143 | + /// Upsert a trusted signer's information or remove them |
| 144 | + public entry fun update_trusted_signer( |
| 145 | + account: &signer, trusted_signer: vector<u8>, expires_at: u64 |
| 146 | + ) acquires Storage { |
| 147 | + // Verify admin capability |
| 148 | + assert!( |
| 149 | + exists<AdminCapability>(signer::address_of(account)), |
| 150 | + ENO_PERMISSIONS |
| 151 | + ); |
| 152 | + |
| 153 | + assert!( |
| 154 | + vector::length(&trusted_signer) == ED25519_PUBLIC_KEY_LENGTH, |
| 155 | + EINVALID_SIGNER |
| 156 | + ); |
| 157 | + |
| 158 | + let storage = borrow_global_mut<Storage>(@pyth_lazer); |
| 159 | + let num_signers = storage.trusted_signers.length(); |
| 160 | + let i = 0; |
| 161 | + let found = false; |
| 162 | + |
| 163 | + while (i < num_signers) { |
| 164 | + let signer_info = vector::borrow(&storage.trusted_signers, (i as u64)); |
| 165 | + if (&signer_info.pubkey == &trusted_signer) { |
| 166 | + found = true; |
| 167 | + break |
| 168 | + }; |
| 169 | + i = i + 1; |
| 170 | + }; |
| 171 | + |
| 172 | + if (expires_at == 0) { |
| 173 | + // Remove signer |
| 174 | + assert!(found, ENO_SUCH_PUBKEY); |
| 175 | + vector::remove(&mut storage.trusted_signers, (i as u64)); |
| 176 | + } else if (found) { |
| 177 | + // Update existing signer |
| 178 | + let signer_info = vector::borrow_mut(&mut storage.trusted_signers, (i as u64)); |
| 179 | + signer_info.expires_at = expires_at; |
| 180 | + } else { |
| 181 | + // Add new signer |
| 182 | + vector::push_back( |
| 183 | + &mut storage.trusted_signers, |
| 184 | + TrustedSignerInfo { pubkey: trusted_signer, expires_at } |
| 185 | + ); |
| 186 | + }; |
| 187 | + } |
| 188 | + |
| 189 | + /// Returns the list of trusted signers |
| 190 | + public fun get_trusted_signers(): vector<TrustedSignerInfo> acquires Storage { |
| 191 | + let storage = borrow_global<Storage>(@pyth_lazer); |
| 192 | + storage.trusted_signers |
| 193 | + } |
| 194 | + |
| 195 | + /// Returns the fee required to verify a message |
| 196 | + public fun get_update_fee(): u64 acquires Storage { |
| 197 | + let storage = borrow_global<Storage>(@pyth_lazer); |
| 198 | + storage.single_update_fee |
| 199 | + } |
| 200 | + |
| 201 | + /// Signer pubkey getter |
| 202 | + public fun get_signer_pubkey(info: &TrustedSignerInfo): vector<u8> { |
| 203 | + info.pubkey |
| 204 | + } |
| 205 | + |
| 206 | + /// Signer expiry getter |
| 207 | + public fun get_signer_expires_at(info: &TrustedSignerInfo): u64 { |
| 208 | + info.expires_at |
| 209 | + } |
| 210 | +} |
0 commit comments