1use dlp::args::{
2 EncryptedBuffer, MaybeEncryptedAccountMeta, MaybeEncryptedIxData,
3 MaybeEncryptedPubkey,
4};
5use solana_program::{instruction::AccountMeta, pubkey::Pubkey};
6
7use crate::{
8 encryption::EncryptionError,
9 instruction_builder::{
10 Encrypt, EncryptableAccountMeta, EncryptableIxData, EncryptablePubkey,
11 PostDelegationInstruction,
12 },
13};
14
15impl Encrypt for EncryptablePubkey {
16 type Output = MaybeEncryptedPubkey;
17 type Error = EncryptionError;
18
19 fn encrypt(self, validator: &Pubkey) -> Result<Self::Output, Self::Error> {
20 if self.is_encryptable {
21 Ok(MaybeEncryptedPubkey::Encrypted(EncryptedBuffer::new(
22 crate::encryption::encrypt_ed25519_recipient(
23 self.pubkey.as_array(),
24 validator.as_array(),
25 )?,
26 )))
27 } else {
28 Ok(MaybeEncryptedPubkey::ClearText(self.pubkey.to_bytes()))
29 }
30 }
31}
32
33impl Encrypt for EncryptableIxData {
34 type Output = MaybeEncryptedIxData;
35 type Error = EncryptionError;
36
37 fn encrypt(self, validator: &Pubkey) -> Result<Self::Output, Self::Error> {
38 if self.encrypt_begin_offset >= self.data.len() {
39 Ok(MaybeEncryptedIxData {
40 prefix: self.data,
41 suffix: EncryptedBuffer::default(),
42 })
43 } else {
44 Ok(MaybeEncryptedIxData {
45 prefix: self.data[0..self.encrypt_begin_offset].into(),
46 suffix: EncryptedBuffer::new(
47 crate::encryption::encrypt_ed25519_recipient(
48 &self.data[self.encrypt_begin_offset..],
49 validator.as_array(),
50 )?,
51 ),
52 })
53 }
54 }
55}
56
57impl Encrypt for dlp::compact::EncryptableAccountMeta {
58 type Output = MaybeEncryptedAccountMeta;
59 type Error = EncryptionError;
60
61 fn encrypt(self, validator: &Pubkey) -> Result<Self::Output, Self::Error> {
62 if self.is_encryptable {
63 Ok(MaybeEncryptedAccountMeta::Encrypted(EncryptedBuffer::new(
64 crate::encryption::encrypt_ed25519_recipient(
65 &[self.account_meta.to_byte()],
66 validator.as_array(),
67 )?,
68 )))
69 } else {
70 Ok(MaybeEncryptedAccountMeta::ClearText(self.account_meta))
71 }
72 }
73}
74
75impl Encrypt for Vec<PostDelegationInstruction> {
76 type Output = (dlp::args::PostDelegationActions, Vec<AccountMeta>);
77 type Error = EncryptionError;
78
79 fn encrypt(self, validator: &Pubkey) -> Result<Self::Output, Self::Error> {
80 use dlp::args::MaybeEncryptedInstruction;
81
82 let mut signers: Vec<AccountMeta> = Vec::new();
83
84 let mut add_to_signers = |meta: &EncryptableAccountMeta| {
85 assert!(
86 meta.account_meta.is_signer,
87 "AccountMeta must be a signer"
88 );
89 assert!(!meta.is_encryptable, "signer must not be encryptable");
90 let Some(found) = signers
91 .iter_mut()
92 .find(|m| m.pubkey == meta.account_meta.pubkey)
93 else {
94 signers.push(meta.account_meta.clone());
95 return;
96 };
97
98 found.is_signer |= meta.account_meta.is_signer;
99 found.is_writable |= meta.account_meta.is_writable;
100 };
101
102 let mut non_signers: Vec<EncryptableAccountMeta> = Vec::new();
103 let mut add_to_non_signers = |meta: &EncryptableAccountMeta| {
104 assert!(
105 !meta.account_meta.is_signer,
106 "AccountMeta must not be a signer"
107 );
108 let Some(found) = non_signers
109 .iter_mut()
110 .find(|m| m.account_meta.pubkey == meta.account_meta.pubkey)
111 else {
112 non_signers.push(meta.clone());
113 return;
114 };
115
116 found.is_encryptable |= meta.is_encryptable;
117 found.account_meta.is_writable |= meta.account_meta.is_writable;
118 };
119
120 for meta in self
121 .iter()
122 .flat_map(|ix| ix.accounts.iter())
123 .filter(|meta| meta.account_meta.is_signer)
124 {
125 add_to_signers(meta);
126 }
127
128 for ix in self.iter() {
129 add_to_non_signers(&EncryptableAccountMeta {
130 account_meta: AccountMeta::new_readonly(
131 ix.program_id.pubkey,
132 false,
133 ),
134 is_encryptable: ix.program_id.is_encryptable,
135 });
136 for meta in ix
137 .accounts
138 .iter()
139 .filter(|meta| !meta.account_meta.is_signer)
140 {
141 let Some(found) = signers
142 .iter_mut()
143 .find(|m| m.pubkey == meta.account_meta.pubkey)
144 else {
145 add_to_non_signers(meta);
146 continue;
147 };
148
149 found.is_writable |= meta.account_meta.is_writable;
150 }
151 }
152
153 if signers.len() + non_signers.len()
154 > dlp::compact::MAX_PUBKEYS as usize
155 {
156 panic!(
157 "delegate_with_actions supports at most {} unique pubkeys",
158 dlp::compact::MAX_PUBKEYS
159 );
160 }
161
162 let index_of = |pk: &Pubkey| -> u8 {
163 if let Some(index) = signers.iter().position(|s| &s.pubkey == pk) {
164 return index as u8;
165 }
166 signers.len() as u8
167 + non_signers
168 .iter()
169 .position(|ns| &ns.account_meta.pubkey == pk)
170 .expect("pubkey must exist in signers or non_signers")
171 as u8
172 };
173
174 let compact_instructions: Vec<MaybeEncryptedInstruction> = self
175 .into_iter()
176 .map(|ix| MaybeEncryptedInstruction {
177 program_id: index_of(&ix.program_id.pubkey),
178
179 accounts: ix
180 .accounts
181 .into_iter()
182 .map(|meta| {
183 let index = index_of(&meta.account_meta.pubkey);
184 meta.to_compact(index)
185 .encrypt(validator)
186 .expect("account metadata encryption failed")
187 })
188 .collect(),
189
190 data: ix
191 .data
192 .encrypt(validator)
193 .expect("instruction data encryption failed"),
194 })
195 .collect();
196
197 Ok((
198 dlp::args::PostDelegationActions {
199 inserted_signers: 0,
200 inserted_non_signers: 0,
201 signers: signers.iter().map(|s| s.pubkey.to_bytes()).collect(),
202
203 non_signers: non_signers
204 .into_iter()
205 .map(|ns| {
206 EncryptablePubkey {
207 pubkey: ns.account_meta.pubkey,
208 is_encryptable: ns.is_encryptable,
209 }
210 .encrypt(validator)
211 .expect("pubkey encryption failed")
212 })
213 .collect(),
214
215 instructions: compact_instructions,
216 },
217 signers,
218 ))
219 }
220}