fix(nft-meta-data-pointer): require the player's NFT token account in chop_tree and reset the tree when it reaches MAX_WOOD_PER_TREE

tokens/token-2022/nft-meta-data-pointer/anchor/programs/extension_nft/src/errors.rs

@@ -10,4 +10,6 @@ pub enum GameErrorCode {
     InvalidMintAccountSpace,
     #[msg("Cant initialize metadata_pointer")]
     CantInitializeMetadataPointer,
+    #[msg("Player does not own this NFT")]
+    NftNotOwned,
 }

tokens/token-2022/nft-meta-data-pointer/anchor/programs/extension_nft/src/instructions/chop_tree.rs

@@ -2,9 +2,9 @@ pub use crate::errors::GameErrorCode;
 pub use crate::state::game_data::GameData;
 use crate::{state::player_data::PlayerData, NftAuthority};
 use anchor_lang::prelude::*;
-use anchor_spl::token_interface::{Token2022};
-use session_keys::{Session, SessionToken};
 use anchor_lang::solana_program::program::invoke_signed;
+use anchor_spl::token_interface::{Mint, Token2022, TokenAccount};
+use session_keys::{Session, SessionToken};
 
 pub fn chop_tree(ctx: Context<ChopTree>, counter: u16, amount: u64) -> Result<()> {
     let account: &mut ChopTree<'_> = ctx.accounts;
@@ -37,7 +37,9 @@ pub fn chop_tree(ctx: Context<ChopTree>, counter: u16, amount: u64) -> Result<()
             &anchor_spl::token_2022::spl_token_2022::id(),
             ctx.accounts.mint.to_account_info().key,
             ctx.accounts.nft_authority.to_account_info().key,
-            anchor_spl::token_2022_extensions::spl_token_metadata_interface::state::Field::Key("wood".to_string()),
+            anchor_spl::token_2022_extensions::spl_token_metadata_interface::state::Field::Key(
+                "wood".to_string(),
+            ),
             ctx.accounts.player.wood.to_string(),
         ),
         &[
@@ -84,10 +86,17 @@ pub struct ChopTree<'info> {
     #[account(mut)]
     pub signer: Signer<'info>,
     pub system_program: Program<'info, System>,
-    /// CHECK: Make sure the ata to the mint is actually owned by the signer
     #[account(mut)]
-    pub mint: AccountInfo<'info>,
-    #[account(  
+    pub mint: InterfaceAccount<'info, Mint>,
+    // The player must hold the NFT they are updating
+    #[account(
+        associated_token::mint = mint,
+        associated_token::authority = player.authority,
+        associated_token::token_program = token_program,
+        constraint = player_token_account.amount == 1 @ GameErrorCode::NftNotOwned,
+    )]
+    pub player_token_account: InterfaceAccount<'info, TokenAccount>,
+    #[account(
         init_if_needed,
         seeds = [b"nft_authority".as_ref()],
         bump,

tokens/token-2022/nft-meta-data-pointer/anchor/programs/extension_nft/src/instructions/mint_nft.rs

@@ -184,7 +184,7 @@ pub struct MintNft<'info> {
     pub token_program: Program<'info, Token2022>,
     /// CHECK: We will create this one for the user
     #[account(mut)]
-    pub token_account: AccountInfo<'info>,
+    pub token_account: UncheckedAccount<'info>,
     #[account(mut)]
     pub mint: Signer<'info>,
     pub rent: Sysvar<'info, Rent>,

tokens/token-2022/nft-meta-data-pointer/anchor/programs/extension_nft/src/state/game_data.rs

@@ -11,7 +11,7 @@ impl GameData {
     pub fn on_tree_chopped(&mut self, amount_chopped: u64) -> Result<()> {
         match self.total_wood_collected.checked_add(amount_chopped) {
             Some(v) => {
-                if self.total_wood_collected >= MAX_WOOD_PER_TREE {
+                if v >= MAX_WOOD_PER_TREE {
                     self.total_wood_collected = 0;
                     msg!("Tree successfully chopped. New Tree coming up.");
                 } else {
@@ -27,3 +27,26 @@ impl GameData {
         Ok(())
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn keeps_counting_below_max() {
+        let mut game = GameData {
+            total_wood_collected: 0,
+        };
+        game.on_tree_chopped(1).unwrap();
+        assert_eq!(game.total_wood_collected, 1);
+    }
+
+    #[test]
+    fn resets_when_total_reaches_max() {
+        let mut game = GameData {
+            total_wood_collected: MAX_WOOD_PER_TREE - 1,
+        };
+        game.on_tree_chopped(1).unwrap();
+        assert_eq!(game.total_wood_collected, 0);
+    }
+}

tokens/token-2022/nft-meta-data-pointer/anchor/tests/lumberjack.ts

@@ -3,49 +3,128 @@ import * as anchor from '@anchor-lang/core';
 import {
     ASSOCIATED_TOKEN_PROGRAM_ID,
     getAssociatedTokenAddressSync,
-    getOrCreateAssociatedTokenAccount,
+    getTokenMetadata,
     TOKEN_2022_PROGRAM_ID,
 } from '@solana/spl-token';
-import { Keypair } from '@solana/web3.js';
+import { Keypair, PublicKey } from '@solana/web3.js';
+import { assert } from 'chai';
 import type { ExtensionNft } from '../target/types/extension_nft';
 
 describe('extension_nft', () => {
     const provider = anchor.AnchorProvider.env();
     anchor.setProvider(provider);
     const program = anchor.workspace.ExtensionNft as Program<ExtensionNft>;
+    const connection = provider.connection;
     const payer = provider.wallet as anchor.Wallet;
 
+    const LEVEL_SEED = 'level_1';
+    const mint = new Keypair();
+
+    const ataOf = (owner: PublicKey) =>
+        getAssociatedTokenAddressSync(mint.publicKey, owner, false, TOKEN_2022_PROGRAM_ID, ASSOCIATED_TOKEN_PROGRAM_ID);
+
+    const playerPdaOf = (authority: PublicKey) =>
+        PublicKey.findProgramAddressSync([Buffer.from('player'), authority.toBuffer()], program.programId)[0];
+
+    const woodMetadataOf = async (mintKey: PublicKey) => {
+        const metadata = await getTokenMetadata(connection, mintKey, 'confirmed', TOKEN_2022_PROGRAM_ID);
+        return metadata?.additionalMetadata.find(([key]) => key === 'wood')?.[1];
+    };
+
+    const airdrop = async (to: PublicKey) => {
+        const sig = await connection.requestAirdrop(to, 1e9);
+        await connection.confirmTransaction(sig, 'confirmed');
+    };
+
     it('Mint nft!', async () => {
-        const balance = await anchor.getProvider().connection.getBalance(payer.publicKey);
+        const balance = await connection.getBalance(payer.publicKey);
 
         if (balance < 1e8) {
-            const res = await anchor.getProvider().connection.requestAirdrop(payer.publicKey, 1e9);
-            await anchor.getProvider().connection.confirmTransaction(res, 'confirmed');
+            await airdrop(payer.publicKey);
         }
 
-        const mint = new Keypair();
         console.log('Mint public key', mint.publicKey.toBase58());
 
-        const destinationTokenAccount = getAssociatedTokenAddressSync(
-            mint.publicKey,
-            payer.publicKey,
-            false,
-            TOKEN_2022_PROGRAM_ID,
-            ASSOCIATED_TOKEN_PROGRAM_ID,
-        );
-
-        getOrCreateAssociatedTokenAccount;
         const tx = await program.methods
             .mintNft()
             .accounts({
                 signer: payer.publicKey,
-                tokenAccount: destinationTokenAccount,
+                tokenAccount: ataOf(payer.publicKey),
                 mint: mint.publicKey,
             })
             .signers([mint])
             .rpc();
 
         console.log('Mint nft tx', tx);
-        await anchor.getProvider().connection.confirmTransaction(tx, 'confirmed');
+        await connection.confirmTransaction(tx, 'confirmed');
+    });
+
+    it('Init player', async () => {
+        await program.methods
+            .initPlayer(LEVEL_SEED)
+            .accounts({ signer: payer.publicKey })
+            .rpc({ commitment: 'confirmed' });
+
+        const player = await program.account.playerData.fetch(playerPdaOf(payer.publicKey), 'confirmed');
+        assert.isTrue(player.authority.equals(payer.publicKey));
+        assert.strictEqual(player.wood.toNumber(), 0);
+    });
+
+    it('Chop tree with own NFT updates the wood metadata field', async () => {
+        for (const counter of [1, 2]) {
+            await program.methods
+                .chopTree(LEVEL_SEED, counter)
+                .accountsPartial({
+                    sessionToken: null,
+                    player: playerPdaOf(payer.publicKey),
+                    signer: payer.publicKey,
+                    mint: mint.publicKey,
+                    playerTokenAccount: ataOf(payer.publicKey),
+                })
+                .rpc({ commitment: 'confirmed' });
+        }
+
+        const player = await program.account.playerData.fetch(playerPdaOf(payer.publicKey), 'confirmed');
+        assert.strictEqual(player.wood.toNumber(), 2);
+        assert.strictEqual(await woodMetadataOf(mint.publicKey), '2');
+    });
+
+    it("Chop tree with another player's NFT is rejected", async () => {
+        const attacker = Keypair.generate();
+        await airdrop(attacker.publicKey);
+
+        await program.methods
+            .initPlayer(LEVEL_SEED)
+            .accounts({ signer: attacker.publicKey })
+            .signers([attacker])
+            .rpc({ commitment: 'confirmed' });
+
+        let rejected = false;
+        try {
+            await program.methods
+                .chopTree(LEVEL_SEED, 1)
+                .accountsPartial({
+                    sessionToken: null,
+                    player: playerPdaOf(attacker.publicKey),
+                    signer: attacker.publicKey,
+                    mint: mint.publicKey,
+                    playerTokenAccount: ataOf(attacker.publicKey),
+                })
+                .signers([attacker])
+                .rpc({ commitment: 'confirmed' });
+        } catch (error) {
+            rejected = true;
+            console.log('chop_tree rejected:', error instanceof Error ? error.message : String(error));
+        }
+
+        assert.strictEqual(
+            await woodMetadataOf(mint.publicKey),
+            '2',
+            'wood metadata of the foreign NFT was overwritten',
+        );
+        assert.isTrue(rejected, "chop_tree with another player's mint must be rejected");
+
+        const attackerPlayer = await program.account.playerData.fetch(playerPdaOf(attacker.publicKey), 'confirmed');
+        assert.strictEqual(attackerPlayer.wood.toNumber(), 0);
     });
 });

tokens/token-2022/nft-meta-data-pointer/app/components/ChopTreeButton.tsx

@@ -2,6 +2,7 @@ import { Button, HStack, VStack } from '@chakra-ui/react';
 import { useSessionWallet } from '@magicblock-labs/gum-react-sdk';
 import { useKitTransactionSigner } from '@solana/connector/react';
 import { createNoopSigner, type Address } from '@solana/kit';
+import { findAssociatedTokenPda, TOKEN_2022_PROGRAM_ADDRESS } from '@solana-program/token-2022';
 import Image from 'next/image';
 import { useCallback, useState } from 'react';
 import { useGameState } from '@/contexts/GameStateProvider';
@@ -24,7 +25,7 @@ const ChopTreeButton = () => {
 
     const onChopClick = useCallback(async () => {
         setIsLoadingSession(true);
-        if (!playerDataPDA || !sessionWallet?.publicKey || !sessionWallet.sessionToken) {
+        if (!signer || !playerDataPDA || !sessionWallet?.publicKey || !sessionWallet.sessionToken) {
             setIsLoadingSession(false);
             return;
         }
@@ -59,11 +60,18 @@ const ChopTreeButton = () => {
         }
 
         try {
+            const [playerTokenAccount] = await findAssociatedTokenPda({
+                owner: signer.address,
+                mint: nft.id as Address,
+                tokenProgram: TOKEN_2022_PROGRAM_ADDRESS,
+            });
+
             const instruction = await getChopTreeInstructionAsync({
                 sessionToken: sessionWallet.sessionToken as Address,
                 player: playerDataPDA,
                 signer: createNoopSigner(sessionWallet.publicKey.toBase58() as Address),
                 mint: nft.id as Address,
+                playerTokenAccount,
                 nftAuthority,
                 levelSeed: GAME_DATA_SEED,
                 counter: transactionCounter,
@@ -83,7 +91,7 @@ const ChopTreeButton = () => {
         } finally {
             setIsLoadingSession(false);
         }
-    }, [sessionWallet, nftState, playerDataPDA, transactionCounter]);
+    }, [signer, sessionWallet, nftState, playerDataPDA, transactionCounter]);
 
     const onChopMainWalletClick = useCallback(async () => {
         if (!signer || !playerDataPDA) return;
@@ -126,10 +134,17 @@ const ChopTreeButton = () => {
                 typeof nft.authorities[0] === 'string' ? nft.authorities[0] : nft.authorities[0].address;
             console.log('NFTid', nft.id, 'NFT authority', nftAuthorityAddress);
 
+            const [playerTokenAccount] = await findAssociatedTokenPda({
+                owner: signer.address,
+                mint: nft.id as Address,
+                tokenProgram: TOKEN_2022_PROGRAM_ADDRESS,
+            });
+
             const instruction = await getChopTreeInstructionAsync({
                 player: playerDataPDA,
                 signer,
                 mint: nft.id as Address,
+                playerTokenAccount,
                 nftAuthority: nftAuthorityAddress as Address,
                 levelSeed: GAME_DATA_SEED,
                 counter: transactionCounter,

tokens/token-2022/nft-meta-data-pointer/app/idl/extension_nft.json

@@ -57,6 +57,33 @@
                     "name": "mint",
                     "writable": true
                 },
+                {
+                    "name": "player_token_account",
+                    "pda": {
+                        "seeds": [
+                            {
+                                "kind": "account",
+                                "path": "player.authority",
+                                "account": "PlayerData"
+                            },
+                            {
+                                "kind": "account",
+                                "path": "token_program"
+                            },
+                            {
+                                "kind": "account",
+                                "path": "mint"
+                            }
+                        ],
+                        "program": {
+                            "kind": "const",
+                            "value": [
+                                140, 151, 37, 143, 78, 36, 137, 241, 187, 61, 16, 41, 20, 142, 13, 131, 11, 90, 19, 153,
+                                218, 255, 16, 132, 4, 142, 123, 216, 219, 233, 248, 89
+                            ]
+                        }
+                    }
+                },
                 {
                     "name": "nft_authority",
                     "writable": true,
@@ -218,6 +245,11 @@
             "code": 6003,
             "name": "CantInitializeMetadataPointer",
             "msg": "Cant initialize metadata_pointer"
+        },
+        {
+            "code": 6004,
+            "name": "NftNotOwned",
+            "msg": "Player does not own this NFT"
         }
     ],
     "types": [