@layerzerolabs/protocol-stellar-v2 0.2.60 → 0.2.61

package/contracts/workers/dvn/src/tests/auth.rs

@@ -1,27 +1,15 @@
 use crate::{
     errors::DvnError,
     tests::setup::{TestSetup, VID},
-    LzDVNClient, Sender, TransactionAuthData,
+    Call, LzDVNClient, Sender, TransactionAuthData,
 };
 use ed25519_dalek::{Signer, SigningKey};
 use rand::thread_rng;
-use soroban_sdk::{auth::Context, vec, xdr::ToXdr, Bytes, BytesN, Env, IntoVal, Vec};
-use utils::buffer_writer::BufferWriter;
-
-const AUTH_NONCE: i64 = 0;
-const AUTH_EXPIRY_LEDGER: u32 = 10000;
-const ENVELOPE_TYPE_SOROBAN_AUTHORIZATION: u32 = 9;
-
-/// Build the full preimage bytes from auth fields.
-fn build_preimage(env: &Env, auth_nonce: i64, auth_expiry_ledger: u32, invocation: &Bytes) -> Bytes {
-    BufferWriter::new(env)
-        .write_u32(ENVELOPE_TYPE_SOROBAN_AUTHORIZATION)
-        .write_bytes_n(&env.ledger().network_id())
-        .write_i64(auth_nonce)
-        .write_u32(auth_expiry_ledger)
-        .write_bytes(invocation)
-        .to_bytes()
-}
+use soroban_sdk::{
+    auth::{ContractContext, Context},
+    testutils::Address as _,
+    vec, Address, BytesN, Env, IntoVal, Symbol, Val, Vec,
+};
 
 struct Ed25519KeyPair {
     signing_key: SigningKey,
@@ -46,18 +34,97 @@ impl Ed25519KeyPair {
     }
 }
 
-/// Compute signature_payload = SHA256(preimage) for the given preimage bytes.
-fn compute_signature_payload(env: &Env, preimage: &Bytes) -> BytesN<32> {
-    env.crypto().sha256(preimage).into()
+/// Creates a single self-call auth context and matching Call for testing.
+fn make_self_call_context(
+    env: &Env,
+    contract_id: &soroban_sdk::Address,
+) -> (Vec<Context>, Vec<Call>) {
+    let fn_name = Symbol::new(env, "set_paused");
+    let args: Vec<Val> = vec![env, true.into_val(env)];
+
+    let context = Context::Contract(ContractContext {
+        contract: contract_id.clone(),
+        fn_name: fn_name.clone(),
+        args: args.clone(),
+    });
+    let auth_contexts = vec![env, context];
+
+    let calls = vec![env, Call { to: contract_id.clone(), func: fn_name, args }];
+
+    (auth_contexts, calls)
 }
 
-/// Compute call hash for multisig: keccak256(vid || expiration || invocation).
-fn compute_call_hash(env: &Env, vid: u32, expiration: u64, invocation: &Bytes) -> BytesN<32> {
-    let mut writer = BufferWriter::new(env);
-    let data = writer.write_u32(vid).write_u64(expiration).write_bytes(invocation).to_bytes();
-    env.crypto().keccak256(&data).into()
+/// Creates upgrade auth contexts (3 entries: upgrader call, upgrade, migrate).
+fn make_upgrade_contexts(
+    env: &Env,
+    contract_id: &soroban_sdk::Address,
+    upgrader_id: &soroban_sdk::Address,
+) -> (Vec<Context>, Vec<Call>) {
+    let wasm_hash_val: Val = BytesN::from_array(env, &[0xABu8; 32]).into_val(env);
+    let migration_data_val: Val = soroban_sdk::Bytes::new(env).into_val(env);
+
+    // [0]: Upgrader call
+    let upgrader_ctx = Context::Contract(ContractContext {
+        contract: upgrader_id.clone(),
+        fn_name: Symbol::new(env, "upgrade_and_migrate"),
+        args: vec![
+            env,
+            contract_id.clone().into_val(env),
+            wasm_hash_val.clone(),
+            migration_data_val.clone(),
+        ],
+    });
+    // [1]: upgrade self-call
+    let upgrade_ctx = Context::Contract(ContractContext {
+        contract: contract_id.clone(),
+        fn_name: Symbol::new(env, "upgrade"),
+        args: vec![env, wasm_hash_val],
+    });
+    // [2]: migrate self-call
+    let migrate_ctx = Context::Contract(ContractContext {
+        contract: contract_id.clone(),
+        fn_name: Symbol::new(env, "migrate"),
+        args: vec![env, migration_data_val],
+    });
+
+    let auth_contexts = vec![env, upgrader_ctx.clone(), upgrade_ctx.clone(), migrate_ctx.clone()];
+
+    // Build matching Call vec
+    let calls = vec![
+        env,
+        Call {
+            to: upgrader_id.clone(),
+            func: Symbol::new(env, "upgrade_and_migrate"),
+            args: match &upgrader_ctx {
+                Context::Contract(c) => c.args.clone(),
+                _ => unreachable!(),
+            },
+        },
+        Call {
+            to: contract_id.clone(),
+            func: Symbol::new(env, "upgrade"),
+            args: match &upgrade_ctx {
+                Context::Contract(c) => c.args.clone(),
+                _ => unreachable!(),
+            },
+        },
+        Call {
+            to: contract_id.clone(),
+            func: Symbol::new(env, "migrate"),
+            args: match &migrate_ctx {
+                Context::Contract(c) => c.args.clone(),
+                _ => unreachable!(),
+            },
+        },
+    ];
+
+    (auth_contexts, calls)
 }
 
+// ============================================================================
+// Single Self-Call Tests
+// ============================================================================
+
 #[test]
 fn test_check_auth_success() {
     extern crate std;
@@ -67,15 +134,14 @@ fn test_check_auth_success() {
     let setup = TestSetup::with_admin_bytes(1, std::vec![admin_bytes]);
     let env = setup.env.clone();
     let expiration = env.ledger().timestamp() + 1000;
-    let auth_contexts: Vec<Context> = Vec::new(&env);
+    let (auth_contexts, calls) = make_self_call_context(&env, &setup.contract_id);
 
-    let invocation = auth_contexts.clone().to_xdr(&env);
-    let preimage = build_preimage(&env, AUTH_NONCE, AUTH_EXPIRY_LEDGER, &invocation);
-    let payload = compute_signature_payload(&env, &preimage);
+    let payload = BytesN::from_array(&env, &[0u8; 32]);
     let public_key = admin_kp.public_key(&env);
     let signature = admin_kp.sign(&env, &payload.to_array());
 
-    let hash = compute_call_hash(&env, VID, expiration, &invocation);
+    let dvn_client = LzDVNClient::new(&env, &setup.contract_id);
+    let hash = dvn_client.hash_call_data(&VID, &expiration, &calls);
     let sig = setup.key_pairs[0].sign_bytes(&env, &hash);
 
     let tx_auth = TransactionAuthData {
@@ -83,7 +149,6 @@ fn test_check_auth_success() {
         expiration,
         signatures: vec![&env, sig],
         sender: Sender::Admin(public_key, signature),
-        preimage,
     };
 
     let res = env.try_invoke_contract_check_auth::<DvnError>(
@@ -104,15 +169,14 @@ fn test_check_auth_not_admin() {
     let setup = TestSetup::new(1);
     let env = setup.env.clone();
     let expiration = env.ledger().timestamp() + 1000;
-    let auth_contexts: Vec<Context> = Vec::new(&env);
+    let (auth_contexts, calls) = make_self_call_context(&env, &setup.contract_id);
 
-    let invocation = auth_contexts.clone().to_xdr(&env);
-    let preimage = build_preimage(&env, AUTH_NONCE, AUTH_EXPIRY_LEDGER, &invocation);
-    let payload = compute_signature_payload(&env, &preimage);
+    let payload = BytesN::from_array(&env, &[0u8; 32]);
     let public_key = non_admin_kp.public_key(&env);
     let signature = non_admin_kp.sign(&env, &payload.to_array());
 
-    let hash = compute_call_hash(&env, VID, expiration, &invocation);
+    let dvn_client = LzDVNClient::new(&env, &setup.contract_id);
+    let hash = dvn_client.hash_call_data(&VID, &expiration, &calls);
     let sig = setup.key_pairs[0].sign_bytes(&env, &hash);
 
     let tx_auth = TransactionAuthData {
@@ -120,7 +184,6 @@ fn test_check_auth_not_admin() {
         expiration,
         signatures: vec![&env, sig],
         sender: Sender::Admin(public_key, signature),
-        preimage,
     };
 
     let res = env.try_invoke_contract_check_auth::<DvnError>(
@@ -142,15 +205,14 @@ fn test_check_auth_wrong_signer_fails() {
     let setup = TestSetup::with_admin_bytes(1, std::vec![admin_bytes]);
     let env = setup.env.clone();
     let expiration = env.ledger().timestamp() + 1000;
-    let auth_contexts: Vec<Context> = Vec::new(&env);
+    let (auth_contexts, calls) = make_self_call_context(&env, &setup.contract_id);
 
-    let invocation = auth_contexts.clone().to_xdr(&env);
-    let preimage = build_preimage(&env, AUTH_NONCE, AUTH_EXPIRY_LEDGER, &invocation);
-    let payload = compute_signature_payload(&env, &preimage);
+    let payload = BytesN::from_array(&env, &[0u8; 32]);
     let public_key = admin_kp.public_key(&env);
     let signature = admin_kp.sign(&env, &payload.to_array());
 
-    let hash = compute_call_hash(&env, VID, expiration, &invocation);
+    let dvn_client = LzDVNClient::new(&env, &setup.contract_id);
+    let hash = dvn_client.hash_call_data(&VID, &expiration, &calls);
     let wrong_sig = crate::tests::key_pair::KeyPair::generate().sign_bytes(&env, &hash);
 
     let tx_auth = TransactionAuthData {
@@ -158,9 +220,9 @@ fn test_check_auth_wrong_signer_fails() {
         expiration,
         signatures: vec![&env, wrong_sig],
         sender: Sender::Admin(public_key, signature),
-        preimage,
     };
 
+    // verify_signatures panics with MultiSigError::SignerNotFound when signer is not found
     let res = env.try_invoke_contract_check_auth::<utils::errors::MultiSigError>(
         &setup.contract_id,
         &payload,
@@ -181,15 +243,14 @@ fn test_check_auth_invalid_vid_fails() {
     let env = setup.env.clone();
     let expiration = env.ledger().timestamp() + 1000;
     let wrong_vid = VID + 1;
-    let auth_contexts: Vec<Context> = Vec::new(&env);
+    let (auth_contexts, calls) = make_self_call_context(&env, &setup.contract_id);
 
-    let invocation = auth_contexts.clone().to_xdr(&env);
-    let preimage = build_preimage(&env, AUTH_NONCE, AUTH_EXPIRY_LEDGER, &invocation);
-    let payload = compute_signature_payload(&env, &preimage);
+    let payload = BytesN::from_array(&env, &[0u8; 32]);
     let public_key = admin_kp.public_key(&env);
     let signature = admin_kp.sign(&env, &payload.to_array());
 
-    let hash = compute_call_hash(&env, wrong_vid, expiration, &invocation);
+    let dvn_client = LzDVNClient::new(&env, &setup.contract_id);
+    let hash = dvn_client.hash_call_data(&wrong_vid, &expiration, &calls);
     let sig = setup.key_pairs[0].sign_bytes(&env, &hash);
 
     let tx_auth = TransactionAuthData {
@@ -197,7 +258,6 @@ fn test_check_auth_invalid_vid_fails() {
         expiration,
         signatures: vec![&env, sig],
         sender: Sender::Admin(public_key, signature),
-        preimage,
     };
 
     let res = env.try_invoke_contract_check_auth::<DvnError>(
@@ -219,15 +279,14 @@ fn test_check_auth_expired_fails() {
     let setup = TestSetup::with_admin_bytes(1, std::vec![admin_bytes]);
     let env = setup.env.clone();
     let expiration = env.ledger().timestamp();
-    let auth_contexts: Vec<Context> = Vec::new(&env);
+    let (auth_contexts, calls) = make_self_call_context(&env, &setup.contract_id);
 
-    let invocation = auth_contexts.clone().to_xdr(&env);
-    let preimage = build_preimage(&env, AUTH_NONCE, AUTH_EXPIRY_LEDGER, &invocation);
-    let payload = compute_signature_payload(&env, &preimage);
+    let payload = BytesN::from_array(&env, &[0u8; 32]);
     let public_key = admin_kp.public_key(&env);
     let signature = admin_kp.sign(&env, &payload.to_array());
 
-    let hash = compute_call_hash(&env, VID, expiration, &invocation);
+    let dvn_client = LzDVNClient::new(&env, &setup.contract_id);
+    let hash = dvn_client.hash_call_data(&VID, &expiration, &calls);
     let sig = setup.key_pairs[0].sign_bytes(&env, &hash);
 
     let tx_auth = TransactionAuthData {
@@ -235,7 +294,6 @@ fn test_check_auth_expired_fails() {
         expiration,
         signatures: vec![&env, sig],
         sender: Sender::Admin(public_key, signature),
-        preimage,
     };
 
     let res = env.try_invoke_contract_check_auth::<DvnError>(
@@ -257,15 +315,14 @@ fn test_check_auth_hash_already_used_fails() {
     let setup = TestSetup::with_admin_bytes(1, std::vec![admin_bytes]);
     let env = setup.env.clone();
     let expiration = env.ledger().timestamp() + 1000;
-    let auth_contexts: Vec<Context> = Vec::new(&env);
+    let (auth_contexts, calls) = make_self_call_context(&env, &setup.contract_id);
 
-    let invocation = auth_contexts.clone().to_xdr(&env);
-    let preimage = build_preimage(&env, AUTH_NONCE, AUTH_EXPIRY_LEDGER, &invocation);
-    let payload = compute_signature_payload(&env, &preimage);
+    let payload = BytesN::from_array(&env, &[0u8; 32]);
     let public_key = admin_kp.public_key(&env);
     let signature = admin_kp.sign(&env, &payload.to_array());
 
-    let hash = compute_call_hash(&env, VID, expiration, &invocation);
+    let dvn_client = LzDVNClient::new(&env, &setup.contract_id);
+    let hash = dvn_client.hash_call_data(&VID, &expiration, &calls);
     let sig = setup.key_pairs[0].sign_bytes(&env, &hash);
 
     let tx_auth = TransactionAuthData {
@@ -273,7 +330,6 @@ fn test_check_auth_hash_already_used_fails() {
         expiration,
         signatures: vec![&env, sig.clone()],
         sender: Sender::Admin(public_key.clone(), signature.clone()),
-        preimage: preimage.clone(),
     };
 
     let res = env.try_invoke_contract_check_auth::<DvnError>(
@@ -289,7 +345,6 @@ fn test_check_auth_hash_already_used_fails() {
         expiration,
         signatures: vec![&env, sig],
         sender: Sender::Admin(public_key, signature),
-        preimage,
     };
 
     let res2 = env.try_invoke_contract_check_auth::<DvnError>(
@@ -305,21 +360,26 @@ fn test_check_auth_hash_already_used_fails() {
 #[test]
 fn test_check_auth_sender_none_fails_when_admin_required() {
     extern crate std;
+    use crate::Sender;
 
     let setup = TestSetup::new(1);
     let env = setup.env.clone();
     let expiration = env.ledger().timestamp() + 1000;
-    let auth_contexts: Vec<Context> = Vec::new(&env);
+    let (auth_contexts, calls) = make_self_call_context(&env, &setup.contract_id);
 
-    let invocation = auth_contexts.clone().to_xdr(&env);
-    let preimage = build_preimage(&env, AUTH_NONCE, AUTH_EXPIRY_LEDGER, &invocation);
-    let payload = compute_signature_payload(&env, &preimage);
+    let payload = BytesN::from_array(&env, &[0u8; 32]);
 
-    let hash = compute_call_hash(&env, VID, expiration, &invocation);
+    let dvn_client = LzDVNClient::new(&env, &setup.contract_id);
+    let hash = dvn_client.hash_call_data(&VID, &expiration, &calls);
     let sig = setup.key_pairs[0].sign_bytes(&env, &hash);
 
-    let tx_auth =
-        TransactionAuthData { vid: VID, expiration, signatures: vec![&env, sig], sender: Sender::None, preimage };
+    // Use Sender::None which should fail when admin is required
+    let tx_auth = TransactionAuthData {
+        vid: VID,
+        expiration,
+        signatures: vec![&env, sig],
+        sender: Sender::None,
+    };
 
     let res = env.try_invoke_contract_check_auth::<DvnError>(
         &setup.contract_id,
@@ -334,29 +394,42 @@ fn test_check_auth_sender_none_fails_when_admin_required() {
 #[test]
 fn test_check_auth_set_admin_bypasses_admin_verification() {
     extern crate std;
-    use soroban_sdk::{auth::ContractContext, Symbol};
+    use crate::Sender;
 
     let setup = TestSetup::new(1);
     let env = setup.env.clone();
     let expiration = env.ledger().timestamp() + 1000;
 
+    // Create a context for set_admin call on the DVN contract
     let set_admin_context = Context::Contract(ContractContext {
         contract: setup.contract_id.clone(),
         fn_name: Symbol::new(&env, "set_admin"),
         args: Vec::new(&env),
     });
-    let auth_contexts: Vec<Context> = vec![&env, set_admin_context.clone()];
+    let auth_contexts: Vec<Context> = vec![&env, set_admin_context];
+
+    let payload = BytesN::from_array(&env, &[0u8; 32]);
 
-    let invocation = auth_contexts.clone().to_xdr(&env);
-    let preimage = build_preimage(&env, AUTH_NONCE, AUTH_EXPIRY_LEDGER, &invocation);
-    let payload = compute_signature_payload(&env, &preimage);
+    let calls: Vec<Call> = vec![
+        &env,
+        Call {
+            to: setup.contract_id.clone(),
+            func: Symbol::new(&env, "set_admin"),
+            args: Vec::new(&env),
+        },
+    ];
 
     let dvn_client = LzDVNClient::new(&env, &setup.contract_id);
-    let hash = dvn_client.hash_call_data(&VID, &expiration, &invocation);
+    let hash = dvn_client.hash_call_data(&VID, &expiration, &calls);
     let sig = setup.key_pairs[0].sign_bytes(&env, &hash);
 
-    let tx_auth =
-        TransactionAuthData { vid: VID, expiration, signatures: vec![&env, sig], sender: Sender::None, preimage };
+    // Use Sender::None - should succeed for set_admin since it bypasses admin verification
+    let tx_auth = TransactionAuthData {
+        vid: VID,
+        expiration,
+        signatures: vec![&env, sig],
+        sender: Sender::None,
+    };
 
     let res = env.try_invoke_contract_check_auth::<DvnError>(
         &setup.contract_id,
@@ -365,11 +438,60 @@ fn test_check_auth_set_admin_bypasses_admin_verification() {
         &auth_contexts,
     );
 
+    // Should succeed because set_admin bypasses admin verification
     assert!(res.is_ok(), "Expected success for set_admin call, got {:?}", res);
 }
 
 #[test]
-fn test_check_auth_invalid_signature_payload_fails() {
+fn test_check_auth_non_contract_context_fails() {
+    extern crate std;
+    use crate::Sender;
+    use soroban_sdk::auth::{ContractExecutable, CreateContractHostFnContext};
+
+    let admin_kp = Ed25519KeyPair::generate();
+    let admin_bytes = admin_kp.public_key_bytes();
+
+    let setup = TestSetup::with_admin_bytes(1, std::vec![admin_bytes]);
+    let env = setup.env.clone();
+    let expiration = env.ledger().timestamp() + 1000;
+
+    // Create a non-Contract context (CreateContractHostFn)
+    let non_contract_context = Context::CreateContractHostFn(CreateContractHostFnContext {
+        executable: ContractExecutable::Wasm(BytesN::from_array(&env, &[0; 32])),
+        salt: BytesN::from_array(&env, &[0; 32]),
+    });
+    let auth_contexts: Vec<Context> = vec![&env, non_contract_context];
+
+    let payload = BytesN::from_array(&env, &[0u8; 32]);
+    let public_key = admin_kp.public_key(&env);
+    let signature = admin_kp.sign(&env, &payload.to_array());
+
+    let dvn_client = LzDVNClient::new(&env, &setup.contract_id);
+    let dummy_calls: Vec<Call> =
+        vec![&env, Call { to: setup.contract_id.clone(), func: Symbol::new(&env, "noop"), args: Vec::new(&env) }];
+    let hash = dvn_client.hash_call_data(&VID, &expiration, &dummy_calls);
+    let sig = setup.key_pairs[0].sign_bytes(&env, &hash);
+
+    let tx_auth = TransactionAuthData {
+        vid: VID,
+        expiration,
+        signatures: vec![&env, sig],
+        sender: Sender::Admin(public_key, signature),
+    };
+
+    let res = env.try_invoke_contract_check_auth::<DvnError>(
+        &setup.contract_id,
+        &payload,
+        tx_auth.into_val(&env),
+        &auth_contexts,
+    );
+
+    // Should fail with NonContractInvoke error
+    assert_eq!(res, Err(Ok(DvnError::NonContractInvoke)));
+}
+
+#[test]
+fn test_check_auth_rejects_empty_contexts() {
     extern crate std;
     let admin_kp = Ed25519KeyPair::generate();
     let admin_bytes = admin_kp.public_key_bytes();
@@ -379,27 +501,68 @@ fn test_check_auth_invalid_signature_payload_fails() {
     let expiration = env.ledger().timestamp() + 1000;
     let auth_contexts: Vec<Context> = Vec::new(&env);
 
-    let invocation = auth_contexts.clone().to_xdr(&env);
-    // Compute the correct payload so admin signature is valid
-    let preimage = build_preimage(&env, AUTH_NONCE, AUTH_EXPIRY_LEDGER, &invocation);
-    let payload = compute_signature_payload(&env, &preimage);
+    let payload = BytesN::from_array(&env, &[0u8; 32]);
     let public_key = admin_kp.public_key(&env);
     let signature = admin_kp.sign(&env, &payload.to_array());
 
-    let hash = compute_call_hash(&env, VID, expiration, &invocation);
+    let dvn_client = LzDVNClient::new(&env, &setup.contract_id);
+    let dummy_calls: Vec<Call> =
+        vec![&env, Call { to: setup.contract_id.clone(), func: Symbol::new(&env, "noop"), args: Vec::new(&env) }];
+    let hash = dvn_client.hash_call_data(&VID, &expiration, &dummy_calls);
     let sig = setup.key_pairs[0].sign_bytes(&env, &hash);
 
-    // Tamper with preimage so it no longer matches the signature_payload
-    let mut tampered_invocation = Bytes::new(&env);
-    tampered_invocation.extend_from_array(&[0u8; 32]);
-    let tampered_preimage = build_preimage(&env, AUTH_NONCE, AUTH_EXPIRY_LEDGER, &tampered_invocation);
+    let tx_auth = TransactionAuthData {
+        vid: VID,
+        expiration,
+        signatures: vec![&env, sig],
+        sender: Sender::Admin(public_key, signature),
+    };
+
+    let res = env.try_invoke_contract_check_auth::<DvnError>(
+        &setup.contract_id,
+        &payload,
+        tx_auth.into_val(&env),
+        &auth_contexts,
+    );
+
+    assert_eq!(res, Err(Ok(DvnError::InvalidAuthContext)));
+}
+
+#[test]
+fn test_check_auth_rejects_external_target() {
+    extern crate std;
+    let admin_kp = Ed25519KeyPair::generate();
+    let admin_bytes = admin_kp.public_key_bytes();
+
+    let setup = TestSetup::with_admin_bytes(1, std::vec![admin_bytes]);
+    let env = setup.env.clone();
+    let expiration = env.ledger().timestamp() + 1000;
+
+    // Create a context targeting a different contract
+    let other_setup = TestSetup::new(1);
+    let other_contract = other_setup.contract_id;
+    let external_context = Context::Contract(ContractContext {
+        contract: other_contract,
+        fn_name: Symbol::new(&env, "some_fn"),
+        args: Vec::new(&env),
+    });
+    let auth_contexts: Vec<Context> = vec![&env, external_context];
+
+    let payload = BytesN::from_array(&env, &[0u8; 32]);
+    let public_key = admin_kp.public_key(&env);
+    let signature = admin_kp.sign(&env, &payload.to_array());
+
+    let dvn_client = LzDVNClient::new(&env, &setup.contract_id);
+    let dummy_calls: Vec<Call> =
+        vec![&env, Call { to: setup.contract_id.clone(), func: Symbol::new(&env, "noop"), args: Vec::new(&env) }];
+    let hash = dvn_client.hash_call_data(&VID, &expiration, &dummy_calls);
+    let sig = setup.key_pairs[0].sign_bytes(&env, &hash);
 
     let tx_auth = TransactionAuthData {
         vid: VID,
         expiration,
         signatures: vec![&env, sig],
         sender: Sender::Admin(public_key, signature),
-        preimage: tampered_preimage,
     };
 
     let res = env.try_invoke_contract_check_auth::<DvnError>(
@@ -409,5 +572,291 @@ fn test_check_auth_invalid_signature_payload_fails() {
         &auth_contexts,
     );
 
-    assert_eq!(res, Err(Ok(DvnError::InvalidPreimage)));
+    assert_eq!(res, Err(Ok(DvnError::InvalidAuthContext)));
+}
+
+// ============================================================================
+// set_upgrader Tests
+// ============================================================================
+
+#[test]
+fn test_set_upgrader_requires_admin() {
+    extern crate std;
+    use crate::Sender;
+
+    let setup = TestSetup::new(1);
+    let env = setup.env.clone();
+    let expiration = env.ledger().timestamp() + 1000;
+
+    let set_upgrader_context = Context::Contract(ContractContext {
+        contract: setup.contract_id.clone(),
+        fn_name: Symbol::new(&env, "set_upgrader"),
+        args: Vec::new(&env),
+    });
+    let auth_contexts: Vec<Context> = vec![&env, set_upgrader_context];
+
+    let payload = BytesN::from_array(&env, &[0u8; 32]);
+
+    let calls: Vec<Call> = vec![
+        &env,
+        Call {
+            to: setup.contract_id.clone(),
+            func: Symbol::new(&env, "set_upgrader"),
+            args: Vec::new(&env),
+        },
+    ];
+
+    let dvn_client = LzDVNClient::new(&env, &setup.contract_id);
+    let hash = dvn_client.hash_call_data(&VID, &expiration, &calls);
+    let sig = setup.key_pairs[0].sign_bytes(&env, &hash);
+
+    let tx_auth = TransactionAuthData {
+        vid: VID,
+        expiration,
+        signatures: vec![&env, sig],
+        sender: Sender::None,
+    };
+
+    let res = env.try_invoke_contract_check_auth::<DvnError>(
+        &setup.contract_id,
+        &payload,
+        tx_auth.into_val(&env),
+        &auth_contexts,
+    );
+
+    assert_eq!(res, Err(Ok(DvnError::OnlyAdmin)));
+}
+
+// ============================================================================
+// Upgrade Path Tests
+// ============================================================================
+
+#[test]
+fn test_check_auth_upgrade_success() {
+    extern crate std;
+
+    let admin_kp = Ed25519KeyPair::generate();
+    let admin_bytes = admin_kp.public_key_bytes();
+
+    let setup = TestSetup::with_admin_bytes(1, std::vec![admin_bytes]);
+    let env = setup.env.clone();
+    let expiration = env.ledger().timestamp() + 1000;
+
+    let upgrader_id = Address::generate(&env);
+    env.as_contract(&setup.contract_id, || {
+        crate::storage::DvnStorage::set_upgrader(&env, &upgrader_id);
+    });
+
+    let (auth_contexts, calls) = make_upgrade_contexts(&env, &setup.contract_id, &upgrader_id);
+
+    let payload = BytesN::from_array(&env, &[0u8; 32]);
+    let public_key = admin_kp.public_key(&env);
+    let signature = admin_kp.sign(&env, &payload.to_array());
+
+    let dvn_client = LzDVNClient::new(&env, &setup.contract_id);
+    let hash = dvn_client.hash_call_data(&VID, &expiration, &calls);
+    let sig = setup.key_pairs[0].sign_bytes(&env, &hash);
+
+    let tx_auth = TransactionAuthData {
+        vid: VID,
+        expiration,
+        signatures: vec![&env, sig],
+        sender: Sender::Admin(public_key, signature),
+    };
+
+    let res = env.try_invoke_contract_check_auth::<DvnError>(
+        &setup.contract_id,
+        &payload,
+        tx_auth.into_val(&env),
+        &auth_contexts,
+    );
+
+    assert!(res.is_ok(), "Expected success for upgrade, got {:?}", res);
+}
+
+#[test]
+fn test_check_auth_upgrade_no_upgrader_set() {
+    extern crate std;
+
+    let setup = TestSetup::new(1);
+    let env = setup.env.clone();
+    let expiration = env.ledger().timestamp() + 1000;
+
+    // Don't register any upgrader
+    let fake_upgrader = TestSetup::new(1).contract_id;
+    let (auth_contexts, calls) = make_upgrade_contexts(&env, &setup.contract_id, &fake_upgrader);
+
+    let payload = BytesN::from_array(&env, &[0u8; 32]);
+
+    let dvn_client = LzDVNClient::new(&env, &setup.contract_id);
+    let hash = dvn_client.hash_call_data(&VID, &expiration, &calls);
+    let sig = setup.key_pairs[0].sign_bytes(&env, &hash);
+
+    let tx_auth = TransactionAuthData {
+        vid: VID,
+        expiration,
+        signatures: vec![&env, sig],
+        sender: Sender::None,
+    };
+
+    let res = env.try_invoke_contract_check_auth::<DvnError>(
+        &setup.contract_id,
+        &payload,
+        tx_auth.into_val(&env),
+        &auth_contexts,
+    );
+
+    assert_eq!(res, Err(Ok(DvnError::UpgraderNotSet)));
+}
+
+#[test]
+fn test_check_auth_upgrade_wrong_upgrader() {
+    extern crate std;
+
+    let setup = TestSetup::new(1);
+    let env = setup.env.clone();
+    let expiration = env.ledger().timestamp() + 1000;
+
+    // Register one upgrader
+    let registered_upgrader = Address::generate(&env);
+    env.as_contract(&setup.contract_id, || {
+        crate::storage::DvnStorage::set_upgrader(&env, &registered_upgrader);
+    });
+
+    // But use a different upgrader in the auth contexts
+    let wrong_upgrader = Address::generate(&env);
+    let (auth_contexts, calls) = make_upgrade_contexts(&env, &setup.contract_id, &wrong_upgrader);
+
+    let payload = BytesN::from_array(&env, &[0u8; 32]);
+
+    let dvn_client = LzDVNClient::new(&env, &setup.contract_id);
+    let hash = dvn_client.hash_call_data(&VID, &expiration, &calls);
+    let sig = setup.key_pairs[0].sign_bytes(&env, &hash);
+
+    let tx_auth = TransactionAuthData {
+        vid: VID,
+        expiration,
+        signatures: vec![&env, sig],
+        sender: Sender::None,
+    };
+
+    let res = env.try_invoke_contract_check_auth::<DvnError>(
+        &setup.contract_id,
+        &payload,
+        tx_auth.into_val(&env),
+        &auth_contexts,
+    );
+
+    assert_eq!(res, Err(Ok(DvnError::InvalidUpgradeContext)));
+}
+
+#[test]
+fn test_check_auth_upgrade_missing_migrate() {
+    extern crate std;
+
+    let setup = TestSetup::new(1);
+    let env = setup.env.clone();
+    let expiration = env.ledger().timestamp() + 1000;
+
+    let upgrader_id = TestSetup::new(1).contract_id;
+    env.as_contract(&setup.contract_id, || {
+        crate::storage::DvnStorage::set_upgrader(&env, &upgrader_id);
+    });
+
+    // Only 2 contexts: upgrader + upgrade (missing migrate) → InvalidAuthContext (len != 1 and != 3)
+    let upgrader_ctx = Context::Contract(ContractContext {
+        contract: upgrader_id.clone(),
+        fn_name: Symbol::new(&env, "upgrade_and_migrate"),
+        args: Vec::new(&env),
+    });
+    let upgrade_ctx = Context::Contract(ContractContext {
+        contract: setup.contract_id.clone(),
+        fn_name: Symbol::new(&env, "upgrade"),
+        args: Vec::new(&env),
+    });
+    let auth_contexts: Vec<Context> = vec![&env, upgrader_ctx, upgrade_ctx];
+
+    let payload = BytesN::from_array(&env, &[0u8; 32]);
+    let dummy_calls: Vec<Call> =
+        vec![&env, Call { to: setup.contract_id.clone(), func: Symbol::new(&env, "noop"), args: Vec::new(&env) }];
+
+    let dvn_client = LzDVNClient::new(&env, &setup.contract_id);
+    let hash = dvn_client.hash_call_data(&VID, &expiration, &dummy_calls);
+    let sig = setup.key_pairs[0].sign_bytes(&env, &hash);
+
+    let tx_auth = TransactionAuthData {
+        vid: VID,
+        expiration,
+        signatures: vec![&env, sig],
+        sender: Sender::None,
+    };
+
+    let res = env.try_invoke_contract_check_auth::<DvnError>(
+        &setup.contract_id,
+        &payload,
+        tx_auth.into_val(&env),
+        &auth_contexts,
+    );
+
+    // 2 contexts is neither 1 (single-call) nor 3 (upgrade), so InvalidAuthContext
+    assert_eq!(res, Err(Ok(DvnError::InvalidAuthContext)));
+}
+
+#[test]
+fn test_check_auth_upgrade_replay_protection() {
+    extern crate std;
+
+    let admin_kp = Ed25519KeyPair::generate();
+    let admin_bytes = admin_kp.public_key_bytes();
+
+    let setup = TestSetup::with_admin_bytes(1, std::vec![admin_bytes]);
+    let env = setup.env.clone();
+    let expiration = env.ledger().timestamp() + 1000;
+
+    let upgrader_id = Address::generate(&env);
+    env.as_contract(&setup.contract_id, || {
+        crate::storage::DvnStorage::set_upgrader(&env, &upgrader_id);
+    });
+
+    let (auth_contexts, calls) = make_upgrade_contexts(&env, &setup.contract_id, &upgrader_id);
+
+    let payload = BytesN::from_array(&env, &[0u8; 32]);
+    let public_key = admin_kp.public_key(&env);
+    let signature = admin_kp.sign(&env, &payload.to_array());
+
+    let dvn_client = LzDVNClient::new(&env, &setup.contract_id);
+    let hash = dvn_client.hash_call_data(&VID, &expiration, &calls);
+    let sig = setup.key_pairs[0].sign_bytes(&env, &hash);
+
+    let tx_auth = TransactionAuthData {
+        vid: VID,
+        expiration,
+        signatures: vec![&env, sig.clone()],
+        sender: Sender::Admin(public_key.clone(), signature.clone()),
+    };
+
+    let res = env.try_invoke_contract_check_auth::<DvnError>(
+        &setup.contract_id,
+        &payload,
+        tx_auth.into_val(&env),
+        &auth_contexts,
+    );
+    assert!(res.is_ok());
+
+    // Second attempt with same data should fail
+    let tx_auth2 = TransactionAuthData {
+        vid: VID,
+        expiration,
+        signatures: vec![&env, sig],
+        sender: Sender::Admin(public_key, signature),
+    };
+
+    let res2 = env.try_invoke_contract_check_auth::<DvnError>(
+        &setup.contract_id,
+        &payload,
+        tx_auth2.into_val(&env),
+        &auth_contexts,
+    );
+
+    assert_eq!(res2, Err(Ok(DvnError::HashAlreadyUsed)));
 }