From 007b01860922ecf3f467eef7208d81d43e06b742 Mon Sep 17 00:00:00 2001
From: brewmaster012 <88689859+brewmaster012@users.noreply.github.com>
Date: Wed, 19 Jun 2024 20:46:55 -0500
Subject: [PATCH] test: add nonce mismatch case

---
 programs/protocol-contracts-solana/src/lib.rs | 24 ++++++++++++------
 tests/protocol-contracts-solana.ts            | 25 ++++++++++++++++---
 2 files changed, 38 insertions(+), 11 deletions(-)

diff --git a/programs/protocol-contracts-solana/src/lib.rs b/programs/protocol-contracts-solana/src/lib.rs
index 2391988..1dec055 100644
--- a/programs/protocol-contracts-solana/src/lib.rs
+++ b/programs/protocol-contracts-solana/src/lib.rs
@@ -6,9 +6,15 @@ use solana_program::secp256k1_recover::secp256k1_recover;
 use std::mem::size_of;
 
 #[error_code]
-pub enum MyError {
-    #[msg("This is a custom error message")]
-    CustomError,
+pub enum Errors {
+    #[msg("SignerIsNotAuthority")]
+    SignerIsNotAuthority,
+    #[msg("InsufficientPoints")]
+    InsufficientPoints,
+    #[msg("NonceMismatch")]
+    NonceMismatch,
+    #[msg("TSSAuthenticationFailed")]
+    TSSAuthenticationFailed,
 }
 
 declare_id!("9WSwbVLthCsJXABeDJcVcw4UQMYuoNLTJTLqueFXU5Q2");
@@ -55,13 +61,13 @@ pub mod gateway {
 
         if nonce != pda.nonce {
             msg!("mismatch nonce");
-            return err!(MyError::CustomError);
+            return err!(Errors::NonceMismatch);
         }
         let address = recover_eth_address(&message_hash, recovery_id, &signature)?; // ethereum address is the last 20 Bytes of the hashed pubkey
         msg!("recovered address {:?}", address);
         if address != pda.tss_address {
             msg!("ECDSA signature error");
-            return err!(MyError::CustomError);
+            return err!(Errors::TSSAuthenticationFailed);
         }
 
         // transfer amount of SOL from PDA to the payer
@@ -85,13 +91,13 @@ pub mod gateway {
         // let program_id = &mut ctx.accounts
         if nonce != pda.nonce {
             msg!("mismatch nonce");
-            return err!(MyError::CustomError);
+            return err!(Errors::NonceMismatch);
         }
         let address = recover_eth_address(&message_hash, recovery_id, &signature)?; // ethereum address is the last 20 Bytes of the hashed pubkey
         msg!("recovered address {:?}", address);
         if address != pda.tss_address {
             msg!("ECDSA signature error");
-            return err!(MyError::CustomError);
+            return err!(Errors::TSSAuthenticationFailed);
         }
 
         let token = &ctx.accounts.token_program;
@@ -102,13 +108,15 @@ pub mod gateway {
             anchor_spl::token::Transfer {
                 from: ctx.accounts.from.to_account_info(),
                 to: ctx.accounts.to.to_account_info(),
-                authority: ctx.accounts.pda.to_account_info(),
+                authority: pda.to_account_info(),
             },
             signer_seeds,
         );
         transfer(xfer_ctx, amount)?;
         msg!("withdraw spl token successfully");
 
+        pda.nonce += 1;
+
         Ok(())
     }
 
diff --git a/tests/protocol-contracts-solana.ts b/tests/protocol-contracts-solana.ts
index 9bd64a6..8523f9c 100644
--- a/tests/protocol-contracts-solana.ts
+++ b/tests/protocol-contracts-solana.ts
@@ -3,6 +3,8 @@ import {Program} from "@coral-xyz/anchor";
 import {Gateway} from "../target/types/gateway";
 import * as spl from "@solana/spl-token";
 import {TOKEN_PROGRAM_ID} from "@solana/spl-token";
+import { expect } from 'chai';
+
 
 const fromHexString = (hexString) =>
     hexString.match(/.{1,2}/g).map((byte) => parseInt(byte, 16));
@@ -97,9 +99,10 @@ describe("some tests", () => {
             wallet,
             1_000_000
         );
-        console.log("xfer tx hash", tx_xfer);
+        // console.log("xfer tx hash", tx_xfer);
         const account2 = await spl.getAccount(conn, pda_ata.address);
-        console.log("B4 withdraw: Account balance:", account2.amount.toString());
+        expect(account2.amount).to.be.eq(1_000_000n);
+        // console.log("B4 withdraw: Account balance:", account2.amount.toString());
 
 
         const pdaAccountData = await gatewayProgram.account.pda.fetch(pdaAccount);
@@ -118,7 +121,23 @@ describe("some tests", () => {
             }).rpc();
 
         const account3 = await spl.getAccount(conn, pda_ata.address);
-        console.log("After withdraw: Account balance:", account3.amount.toString());
+        expect(account3.amount).to.be.eq(500_000n);
+
+
+        try {
+            (await gatewayProgram.methods.withdrawSplToken(new anchor.BN(500_000), signature, 0, message_hash, nonce)
+                .accounts({
+                    from: pda_ata.address,
+                    to: wallet_ata,
+                }).rpc());
+            throw new Error("Expected error not thrown"); // This line will make the test fail if no error is thrown
+        } catch (err) {
+            expect(err).to.be.instanceof(anchor.AnchorError);
+            expect(err.message).to.include("NonceMismatch");
+            const account4 = await spl.getAccount(conn, pda_ata.address);
+            console.log("After 2nd withdraw: Account balance:", account4.amount.toString());
+            expect(account4.amount).to.be.eq(500_000n);
+        }
 
     });