remove attack vectors

AztecProtocol · Sep 14, 2023 · d6c4f33 · d6c4f33
1 parent 5ecec2f
commit d6c4f33
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 23 deletions.
diff --git a/yarn-project/end-to-end/src/e2e_token_bridge.test.ts b/yarn-project/end-to-end/src/e2e_token_bridge.test.ts
@@ -139,9 +139,7 @@ describe('e2e_token_bridge_contract', () => {
 
     // 3. Consume message on aztec and mint publicly
     logger('Consuming messages on L2');
-    const tx = bridge.methods
-      .mint_public({ address: ownerAddress }, bridgeAmount, messageKey, secret, { address: ethAccount.toField() })
-      .send();
+    const tx = bridge.methods.mint_public(bridgeAmount, messageKey, secret, { address: ethAccount.toField() }).send();
     const receipt = await tx.wait();
     expect(receipt.status).toBe(TxStatus.MINED);
     const afterBalance = await token.methods.balance_of_public({ address: ownerAddress }).view();
@@ -174,13 +172,7 @@ describe('e2e_token_bridge_contract', () => {
 
     // 5. Withdraw from L2 bridge
     const withdrawTx = bridge.methods
-      .withdraw_public(
-        { address: ownerAddress },
-        { address: ethAccount.toField() },
-        withdrawAmount,
-        { address: EthAddress.ZERO.toField() },
-        nonce,
-      )
+      .withdraw_public({ address: ethAccount.toField() }, withdrawAmount, { address: EthAddress.ZERO.toField() }, nonce)
       .send();
     const withdrawReceipt = await withdrawTx.wait();
     expect(withdrawReceipt.status).toBe(TxStatus.MINED);
@@ -213,9 +205,7 @@ describe('e2e_token_bridge_contract', () => {
 
     // 3. Consume message on aztec and mint publicly
     logger('Consuming messages on L2');
-    const tx = bridge.methods
-      .mint({ address: ownerAddress }, bridgeAmount, messageKey, secret, { address: ethAccount.toField() })
-      .send();
+    const tx = bridge.methods.mint(bridgeAmount, messageKey, secret, { address: ethAccount.toField() }).send();
     const receipt = await tx.wait();
     expect(receipt.status).toBe(TxStatus.MINED);
     const txClaim = token.methods.redeem_shield({ address: ownerAddress }, bridgeAmount, secret).send();
@@ -248,7 +238,6 @@ describe('e2e_token_bridge_contract', () => {
     const withdrawTx = bridge.methods
       .withdraw(
         { address: token.address },
-        { address: ownerAddress },
         { address: ethAccount.toField() },
         withdrawAmount,
         { address: EthAddress.ZERO.toField() },

diff --git a/yarn-project/noir-contracts/src/contracts/token_bridge_contract/src/main.nr b/yarn-project/noir-contracts/src/contracts/token_bridge_contract/src/main.nr
@@ -47,37 +47,35 @@ contract TokenBridge {
     // Consumes a L1->L2 message and calls the token contract to mint the appropriate amount publicly
     #[aztec(public)]
     fn mint_public(
-        recipient: AztecAddress,
         amount: Field,
         msg_key: Field, // L1 to L2 message key as derived from the inbox contract
         secret: Field,
         canceller: EthereumAddress,
     ) -> Field {
         let storage = Storage::init(Context::public(&mut context));
 
-        let content_hash = get_mint_content_hash(amount, recipient.address, canceller.address);
+        let content_hash = get_mint_content_hash(amount, context.msg_sender(), canceller.address);
 
         // Consume message and emit nullifier
         context.consume_l1_to_l2_message(msg_key, content_hash, secret);
 
         // Mint token on L2
-        Token::at(storage.token.read()).mint_public(context, recipient.address, amount)
+        Token::at(storage.token.read()).mint_public(context, context.msg_sender(), amount)
     }
 
     // Consumes a L1->L2 message and calls the token contract to mint the appropriate amount in private assets
     // User needs to call token.redeem_shield() to get the private assets
     // This method is public because it accesses public storage. For similar reasons, the corresponding call on the token is also public
     #[aztec(public)]
     fn mint(
-        recipient: AztecAddress,
         amount: Field,
         msg_key: Field, // L1 to L2 message key as derived from the inbox contract
         secret: Field,
         canceller: EthereumAddress,
     ) -> Field {
         let storage = Storage::init(Context::public(&mut context));
 
-        let content_hash = get_mint_content_hash(amount, recipient.address, canceller.address);
+        let content_hash = get_mint_content_hash(amount, context.msg_sender(), canceller.address);
 
         // Consume message and emit nullifier
         context.consume_l1_to_l2_message(msg_key, content_hash, secret);
@@ -91,7 +89,6 @@ contract TokenBridge {
     // Requires `from` to give approval to the bridge to burn tokens on their behalf using witness signatures 
     #[aztec(public)]
     fn withdraw_public(
-        from: AztecAddress, // aztec address to withdraw from
         recipient: EthereumAddress, // ethereum address to withdraw to
         amount: Field,
         callerOnL1: EthereumAddress, // ethereum address that can call this function on the L1 portal (0x0 if anyone can call)
@@ -100,7 +97,7 @@ contract TokenBridge {
         let storage = Storage::init(Context::public(&mut context));
 
         // Burn tokens on L2
-        let return_value = Token::at(storage.token.read()).burn_public(context, from.address, amount, nonce);
+        let return_value = Token::at(storage.token.read()).burn_public(context, context.msg_sender(), amount, nonce);
 
         let content = get_withdraw_content_hash(amount, recipient.address, callerOnL1.address);
 
@@ -115,15 +112,14 @@ contract TokenBridge {
     #[aztec(private)]
     fn withdraw(
         token:AztecAddress, // can't read public storage in private, so pass the token and call an internal public fn to check if provided token is as expected.
-        from: AztecAddress, // aztec address to withdraw from
         recipient: EthereumAddress, // ethereum address to withdraw to
         amount: Field,
         callerOnL1: EthereumAddress, // ethereum address that can call this function on the L1 portal (0x0 if anyone can call)
         nonce: Field, // used in creating the approval message (to prevent replay attacks)
     ) -> Field {
         // can't read public storage (`storage.token`) in private so let the user pass it in
         // later assert that this token address is as expected
-        let return_value = Token::at(token.address).burn(&mut context, from.address, amount, nonce);
+        let return_value = Token::at(token.address).burn(&mut context, context.msg_sender(), amount, nonce);
 
         let content = get_withdraw_content_hash(amount, recipient.address, callerOnL1.address);
         // Emit the l2 to l1 message