diff --git a/programs/gateway/src/lib.rs b/programs/gateway/src/lib.rs index f3e14c2..70b8b63 100644 --- a/programs/gateway/src/lib.rs +++ b/programs/gateway/src/lib.rs @@ -31,6 +31,15 @@ pub enum Errors { EmptyReceiver, } +/// Enumeration for instruction identifiers in message hashes. +#[repr(u8)] +enum InstructionId { + Withdraw = 1, + WithdrawSplToken = 2, + WhitelistSplToken = 3, + UnwhitelistSplToken = 4, +} + declare_id!("ZETAjseVjuFsxdRxo6MmTCvqFwb3ZHUx56Co3vCmGis"); #[program] @@ -159,7 +168,7 @@ pub mod gateway { signature, recovery_id, nonce, - "whitelist_spl_mint", + InstructionId::WhitelistSplToken as u8, )?; } else { require!( @@ -202,7 +211,7 @@ pub mod gateway { signature, recovery_id, nonce, - "unwhitelist_spl_mint", + InstructionId::UnwhitelistSplToken as u8, )?; } else { require!( @@ -377,7 +386,8 @@ pub mod gateway { } let mut concatenated_buffer = Vec::new(); - concatenated_buffer.extend_from_slice("withdraw".as_bytes()); + concatenated_buffer.extend_from_slice(b"ZETACHAIN"); + concatenated_buffer.push(InstructionId::Withdraw as u8); concatenated_buffer.extend_from_slice(&pda.chain_id.to_be_bytes()); concatenated_buffer.extend_from_slice(&nonce.to_be_bytes()); concatenated_buffer.extend_from_slice(&amount.to_be_bytes()); @@ -435,7 +445,8 @@ pub mod gateway { } let mut concatenated_buffer = Vec::new(); - concatenated_buffer.extend_from_slice("withdraw_spl_token".as_bytes()); + concatenated_buffer.extend_from_slice(b"ZETACHAIN"); + concatenated_buffer.push(InstructionId::WithdrawSplToken as u8); concatenated_buffer.extend_from_slice(&pda.chain_id.to_be_bytes()); concatenated_buffer.extend_from_slice(&nonce.to_be_bytes()); concatenated_buffer.extend_from_slice(&amount.to_be_bytes()); @@ -581,7 +592,7 @@ fn validate_whitelist_tss_signature( signature: [u8; 64], recovery_id: u8, nonce: u64, - instruction_name: &str, + instruction: u8, ) -> Result<()> { if nonce != pda.nonce { msg!( @@ -593,10 +604,11 @@ fn validate_whitelist_tss_signature( } let mut concatenated_buffer = Vec::new(); - concatenated_buffer.extend_from_slice(instruction_name.as_bytes()); + concatenated_buffer.extend_from_slice(b"ZETACHAIN"); + concatenated_buffer.push(instruction); concatenated_buffer.extend_from_slice(&pda.chain_id.to_be_bytes()); - concatenated_buffer.extend_from_slice(&whitelist_candidate.key().to_bytes()); concatenated_buffer.extend_from_slice(&nonce.to_be_bytes()); + concatenated_buffer.extend_from_slice(&whitelist_candidate.key().to_bytes()); let computed_message_hash = hash(&concatenated_buffer[..]).to_bytes(); msg!("Computed message hash: {:?}", computed_message_hash); diff --git a/tests/gateway.ts b/tests/gateway.ts index 2a3e5fa..4c20590 100644 --- a/tests/gateway.ts +++ b/tests/gateway.ts @@ -96,7 +96,8 @@ async function withdrawSplToken( gatewayProgram: Program ) { const buffer = Buffer.concat([ - Buffer.from("withdraw_spl_token", "utf-8"), + Buffer.from("ZETACHAIN", "utf-8"), + Buffer.from([0x02]), chain_id_bn.toArrayLike(Buffer, "be", 8), nonce.toArrayLike(Buffer, "be", 8), amount.toArrayLike(Buffer, "be", 8), @@ -373,7 +374,8 @@ describe("Gateway", () => { try { const nonce2 = nonce.addn(1); const buffer = Buffer.concat([ - Buffer.from("withdraw_spl_token", "utf-8"), + Buffer.from("ZETACHAIN", "utf-8"), + Buffer.from([0x02]), chain_id_bn.toArrayLike(Buffer, "be", 8), nonce2.toArrayLike(Buffer, "be", 8), amount.toArrayLike(Buffer, "be", 8), @@ -438,7 +440,8 @@ describe("Gateway", () => { wallet.publicKey ); const buffer = Buffer.concat([ - Buffer.from("withdraw", "utf-8"), + Buffer.from("ZETACHAIN", "utf-8"), + Buffer.from([0x01]), chain_id_bn.toArrayLike(Buffer, "be", 8), nonce.toArrayLike(Buffer, "be", 8), amount.toArrayLike(Buffer, "be", 8), @@ -479,7 +482,8 @@ describe("Gateway", () => { ); const buffer = Buffer.concat([ - Buffer.from("withdraw", "utf-8"), + Buffer.from("ZETACHAIN", "utf-8"), + Buffer.from([0x01]), chain_id_bn.toArrayLike(Buffer, "be", 8), nonce.toArrayLike(Buffer, "be", 8), amount.toArrayLike(Buffer, "be", 8), @@ -522,7 +526,8 @@ describe("Gateway", () => { ); const buffer = Buffer.concat([ - Buffer.from("withdraw", "utf-8"), + Buffer.from("ZETACHAIN", "utf-8"), + Buffer.from([0x01]), chain_id_bn.toArrayLike(Buffer, "be", 8), nonce.subn(1).toArrayLike(Buffer, "be", 8), // wrong nonce amount.toArrayLike(Buffer, "be", 8), @@ -640,7 +645,8 @@ describe("Gateway", () => { ); const buffer = Buffer.concat([ - Buffer.from("withdraw_spl_token", "utf-8"), + Buffer.from("ZETACHAIN", "utf-8"), + Buffer.from([0x02]), chain_id_bn.toArrayLike(Buffer, "be", 8), nonce.subn(1).toArrayLike(Buffer, "be", 8), // wrong nonce amount.toArrayLike(Buffer, "be", 8), @@ -754,10 +760,11 @@ describe("Gateway", () => { const nonce = pdaAccountData.nonce; const buffer = Buffer.concat([ - Buffer.from("unwhitelist_spl_mint", "utf-8"), + Buffer.from("ZETACHAIN", "utf-8"), + Buffer.from([0x04]), chain_id_bn.toArrayLike(Buffer, "be", 8), - mint.publicKey.toBuffer(), nonce.toArrayLike(Buffer, "be", 8), + mint.publicKey.toBuffer(), ]); const message_hash = keccak256(buffer); const signature = keyPair.sign(message_hash, "hex"); @@ -792,10 +799,11 @@ describe("Gateway", () => { const nonce = pdaAccountData.nonce; const buffer = Buffer.concat([ - Buffer.from("whitelist_spl_mint", "utf-8"), + Buffer.from("ZETACHAIN", "utf-8"), + Buffer.from([0x03]), chain_id_bn.toArrayLike(Buffer, "be", 8), - mint.publicKey.toBuffer(), nonce.toArrayLike(Buffer, "be", 8), + mint.publicKey.toBuffer(), ]); const message_hash = keccak256(buffer); const signature = keyPair.sign(message_hash, "hex"); @@ -823,10 +831,11 @@ describe("Gateway", () => { const nonce = pdaAccountData.nonce; const buffer = Buffer.concat([ - Buffer.from("whitelist_spl_mint", "utf-8"), + Buffer.from("ZETACHAIN", "utf-8"), + Buffer.from([0x03]), chain_id_bn.toArrayLike(Buffer, "be", 8), - mint.publicKey.toBuffer(), nonce.subn(1).toArrayLike(Buffer, "be", 8), // wrong nonce + mint.publicKey.toBuffer(), ]); const message_hash = keccak256(buffer); const signature = keyPair.sign(message_hash, "hex"); @@ -858,10 +867,11 @@ describe("Gateway", () => { const nonce = pdaAccountData.nonce; const buffer = Buffer.concat([ - Buffer.from("whitelist_spl_mint", "utf-8"), + Buffer.from("ZETACHAIN", "utf-8"), + Buffer.from([0x03]), chain_id_bn.toArrayLike(Buffer, "be", 8), + nonce.toArrayLike(Buffer, "be", 8), mint.publicKey.toBuffer(), - nonce.toArrayLike(Buffer, "be", 8), // wrong nonce ]); const message_hash = keccak256(buffer); const signature = keyPair.sign(message_hash, "hex");