From d28fc35a1d3dc28dcc3723e5971fc14f636e08a4 Mon Sep 17 00:00:00 2001
From: Gabriel Hicks <60371583+gabrielhicks@users.noreply.github.com>
Date: Fri, 13 Sep 2024 14:53:29 -0500
Subject: [PATCH] feat(add): t22 support

---
 programs/org-guard/Cargo.toml                 |  2 +-
 .../initialize_proposal_by_token_v0.rs        | 39 ++++++++++++++++---
 2 files changed, 35 insertions(+), 6 deletions(-)

diff --git a/programs/org-guard/Cargo.toml b/programs/org-guard/Cargo.toml
index 642e1fa..6d6ca50 100644
--- a/programs/org-guard/Cargo.toml
+++ b/programs/org-guard/Cargo.toml
@@ -17,6 +17,6 @@ default = []
 
 [dependencies]
 anchor-lang = { git = "https://github.com/ChewingGlass/anchor", branch = "bugfix/cpi-polymorphism" }
-anchor-spl = { git = "https://github.com/ChewingGlass/anchor", branch = "bugfix/cpi-polymorphism", features = ["token"] }
+anchor-spl = { git = "https://github.com/ChewingGlass/anchor", branch = "bugfix/cpi-polymorphism", features = ["token", "token_2022"] }
 organization = { path = "../organization", features = ["no-entrypoint", "cpi"] }
 
diff --git a/programs/org-guard/src/instructions/initialize_proposal_by_token_v0.rs b/programs/org-guard/src/instructions/initialize_proposal_by_token_v0.rs
index e9a33a7..5f9937a 100644
--- a/programs/org-guard/src/instructions/initialize_proposal_by_token_v0.rs
+++ b/programs/org-guard/src/instructions/initialize_proposal_by_token_v0.rs
@@ -2,17 +2,21 @@ use crate::error::ErrorCode;
 use crate::initialize_proposal_base_v0::*;
 use crate::state::*;
 use anchor_lang::prelude::*;
-use anchor_spl::token::TokenAccount;
+use anchor_spl::token::{Token, TokenAccount as SPLTokenAccount};
+use anchor_spl::token_2022::Token2022;
+use anchor_spl::token_interface::TokenAccount as Token2022Account;
 
 #[derive(Accounts)]
 pub struct InitializeProposalByTokenV0<'info> {
     pub initialize_proposal_base: InitializeProposalBaseV0<'info>,
 
     pub proposer: Signer<'info>,
+
+    /// CHECK: The token account must be owned by the Token program or the Token 2022 program.
     #[account(
-        token::authority = proposer,
+        constraint = token_account.owner == &Token::id() || token_account.owner == &Token2022::id()
     )]
-    pub token_account: Box<Account<'info, TokenAccount>>,
+    pub token_account: AccountInfo<'info>,
 }
 
 pub fn handler<'info>(
@@ -20,13 +24,38 @@ pub fn handler<'info>(
     args: InitializeProposalArgsV0,
 ) -> Result<()> {
     let base = &ctx.accounts.initialize_proposal_base;
+    let token_account = &ctx.accounts.token_account;
 
-    assert_sufficient_weight(&base.guard.guard_type, &ctx.accounts.token_account)?;
+    if token_account.owner == &anchor_spl::token::ID {
+        let token_account = SPLTokenAccount::try_deserialize(&mut &token_account.data.borrow()[..])?;
+        assert_sufficient_weight(&base.guard.guard_type, &token_account)?;
+    } else {
+        let token_account = Token2022Account::try_deserialize(&mut &token_account.data.borrow()[..])?;
+        assert_sufficient_weight_t22(&base.guard.guard_type, &token_account)?;
+    }
 
     cpi_initialize_proposal(&base, args)
 }
 
-fn assert_sufficient_weight(guard_type: &GuardType, token: &TokenAccount) -> Result<()> {
+fn assert_sufficient_weight(guard_type: &GuardType, token: &SPLTokenAccount) -> Result<()> {
+    let config = match guard_type {
+        GuardType::MintList { guard_data } => guard_data
+            .iter()
+            .find(|config| config.address == token.mint)
+            .ok_or(ErrorCode::MintNotValid)
+            .cloned(),
+
+        _ => Err(ErrorCode::InstructionNotAllowed.into()),
+    }?;
+
+    if token.amount >= config.divisor {
+        Ok(())
+    } else {
+        Err(ErrorCode::InsufficientWeight.into())
+    }
+}
+
+fn assert_sufficient_weight_t22(guard_type: &GuardType, token: &Token2022Account) -> Result<()> {
     let config = match guard_type {
         GuardType::MintList { guard_data } => guard_data
             .iter()