Skip to content

Commit

Permalink
feat: enable tss auth on spl whitelisting (#60)
Browse files Browse the repository at this point in the history
  • Loading branch information
@skosito
skosito authored Nov 3, 2024
1 parent 6367073 commit 177f563
Show file tree
Hide file tree
Showing 3 changed files with 172 additions and 14 deletions.
1 change: 1 addition & 0 deletions .gitignore
View file
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
.anchor/
/target
**/.DS_Store
node_modules
109 changes: 104 additions & 5 deletions programs/protocol-contracts-solana/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -90,12 +90,75 @@ pub mod gateway {
Ok(())
}

// whitelisting SPL tokens
pub fn whitelist_spl_mint(_ctx: Context<Whitelist>) -> Result<()> {
// whitelist new spl token
// in case signature is provided, check if tss is the signer, otherwise check if authority is pda.authority
// if succeeds, new whitelist entry account is created
pub fn whitelist_spl_mint(
ctx: Context<Whitelist>,
signature: [u8; 64],
recovery_id: u8,
message_hash: [u8; 32],
nonce: u64,
) -> Result<()> {
let pda = &mut ctx.accounts.pda;
let whitelist_candidate = &mut ctx.accounts.whitelist_candidate;
let authority = &ctx.accounts.authority;

// signature provided, recover and verify that tss is the signer
if signature != [0u8; 64] {
validate_whitelist_tss_signature(
pda,
whitelist_candidate,
signature,
recovery_id,
message_hash,
nonce,
"whitelist_spl_mint",
)?;
} else {
// no signature provided, fallback to authority check
require!(
authority.key() == pda.authority,
Errors::SignerIsNotAuthority
);
}

Ok(())
}

pub fn unwhitelist_spl_mint(_ctx: Context<Unwhitelist>) -> Result<()> {
// unwhitelist new spl token
// in case signature is provided, check if tss is the signer, otherwise check if authority is pda.authority
// if succeeds, whitelist entry account is deleted
pub fn unwhitelist_spl_mint(
ctx: Context<Unwhitelist>,
signature: [u8; 64],
recovery_id: u8,
message_hash: [u8; 32],
nonce: u64,
) -> Result<()> {
let pda = &mut ctx.accounts.pda;
let whitelist_candidate: &mut Account<'_, Mint> = &mut ctx.accounts.whitelist_candidate;
let authority = &ctx.accounts.authority;

// signature provided, recover and verify that tss is the signer
if signature != [0u8; 64] {
validate_whitelist_tss_signature(
pda,
whitelist_candidate,
signature,
recovery_id,
message_hash,
nonce,
"unwhitelist_spl_mint",
)?;
} else {
// no signature provided, fallback to authority check
require!(
authority.key() == pda.authority,
Errors::SignerIsNotAuthority
);
}

Ok(())
}

Expand Down Expand Up @@ -393,6 +456,42 @@ fn recover_eth_address(
Ok(eth_address)
}

// recover and verify tss signature for whitelist and unwhitelist instructions
fn validate_whitelist_tss_signature(
pda: &mut Account<Pda>,
whitelist_candidate: &mut Account<Mint>,
signature: [u8; 64],
recovery_id: u8,
message_hash: [u8; 32],
nonce: u64,
instruction_name: &str,
) -> Result<()> {
if nonce != pda.nonce {
msg!("mismatch nonce");
return err!(Errors::NonceMismatch);
}

let mut concatenated_buffer = Vec::new();
concatenated_buffer.extend_from_slice(instruction_name.as_bytes());
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());
require!(
message_hash == hash(&concatenated_buffer[..]).to_bytes(),
Errors::MessageHashMismatch
);

let address = recover_eth_address(&message_hash, recovery_id, &signature)?;
if address != pda.tss_address {
msg!("ECDSA signature error");
return err!(Errors::TSSAuthenticationFailed);
}

pda.nonce += 1;

Ok(())
}

#[derive(Accounts)]
pub struct Initialize<'info> {
#[account(mut)]
Expand Down Expand Up @@ -513,7 +612,7 @@ pub struct Whitelist<'info> {
pub whitelist_entry: Account<'info, WhitelistEntry>,
pub whitelist_candidate: Account<'info, Mint>,

#[account(mut, seeds = [b"meta"], bump, has_one = authority)]
#[account(mut, seeds = [b"meta"], bump)]
pub pda: Account<'info, Pda>,
#[account(mut)]
pub authority: Signer<'info>,
Expand All @@ -535,7 +634,7 @@ pub struct Unwhitelist<'info> {
pub whitelist_entry: Account<'info, WhitelistEntry>,
pub whitelist_candidate: Account<'info, Mint>,

#[account(mut, seeds = [b"meta"], bump, has_one = authority)]
#[account(mut, seeds = [b"meta"], bump)]
pub pda: Account<'info, Pda>,
#[account(mut)]
pub authority: Signer<'info>,
Expand Down
76 changes: 67 additions & 9 deletions tests/protocol-contracts-solana.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -190,7 +190,7 @@ describe("some tests", () => {
})

it("whitelist USDC spl token", async () => {
await gatewayProgram.methods.whitelistSplMint().accounts({
await gatewayProgram.methods.whitelistSplMint([], 0, [], new anchor.BN(0)).accounts({
whitelistCandidate: mint.publicKey,
}).signers([]).rpc();

Expand All @@ -199,15 +199,14 @@ describe("some tests", () => {
seeds,
gatewayProgram.programId,
);
let entry = await gatewayProgram.account.whitelistEntry.fetch(entryAddress)

try {
seeds = [Buffer.from("whitelist", "utf-8"), mint_fake.publicKey.toBuffer()];
[entryAddress] = anchor.web3.PublicKey.findProgramAddressSync(
seeds,
gatewayProgram.programId,
);
entry = await gatewayProgram.account.whitelistEntry.fetch(entryAddress);
await gatewayProgram.account.whitelistEntry.fetch(entryAddress);
} catch(err) {
expect(err.message).to.include("Account does not exist or has no data");
}
Expand Down Expand Up @@ -289,7 +288,6 @@ describe("some tests", () => {
// expect(account2.amount).to.be.eq(1_000_000n);

const pdaAccountData = await gatewayProgram.account.pda.fetch(pdaAccount);
const hexAddr = bufferToHex(Buffer.from(pdaAccountData.tssAddress));
const amount = new anchor.BN(500_000);
const nonce = pdaAccountData.nonce;
await withdrawSplToken(mint, usdcDecimals, amount, nonce, pda_ata, wallet_ata, wallet.publicKey, keyPair, gatewayProgram);
Expand Down Expand Up @@ -414,7 +412,7 @@ describe("some tests", () => {
})

it("unwhitelist SPL token and deposit should fail", async () => {
await gatewayProgram.methods.unwhitelistSplMint().accounts({
await gatewayProgram.methods.unwhitelistSplMint([], 0, [], new anchor.BN(0)).accounts({
whitelistCandidate: mint.publicKey,
}).rpc();

Expand All @@ -427,12 +425,75 @@ describe("some tests", () => {
});

it("re-whitelist SPL token and deposit should succeed", async () => {
await gatewayProgram.methods.whitelistSplMint().accounts({
await gatewayProgram.methods.whitelistSplMint([], 0, [], new anchor.BN(0)).accounts({
whitelistCandidate: mint.publicKey,
}).rpc();
await depositSplTokens(gatewayProgram, conn, wallet, mint, address);
});

it("unwhitelist SPL token using TSS signature and deposit should fail", async () => {
const pdaAccountData = await gatewayProgram.account.pda.fetch(pdaAccount);
const nonce = pdaAccountData.nonce;

const buffer = Buffer.concat([
Buffer.from("unwhitelist_spl_mint","utf-8"),
chain_id_bn.toArrayLike(Buffer, 'be', 8),
mint.publicKey.toBuffer(),
nonce.toArrayLike(Buffer, 'be', 8),
]);
const message_hash = keccak256(buffer);
const signature = keyPair.sign(message_hash, 'hex');
const { r, s, recoveryParam } = signature;
const signatureBuffer = Buffer.concat([
r.toArrayLike(Buffer, 'be', 32),
s.toArrayLike(Buffer, 'be', 32),
]);

await gatewayProgram.methods.unwhitelistSplMint(
Array.from(signatureBuffer),
Number(recoveryParam),
Array.from(message_hash),
nonce,
).accounts({
whitelistCandidate: mint.publicKey,
}).rpc();

try {
await depositSplTokens(gatewayProgram, conn, wallet, mint, address)
} catch (err) {
expect(err).to.be.instanceof(anchor.AnchorError);
expect(err.message).to.include("AccountNotInitialized");
}
});

it("re-whitelist SPL token using TSS signature and deposit should succeed", async () => {
const pdaAccountData = await gatewayProgram.account.pda.fetch(pdaAccount);
const nonce = pdaAccountData.nonce;

const buffer = Buffer.concat([
Buffer.from("whitelist_spl_mint","utf-8"),
chain_id_bn.toArrayLike(Buffer, 'be', 8),
mint.publicKey.toBuffer(),
nonce.toArrayLike(Buffer, 'be', 8),
]);
const message_hash = keccak256(buffer);
const signature = keyPair.sign(message_hash, 'hex');
const { r, s, recoveryParam } = signature;
const signatureBuffer = Buffer.concat([
r.toArrayLike(Buffer, 'be', 32),
s.toArrayLike(Buffer, 'be', 32),
]);

await gatewayProgram.methods.whitelistSplMint(
Array.from(signatureBuffer),
Number(recoveryParam),
Array.from(message_hash),
nonce,
).accounts({
whitelistCandidate: mint.publicKey,
}).rpc();
await depositSplTokens(gatewayProgram, conn, wallet, mint, address);
});

it("update TSS address", async () => {
const newTss = new Uint8Array(20);
Expand Down Expand Up @@ -470,9 +531,6 @@ describe("some tests", () => {
}
});




const newAuthority = anchor.web3.Keypair.generate();
it("update authority", async () => {
await gatewayProgram.methods.updateAuthority(newAuthority.publicKey).accounts({
Expand Down

0 comments on commit 177f563

Please sign in to comment.