Merge pull request #29 from lazor-kit/fix/wallet-discriminator-check

Fix/wallet discriminator check

Author: onspeedhp

program/src/processor/create_session.rs

@@ -110,6 +110,12 @@ pub fn process(
         return Err(ProgramError::IllegalOwner);
     }
 
+    // Validate Wallet Discriminator (Issue #7)
+    let wallet_data = unsafe { wallet_pda.borrow_data_unchecked() };
+    if wallet_data.is_empty() || wallet_data[0] != AccountDiscriminator::Wallet as u8 {
+        return Err(ProgramError::InvalidAccountData);
+    }
+
     // Verify Authorizer
     // Check removed: conditional writable check inside match
 

program/src/processor/execute.rs

@@ -4,7 +4,7 @@ use crate::{
     },
     compact::parse_compact_instructions,
     error::AuthError,
-    state::authority::AuthorityAccountHeader,
+    state::{authority::AuthorityAccountHeader, AccountDiscriminator},
 };
 use pinocchio::{
     account_info::AccountInfo,
@@ -61,6 +61,11 @@ pub fn process(
     if wallet_pda.owner() != program_id || authority_pda.owner() != program_id {
         return Err(ProgramError::IllegalOwner);
     }
+    // Validate Wallet Discriminator (Issue #7)
+    let wallet_data = unsafe { wallet_pda.borrow_data_unchecked() };
+    if wallet_data.is_empty() || wallet_data[0] != AccountDiscriminator::Wallet as u8 {
+        return Err(ProgramError::InvalidAccountData);
+    }
 
     if !authority_pda.is_writable() {
         return Err(ProgramError::InvalidAccountData);

program/src/processor/manage_authority.rs

@@ -133,6 +133,17 @@ pub fn process_add_authority(
     if admin_auth_pda.owner() != program_id {
         return Err(ProgramError::IllegalOwner);
     }
+    // Validate Wallet Discriminator (Issue #7)
+    let wallet_data = unsafe { wallet_pda.borrow_data_unchecked() };
+    if wallet_data.is_empty() || wallet_data[0] != AccountDiscriminator::Wallet as u8 {
+        return Err(ProgramError::InvalidAccountData);
+    }
+
+    // Validate Wallet Discriminator (Issue #7)
+    let wallet_data = unsafe { wallet_pda.borrow_data_unchecked() };
+    if wallet_data.is_empty() || wallet_data[0] != AccountDiscriminator::Wallet as u8 {
+        return Err(ProgramError::InvalidAccountData);
+    }
 
     // Validate system_program is the correct System Program (audit N2)
     if !sol_assert_bytes_eq(

program/src/processor/transfer_ownership.rs

@@ -111,6 +111,11 @@ pub fn process(
     if wallet_pda.owner() != program_id || current_owner.owner() != program_id {
         return Err(ProgramError::IllegalOwner);
     }
+    // Validate Wallet Discriminator (Issue #7)
+    let wallet_data = unsafe { wallet_pda.borrow_data_unchecked() };
+    if wallet_data.is_empty() || wallet_data[0] != AccountDiscriminator::Wallet as u8 {
+        return Err(ProgramError::InvalidAccountData);
+    }
 
     // Validate system_program is the correct System Program (audit N2)
     if !sol_assert_bytes_eq(

tests-e2e/TEST_ISSUES.md

@@ -30,10 +30,15 @@
 **Status**: ✅ Fixed
 **Fix**: Replaced hardcoded rent calculations with `Rent::minimum_balance(space)` in `create_wallet.rs` and `manage_authority.rs`. Verified by tests.
 
+### Issue #8 (Validation): Wallet Discriminator Check
+**Status**: ✅ Fixed
+**Fix**: Added `wallet_data[0] == AccountDiscriminator::Wallet` check in `create_session.rs`, `manage_authority.rs`, `execute.rs`, and `transfer_ownership.rs`.
+**Verification**: Added `Scenario 6: Wallet Discriminator Check` in `failures.rs`. Tested passing Authority PDA as Wallet PDA (Rejected).
+
 ## Current Status
 All E2E scenarios are PASSING.
 - Happy Path
-- Failures (5/5)
+- Failures (6/6)
 - Cross Wallet (3/3)
 - DoS Attack
 - Audit Validations

tests-e2e/src/scenarios/failures.rs

@@ -360,5 +360,56 @@ pub fn run(ctx: &mut TestContext) -> Result<()> {
     ctx.execute_tx_expect_error(remove_owner_tx)?;
     println!("✅ Admin Removing Owner Rejected.");
 
+    // Scenario 6: Wallet Discriminator Check (Issue #7)
+    // Attempt to use an Authority PDA (owned by program, but wrong discriminator) as the Wallet PDA
+    println!("\n[6/6] Testing Wallet Discriminator Validation...");
+
+    // We will try to call CreateSession using the Owner Authority PDA as the "Wallet PDA"
+    // The Owner Authority PDA is owned by the program, so it passes the owner check.
+    // However, it has Discriminator::Authority (2), not Wallet (1), so it should fail the new check.
+
+    let fake_wallet_pda = owner_auth_pda; // This is actually an Authority account
+
+    let bad_session_keypair = Keypair::new();
+    let (bad_session_pda, _) = Pubkey::find_program_address(
+        &[
+            b"session",
+            fake_wallet_pda.as_ref(), // Derived from the "fake" wallet
+            bad_session_keypair.pubkey().as_ref(),
+        ],
+        &ctx.program_id,
+    );
+
+    let mut bad_session_data = Vec::new();
+    bad_session_data.push(5); // CreateSession
+    bad_session_data.extend_from_slice(bad_session_keypair.pubkey().as_ref());
+    bad_session_data.extend_from_slice(&(current_slot + 100).to_le_bytes());
+
+    let bad_discriminator_ix = Instruction {
+        program_id: ctx.program_id.to_address(),
+        accounts: vec![
+            AccountMeta::new(Signer::pubkey(&ctx.payer).to_address(), true),
+            AccountMeta::new_readonly(fake_wallet_pda.to_address(), false), // FAKE WALLET (Authority Account)
+            AccountMeta::new_readonly(owner_auth_pda.to_address(), false), // Authorizer (Using same account as auth is technically weird but valid for this test)
+            AccountMeta::new(bad_session_pda.to_address(), false),
+            AccountMeta::new_readonly(solana_system_program::id().to_address(), false),
+            AccountMeta::new_readonly(solana_sysvar::rent::ID.to_address(), false),
+            AccountMeta::new_readonly(Signer::pubkey(&owner_keypair).to_address(), true),
+        ],
+        data: bad_session_data,
+    };
+
+    let message = Message::new(
+        &[bad_discriminator_ix],
+        Some(&Signer::pubkey(&ctx.payer).to_address()),
+    );
+    let mut bad_disc_tx = Transaction::new_unsigned(message);
+    bad_disc_tx.sign(&[&ctx.payer, &owner_keypair], ctx.svm.latest_blockhash());
+
+    // Expect InvalidAccountData (which is often generic error or custom depending on implementation)
+    // Our fix returns InvalidAccountData
+    ctx.execute_tx_expect_error(bad_disc_tx)?;
+    println!("✅ Invalid Wallet Discriminator Rejected.");
+
     Ok(())
 }