1use dlp::{
2 args::{
3 MaybeEncryptedAccountMeta, MaybeEncryptedIxData, MaybeEncryptedPubkey,
4 PostDelegationActions,
5 },
6 compact,
7};
8use solana_sdk::{
9 instruction::{AccountMeta, Instruction},
10 signer::Signer,
11};
12use thiserror::Error;
13
14use crate::encryption::{self, EncryptionError, KEY_LEN};
15
16#[derive(Debug, Error)]
17pub enum DecryptError {
18 #[error(transparent)]
19 DecryptFailed(#[from] EncryptionError),
20 #[error("invalid decrypted pubkey length: {0}")]
21 InvalidPubkeyLength(usize),
22 #[error("invalid decrypted compact account meta length: {0}")]
23 InvalidAccountMetaLength(usize),
24 #[error("invalid decrypted compact account meta value: {0}")]
25 InvalidAccountMetaValue(u8),
26 #[error("invalid program_id index {index} for pubkey table len {len}")]
27 InvalidProgramIdIndex { index: u8, len: usize },
28 #[error("invalid account index {index} for pubkey table len {len}")]
29 InvalidAccountIndex { index: u8, len: usize },
30}
31
32pub trait Decrypt: Sized {
33 type Output;
34
35 fn decrypt(
36 self,
37 recipient_x25519_pubkey: &[u8; KEY_LEN],
38 recipient_x25519_secret: &[u8; KEY_LEN],
39 ) -> Result<Self::Output, DecryptError>;
40
41 fn decrypt_with_keypair(
42 self,
43 recipient_keypair: &solana_sdk::signature::Keypair,
44 ) -> Result<Self::Output, DecryptError>
45 where
46 Self: Sized,
47 {
48 let recipient_x25519_secret =
49 encryption::keypair_to_x25519_secret(recipient_keypair)?;
50 let recipient_x25519_pubkey = encryption::ed25519_pubkey_to_x25519(
51 recipient_keypair.pubkey().as_array(),
52 )?;
53 self.decrypt(&recipient_x25519_pubkey, &recipient_x25519_secret)
54 }
55}
56
57impl Decrypt for MaybeEncryptedPubkey {
58 type Output = [u8; 32];
59
60 fn decrypt(
61 self,
62 recipient_x25519_pubkey: &[u8; KEY_LEN],
63 recipient_x25519_secret: &[u8; KEY_LEN],
64 ) -> Result<Self::Output, DecryptError> {
65 match self {
66 Self::ClearText(pubkey) => Ok(pubkey),
67 Self::Encrypted(buffer) => {
68 let plaintext = encryption::decrypt(
69 buffer.as_bytes(),
70 recipient_x25519_pubkey,
71 recipient_x25519_secret,
72 )
73 .map_err(DecryptError::DecryptFailed)?;
74 Self::Output::try_from(plaintext.as_slice()).map_err(|_| {
75 DecryptError::InvalidPubkeyLength(plaintext.len())
76 })
77 }
78 }
79 }
80}
81
82impl Decrypt for MaybeEncryptedAccountMeta {
83 type Output = compact::AccountMeta;
84
85 fn decrypt(
86 self,
87 recipient_x25519_pubkey: &[u8; KEY_LEN],
88 recipient_x25519_secret: &[u8; KEY_LEN],
89 ) -> Result<Self::Output, DecryptError> {
90 match self {
91 Self::ClearText(account_meta) => Ok(account_meta),
92 Self::Encrypted(buffer) => {
93 let plaintext = encryption::decrypt(
94 buffer.as_bytes(),
95 recipient_x25519_pubkey,
96 recipient_x25519_secret,
97 )
98 .map_err(DecryptError::DecryptFailed)?;
99 if plaintext.len() != 1 {
100 return Err(DecryptError::InvalidAccountMetaLength(
101 plaintext.len(),
102 ));
103 }
104 compact::AccountMeta::from_byte(plaintext[0])
105 .ok_or(DecryptError::InvalidAccountMetaValue(plaintext[0]))
106 }
107 }
108 }
109}
110
111impl Decrypt for MaybeEncryptedIxData {
112 type Output = Vec<u8>;
113
114 fn decrypt(
115 self,
116 recipient_x25519_pubkey: &[u8; KEY_LEN],
117 recipient_x25519_secret: &[u8; KEY_LEN],
118 ) -> Result<Self::Output, DecryptError> {
119 let mut data = self.prefix;
120 if !self.suffix.as_bytes().is_empty() {
121 let suffix = encryption::decrypt(
122 self.suffix.as_bytes(),
123 recipient_x25519_pubkey,
124 recipient_x25519_secret,
125 )
126 .map_err(DecryptError::DecryptFailed)?;
127 data.extend_from_slice(&suffix);
128 }
129 Ok(data)
130 }
131}
132
133impl Decrypt for PostDelegationActions {
134 type Output = Vec<Instruction>;
135
136 fn decrypt(
137 self,
138 recipient_x25519_pubkey: &[u8; KEY_LEN],
139 recipient_x25519_secret: &[u8; KEY_LEN],
140 ) -> Result<Self::Output, DecryptError> {
141 let actions = self;
142
143 let pubkeys = {
144 let mut pubkeys = actions.signers;
145 for non_signer in actions.non_signers {
146 pubkeys.push(non_signer.decrypt(
147 recipient_x25519_pubkey,
148 recipient_x25519_secret,
149 )?);
150 }
151 pubkeys
152 };
153
154 let instructions = actions
155 .instructions
156 .into_iter()
157 .map(|ix| {
158 Ok(Instruction {
159 program_id: pubkeys
160 .get(ix.program_id as usize)
161 .copied()
162 .ok_or(DecryptError::InvalidProgramIdIndex {
163 index: ix.program_id,
164 len: pubkeys.len(),
165 })?
166 .into(),
167
168 accounts: ix
169 .accounts
170 .into_iter()
171 .map(|maybe_compact_meta| {
172 let compact_meta = maybe_compact_meta.decrypt(
173 recipient_x25519_pubkey,
174 recipient_x25519_secret,
175 )?;
176 let account_pubkey = pubkeys
177 .get(compact_meta.key() as usize)
178 .copied()
179 .ok_or(DecryptError::InvalidAccountIndex {
180 index: compact_meta.key(),
181 len: pubkeys.len(),
182 })?;
183
184 Ok(AccountMeta {
185 pubkey: account_pubkey.into(),
186 is_signer: compact_meta.is_signer(),
187 is_writable: compact_meta.is_writable(),
188 })
189 })
190 .collect::<Result<Vec<_>, DecryptError>>()?,
191
192 data: ix.data.decrypt(
193 recipient_x25519_pubkey,
194 recipient_x25519_secret,
195 )?,
196 })
197 })
198 .collect::<Result<Vec<_>, DecryptError>>()?;
199
200 Ok(instructions)
201 }
202}
203
204#[cfg(test)]
205mod tests {
206 use solana_program::instruction::AccountMeta;
207 use solana_sdk::{pubkey::Pubkey, signature::Keypair, signer::Signer};
208
209 use super::*;
210 use crate::instruction_builder::{
211 Encrypt, Encryptable, EncryptableFrom, PostDelegationInstruction,
212 };
213
214 #[test]
215 fn test_post_delegation_actions_decrypt_roundtrip() {
216 let validator = Keypair::new();
217 let signer = Pubkey::new_unique();
218 let nonsigner = Pubkey::new_unique();
219 let program_id = Pubkey::new_unique();
220
221 let instructions = vec![PostDelegationInstruction {
222 program_id: program_id.cleartext(),
223 accounts: vec![
224 AccountMeta::new_readonly(signer, true).cleartext(),
225 AccountMeta::new_readonly(nonsigner, false).encrypted(),
226 ],
227 data: vec![1, 2, 3, 4].encrypted_from(2),
228 }];
229
230 let (actions, signers) = instructions
231 .encrypt(&validator.pubkey())
232 .expect("post-delegation actions encryption failed");
233
234 assert_eq!(signers, vec![AccountMeta::new_readonly(signer, true)]);
235
236 let decrypted = actions.decrypt_with_keypair(&validator).unwrap();
237
238 assert_eq!(
239 decrypted,
240 vec![Instruction {
241 program_id,
242 accounts: vec![
243 AccountMeta::new_readonly(signer, true),
244 AccountMeta::new_readonly(nonsigner, false)
245 ],
246 data: vec![1, 2, 3, 4]
247 }]
248 );
249 }
250}