From b659ba05fdce2de2ab2506121504cb6bf5503a4d Mon Sep 17 00:00:00 2001
From: Giannis Chatziveroglou <gchatz@mit.edu>
Date: Thu, 15 Dec 2022 11:09:40 -0800
Subject: [PATCH] expose reclaim funds for v2 pools

---
 components/ReclaimFunds.tsx       |  4 ++--
 handlers/useHandleReclaimFunds.ts | 39 ++++++++++++++++++++++---------
 hooks/useAllowedTokenDatas.tsx    |  1 +
 package.json                      |  2 +-
 yarn.lock                         |  8 +++----
 5 files changed, 36 insertions(+), 18 deletions(-)

diff --git a/components/ReclaimFunds.tsx b/components/ReclaimFunds.tsx
index d965caa2a..e2d8d364f 100644
--- a/components/ReclaimFunds.tsx
+++ b/components/ReclaimFunds.tsx
@@ -19,8 +19,8 @@ export const ReclaimFunds = () => {
   const handleReclaimFunds = useHandleReclaimFunds()
   if (
     !rewardDistributor.data ||
-    rewardDistributor.data.parsed?.kind !== RewardDistributorKind.Treasury ||
-    isRewardDistributorV2(rewardDistributor.data.parsed)
+    (rewardDistributor.data.parsed?.kind !== RewardDistributorKind.Treasury &&
+      !isRewardDistributorV2(rewardDistributor.data.parsed))
   ) {
     return <></>
   }
diff --git a/handlers/useHandleReclaimFunds.ts b/handlers/useHandleReclaimFunds.ts
index a0a5dd31b..f9ee80782 100644
--- a/handlers/useHandleReclaimFunds.ts
+++ b/handlers/useHandleReclaimFunds.ts
@@ -1,3 +1,5 @@
+import { findAta, withFindOrInitAssociatedTokenAccount } from '@cardinal/common'
+import { rewardsCenterProgram } from '@cardinal/rewards-center'
 import { executeTransaction } from '@cardinal/staking'
 import { withReclaimFunds } from '@cardinal/staking/dist/cjs/programs/rewardDistributor/transaction'
 import { BN } from '@project-serum/anchor'
@@ -7,6 +9,7 @@ import { notify } from 'common/Notification'
 import { asWallet } from 'common/Wallets'
 import { useRewardDistributorData } from 'hooks/useRewardDistributorData'
 import { useMutation } from 'react-query'
+import { TOKEN_PROGRAM_ID } from 'spl-token-v3'
 
 import { isStakePoolV2, useStakePoolData } from '../hooks/useStakePoolData'
 import { useEnvironmentCtx } from '../providers/EnvironmentProvider'
@@ -28,18 +31,32 @@ export const useHandleReclaimFunds = () => {
       if (!rewardDistributor.data) throw 'No reward distributor found'
 
       const transaction = new Transaction()
-      // const program = rewardsCenterProgram(connection, wallet)
       if (isStakePoolV2(stakePool.data.parsed)) {
-        // to do add instruction to lib.rs of reward center
-        // const ix = await program.methods
-        //   .reclaimFunds()
-        //   .accounts({
-        //     stakePool: stakePool.pubkey,
-        //     stakeAuthorizationRecord: stakeAuthorizationId,
-        //     authority: wallet.publicKey,
-        //   })
-        //   .instruction()
-        // transaction.add(ix)
+        const rewardDistributorTokenAccount = await findAta(
+          rewardDistributor.data.parsed.rewardMint,
+          rewardDistributor.data.pubkey,
+          true
+        )
+        const authorityTokenAccount =
+          await withFindOrInitAssociatedTokenAccount(
+            transaction,
+            connection,
+            rewardDistributor.data.parsed.rewardMint,
+            wallet.publicKey,
+            true
+          )
+        const program = rewardsCenterProgram(connection, wallet)
+        const ix = await program.methods
+          .reclaimFunds(new BN(reclaimAmount || 0))
+          .accounts({
+            rewardDistributor: rewardDistributor.data.pubkey,
+            rewardDistributorTokenAccount: rewardDistributorTokenAccount,
+            authorityTokenAccount: authorityTokenAccount,
+            authority: wallet.publicKey,
+            tokenProgram: TOKEN_PROGRAM_ID,
+          })
+          .instruction()
+        transaction.add(ix)
       } else {
         await withReclaimFunds(transaction, connection, wallet, {
           stakePoolId: stakePool.data.pubkey,
diff --git a/hooks/useAllowedTokenDatas.tsx b/hooks/useAllowedTokenDatas.tsx
index 37ba7883c..45dc023ef 100644
--- a/hooks/useAllowedTokenDatas.tsx
+++ b/hooks/useAllowedTokenDatas.tsx
@@ -40,6 +40,7 @@ export const allowedTokensForPool = (
   tokenDatas.filter((token) => {
     let isAllowed = true
     if (!stakePool.parsed) throw 'Stake pool data are unknown'
+    if (!token?.tokenAccount?.parsed) throw 'No token account found'
     const creatorAddresses = stakePool.parsed.allowedCreators
     const collectionAddresses = stakePool.parsed.allowedCollections
     const requiresAuthorization = stakePool.parsed.requiresAuthorization
diff --git a/package.json b/package.json
index ebffadea8..96404675e 100644
--- a/package.json
+++ b/package.json
@@ -26,7 +26,7 @@
     "@apollo/client": "^3.6.9",
     "@cardinal/common": "^4.0.1",
     "@cardinal/namespaces-components": "^4.1.31",
-    "@cardinal/rewards-center": "^2.2.1",
+    "@cardinal/rewards-center": "^2.2.2",
     "@cardinal/staking": "^1.13.1",
     "@cardinal/stats": "^1.1.0",
     "@emotion/babel-plugin": "^11.7.2",
diff --git a/yarn.lock b/yarn.lock
index 861c9482b..93b826757 100644
--- a/yarn.lock
+++ b/yarn.lock
@@ -339,10 +339,10 @@
     "@solana/spl-token" "^0.1.8"
     "@solana/web3.js" "^1.66.2"
 
-"@cardinal/rewards-center@^2.2.1":
-  version "2.2.1"
-  resolved "https://registry.yarnpkg.com/@cardinal/rewards-center/-/rewards-center-2.2.1.tgz#cc95d4f8c884a2d8c8f13874fa4908992036d8a1"
-  integrity sha512-cZ93FvFfB0fLQ40gwZLfd4PzIDN5zWyUF22+myudTN+kaC/QZpRdD1aOWMmzTmezn15I3YMlPsWlk6qo1nuc1w==
+"@cardinal/rewards-center@^2.2.2":
+  version "2.2.2"
+  resolved "https://registry.yarnpkg.com/@cardinal/rewards-center/-/rewards-center-2.2.2.tgz#5a2b8399388678ce20868e993ebbb576b99fb11e"
+  integrity sha512-fQRLu1jSCJ1pJ7SkEj8x+sLMPsCQpWi5W8OXRd8c7oRG/sFCgQaZ8sclsphriRdyWcKoKbXOg0Wgs9NktrP1lw==
   dependencies:
     "@cardinal/common" "^4.0.1"
     "@cardinal/creator-standard" "^2.1.11"