yellowstone_shield_cli/command/
identity.rs

1use super::{CommandComplete, RunCommand, RunResult, SolanaAccount};
2use crate::{command::CommandContext, policy::PolicyVersion};
3use borsh::BorshDeserialize;
4use log::info;
5use solana_sdk::commitment_config::CommitmentConfig;
6use solana_sdk::pubkey::Pubkey;
7use solana_sdk::signature::Signer;
8use spl_associated_token_account::get_associated_token_address_with_program_id;
9use spl_token_2022::{
10    extension::{BaseStateWithExtensions, PodStateWithExtensions},
11    pod::PodMint,
12};
13use spl_token_metadata_interface::{
14    borsh::BorshDeserialize as TokenBorshDeserialize, state::TokenMetadata,
15};
16use yellowstone_shield_client::{
17    accounts::{Policy, PolicyV2},
18    types::Kind,
19};
20use yellowstone_shield_client::{
21    instructions::{AddIdentityBuilder, RemoveIdentityBuilder},
22    PolicyTrait, TransactionBuilder,
23};
24
25const CHUNK_SIZE: usize = 20;
26
27/// Builder for adding a identities to a policy
28#[derive(Debug, Clone)]
29pub struct AddBatchCommandBuilder<'a> {
30    mint: Option<&'a Pubkey>,
31    identities: Option<Vec<Pubkey>>,
32}
33
34impl Default for AddBatchCommandBuilder<'_> {
35    fn default() -> Self {
36        Self::new()
37    }
38}
39
40impl<'a> AddBatchCommandBuilder<'a> {
41    /// Create a new AddCommandBuilder
42    pub fn new() -> Self {
43        Self {
44            mint: None,
45            identities: None,
46        }
47    }
48
49    /// Set the mint address
50    pub fn mint(mut self, mint: &'a Pubkey) -> Self {
51        self.mint = Some(mint);
52        self
53    }
54
55    /// Set the identities to add
56    pub fn identities(mut self, identities: Vec<Pubkey>) -> Self {
57        self.identities = Some(identities);
58        self
59    }
60}
61
62#[async_trait::async_trait]
63impl RunCommand for AddBatchCommandBuilder<'_> {
64    /// Execute the addition of a identity to the policy
65    async fn run(&mut self, context: CommandContext) -> RunResult {
66        let CommandContext { keypair, client } = context;
67
68        let mint = self.mint.expect("mint must be set");
69
70        // PDA seeds are same for both Policy and PolicyV2
71        let (address, _) = Policy::find_pda(mint);
72        let identities = self.identities.take().expect("identities must be set");
73        let token_account = get_associated_token_address_with_program_id(
74            &keypair.pubkey(),
75            mint,
76            &spl_token_2022::ID,
77        );
78
79        let account_data = client.get_account(&address).await?;
80        let account_data: &[u8] = &account_data.data;
81
82        let policy_version = Kind::try_from_slice(&[account_data[0]])?;
83
84        let current = match policy_version {
85            Kind::Policy => Policy::try_deserialize_identities(account_data),
86            Kind::PolicyV2 => PolicyV2::try_deserialize_identities(account_data),
87        }?;
88
89        let add: Vec<Pubkey> = identities
90            .into_iter()
91            .filter(|identity| !current.contains(identity))
92            .collect();
93
94        for batch in add.chunks(CHUNK_SIZE) {
95            let mut instructions = Vec::new();
96            for identity in batch {
97                let add_identity_ix = AddIdentityBuilder::new()
98                    .policy(address)
99                    .mint(*mint)
100                    .token_account(token_account)
101                    .payer(keypair.pubkey())
102                    .owner(keypair.pubkey())
103                    .identity(*identity)
104                    .instruction();
105                instructions.push(add_identity_ix);
106            }
107
108            let last_blockhash = client.get_latest_blockhash().await?;
109
110            let tx = TransactionBuilder::build()
111                .instructions(instructions)
112                .signer(&keypair)
113                .payer(&keypair.pubkey())
114                .recent_blockhash(last_blockhash)
115                .transaction();
116
117            let signature = client
118                .send_and_confirm_transaction_with_spinner_and_commitment(
119                    &tx,
120                    CommitmentConfig::confirmed(),
121                )
122                .await?;
123
124            info!("Transaction signature: {}", signature);
125        }
126
127        let account_data = client.get_account(&address).await?;
128        let account_data: &[u8] = &account_data.data;
129
130        let policy_version = Kind::try_from_slice(&[account_data[0]])?;
131
132        let policy = match policy_version {
133            Kind::Policy => PolicyVersion::V1(Policy::from_bytes(account_data)?),
134            Kind::PolicyV2 => PolicyVersion::V2(PolicyV2::from_bytes(account_data)?),
135        };
136
137        let mint_data = client.get_account(mint).await?;
138        let account_data: &[u8] = &mint_data.data;
139
140        let mint_pod = PodStateWithExtensions::<PodMint>::unpack(account_data).unwrap();
141        let mint_bytes = mint_pod.get_extension_bytes::<TokenMetadata>().unwrap();
142        let token_metadata = TokenMetadata::try_from_slice(mint_bytes).unwrap();
143
144        Ok(CommandComplete(
145            SolanaAccount(*mint, Some(token_metadata)),
146            SolanaAccount(address, Some(policy)),
147        ))
148    }
149}
150
151/// Builder for removing identities from a policy
152pub struct RemoveBatchCommandBuilder<'a> {
153    mint: Option<&'a Pubkey>,
154    identities: Option<Vec<Pubkey>>,
155}
156
157impl Default for RemoveBatchCommandBuilder<'_> {
158    fn default() -> Self {
159        Self::new()
160    }
161}
162
163impl<'a> RemoveBatchCommandBuilder<'a> {
164    /// Create a new RemoveCommandBuilder
165    pub fn new() -> Self {
166        Self {
167            mint: None,
168            identities: None,
169        }
170    }
171
172    /// Set the mint address
173    pub fn mint(mut self, mint: &'a Pubkey) -> Self {
174        self.mint = Some(mint);
175        self
176    }
177
178    /// Set the identities to remove
179    pub fn identities(mut self, identities: Vec<Pubkey>) -> Self {
180        self.identities = Some(identities);
181        self
182    }
183}
184
185#[async_trait::async_trait]
186impl RunCommand for RemoveBatchCommandBuilder<'_> {
187    /// Execute the removal of an identity from the policy
188    async fn run(&mut self, context: CommandContext) -> RunResult {
189        let CommandContext { keypair, client } = context;
190
191        let mint = self.mint.expect("mint must be set");
192        // PDA seeds are same for both Policy and PolicyV2
193        let (address, _) = Policy::find_pda(mint);
194        let identities = self.identities.take().expect("identity must be set");
195
196        let token_account = get_associated_token_address_with_program_id(
197            &keypair.pubkey(),
198            mint,
199            &spl_token_2022::ID,
200        );
201
202        let account_data = client.get_account(&address).await?;
203        let account_data: &[u8] = &account_data.data;
204
205        let policy_version = Kind::try_from_slice(&[account_data[0]])?;
206
207        let current = match policy_version {
208            Kind::Policy => Policy::try_deserialize_identities(account_data),
209            Kind::PolicyV2 => PolicyV2::try_deserialize_identities(account_data),
210        }?;
211
212        let remove: Vec<usize> = identities
213            .into_iter()
214            .filter_map(|identity| {
215                current
216                    .iter()
217                    .position(|&current_identity| current_identity == identity)
218            })
219            .collect();
220
221        for batch in remove.chunks(CHUNK_SIZE) {
222            let mut instructions = Vec::new();
223            for &index in batch {
224                let remove_identity_ix = RemoveIdentityBuilder::new()
225                    .policy(address)
226                    .mint(*mint)
227                    .token_account(token_account)
228                    .owner(keypair.pubkey())
229                    .index(index as u64)
230                    .instruction();
231                instructions.push(remove_identity_ix);
232            }
233
234            let last_blockhash = client.get_latest_blockhash().await?;
235
236            let tx = TransactionBuilder::build()
237                .instructions(instructions)
238                .signer(&keypair)
239                .payer(&keypair.pubkey())
240                .recent_blockhash(last_blockhash)
241                .transaction();
242
243            let signature = client
244                .send_and_confirm_transaction_with_spinner_and_commitment(
245                    &tx,
246                    CommitmentConfig::confirmed(),
247                )
248                .await?;
249
250            info!("Transaction signature: {}", signature);
251        }
252
253        let account_data = client.get_account(&address).await?;
254        let account_data: &[u8] = &account_data.data;
255
256        let policy_version = Kind::try_from_slice(&[account_data[0]])?;
257
258        let policy = match policy_version {
259            Kind::Policy => PolicyVersion::V1(Policy::from_bytes(&account_data[..Policy::LEN])?),
260            Kind::PolicyV2 => {
261                PolicyVersion::V2(PolicyV2::from_bytes(&account_data[..PolicyV2::LEN])?)
262            }
263        };
264
265        let mint_data = client.get_account(mint).await?;
266        let account_data: &[u8] = &mint_data.data;
267
268        let mint_pod = PodStateWithExtensions::<PodMint>::unpack(account_data).unwrap();
269        let mint_bytes = mint_pod.get_extension_bytes::<TokenMetadata>().unwrap();
270        let token_metadata = TokenMetadata::try_from_slice(mint_bytes).unwrap();
271
272        Ok(CommandComplete(
273            SolanaAccount(*mint, Some(token_metadata)),
274            SolanaAccount(address, Some(policy)),
275        ))
276    }
277}
yellowstone_shield_cli/command/identity.rs

yellowstone_shield_cli/command/
identity.rs
