From ff111054590dea885959b269c523df50e11ee6e1 Mon Sep 17 00:00:00 2001
From: none00y <mich_smyk@wp.pl>
Date: Wed, 6 Nov 2024 19:26:37 +0100
Subject: [PATCH] add filter for token 2022

---
 lib/dex-invariant/src/invariant_dex.rs | 21 ++++++++++++++++++---
 1 file changed, 18 insertions(+), 3 deletions(-)

diff --git a/lib/dex-invariant/src/invariant_dex.rs b/lib/dex-invariant/src/invariant_dex.rs
index ba3fa92..d8063d2 100644
--- a/lib/dex-invariant/src/invariant_dex.rs
+++ b/lib/dex-invariant/src/invariant_dex.rs
@@ -4,7 +4,13 @@ use std::{
 };
 
 use anchor_lang::{AnchorDeserialize, Id};
-use anchor_spl::{token::spl_token::{self, state::AccountState}, token_2022::Token2022};
+use anchor_spl::{
+    token::{
+        self,
+        spl_token::{self, state::AccountState},
+    },
+    token_2022::{self, spl_token_2022, Token2022},
+};
 use anyhow::{Context, Ok};
 use async_trait::async_trait;
 use invariant_types::{
@@ -242,11 +248,20 @@ impl DexInterface for InvariantDex {
         let vaults = rpc.get_multiple_accounts(&reserves).await?;
         let banned_reserves = vaults
             .iter()
-            .filter(|(_, reserve)| {
-                        spl_token::state::Account::unpack(reserve.data())
+            .filter(|(_, reserve)| match reserve.owner {
+                token_2022::ID => {
+                    spl_token_2022::state::Account::unpack(reserve.data())
+                        .unwrap()
+                        .state
+                        == spl_token_2022::state::AccountState::Frozen
+                }
+                token::ID => {
+                    spl_token::state::Account::unpack(reserve.data())
                         .unwrap()
                         .state
                         == AccountState::Frozen
+                }
+                _ => true,
             })
             .map(|(pk, _)| pk)
             .collect::<HashSet<_>>();