kora_lib/lighthouse/
assertion.rs

1use solana_client::nonblocking::rpc_client::RpcClient;
2use solana_message::{compiled_instruction::CompiledInstruction, MessageHeader, VersionedMessage};
3use solana_sdk::{
4    instruction::{AccountMeta, Instruction},
5    pubkey::Pubkey,
6    transaction::VersionedTransaction,
7};
8
9use crate::{
10    config::LighthouseConfig,
11    constant::{LIGHTHOUSE_PROGRAM_ID, MAX_TRANSACTION_SIZE},
12    error::KoraError,
13    sanitize_error,
14};
15
16/// Lighthouse instruction discriminators
17const ASSERT_ACCOUNT_INFO_DISCRIMINATOR: u8 = 5;
18
19/// LogLevel::Silent value
20const LOG_LEVEL_SILENT: u8 = 0;
21
22/// IntegerOperator::GreaterThanOrEqual value (matches Lighthouse SDK)
23const INTEGER_OPERATOR_GTE: u8 = 4;
24
25/// AccountInfoAssertion::Lamports variant (index 0 in the enum)
26const ACCOUNT_INFO_ASSERTION_LAMPORTS: u8 = 0;
27
28pub struct LighthouseUtil {}
29
30impl LighthouseUtil {
31    /// Add a fee payer balance assertion to a transaction if lighthouse is enabled and not sending.
32    /// Asserts that fee payer balance >= (current_balance - estimated_fee) at transaction end.
33    ///
34    /// The `will_send` parameter indicates if the transaction will be sent to the network directly.
35    /// When `will_send` is true, the assertion is skipped because modifying the message would
36    /// invalidate existing client signatures.
37    pub async fn add_fee_payer_assertion(
38        transaction: &mut VersionedTransaction,
39        rpc_client: &RpcClient,
40        fee_payer: &Pubkey,
41        estimated_fee: u64,
42        config: &LighthouseConfig,
43        will_send: bool,
44    ) -> Result<(), KoraError> {
45        if !config.enabled || will_send {
46            return Ok(());
47        }
48
49        let current_balance = rpc_client.get_balance(fee_payer).await.map_err(|e| {
50            KoraError::RpcError(format!(
51                "Failed to fetch fee payer balance for Lighthouse assertion: {}",
52                sanitize_error!(e)
53            ))
54        })?;
55        let min_expected = current_balance.saturating_sub(estimated_fee);
56
57        if min_expected == 0 {
58            log::warn!(
59                "Fee payer {} has balance {} which may be insufficient for estimated fee {}",
60                fee_payer,
61                current_balance,
62                estimated_fee
63            );
64        }
65
66        let assertion_ix = Self::build_fee_payer_assertion(fee_payer, min_expected);
67        Self::append_lighthouse_assertion(transaction, assertion_ix, config)
68    }
69
70    /// Build instruction data for AssertAccountInfo with Lamports assertion
71    fn build_assert_account_info_data(min_lamports: u64) -> Vec<u8> {
72        let mut data = Vec::with_capacity(12);
73
74        // Instruction discriminator
75        data.push(ASSERT_ACCOUNT_INFO_DISCRIMINATOR);
76
77        // LogLevel::Silent
78        data.push(LOG_LEVEL_SILENT);
79
80        // AccountInfoAssertion::Lamports variant
81        data.push(ACCOUNT_INFO_ASSERTION_LAMPORTS);
82
83        // Lamports value (u64 little-endian)
84        data.extend_from_slice(&min_lamports.to_le_bytes());
85
86        // IntegerOperator::GreaterThanOrEqual
87        data.push(INTEGER_OPERATOR_GTE);
88
89        data
90    }
91
92    /// Build a Lighthouse assertion instruction that asserts the fee payer's balance
93    /// is >= min_lamports at the end of the transaction.
94    fn build_fee_payer_assertion(fee_payer: &Pubkey, min_lamports: u64) -> Instruction {
95        let data = Self::build_assert_account_info_data(min_lamports);
96
97        Instruction {
98            program_id: LIGHTHOUSE_PROGRAM_ID,
99            accounts: vec![AccountMeta::new_readonly(*fee_payer, false)],
100            data,
101        }
102    }
103
104    /// Find an account in the account keys list or add it
105    fn find_or_add_account(
106        account_keys: &mut Vec<Pubkey>,
107        pubkey: &Pubkey,
108    ) -> Result<(u8, bool), KoraError> {
109        if let Some(index) = account_keys.iter().position(|k| k == pubkey) {
110            Ok((index as u8, false))
111        } else {
112            if account_keys.len() >= 256 {
113                return Err(KoraError::ValidationError(
114                    "Transaction has too many accounts (max 256)".to_string(),
115                ));
116            }
117            let index = account_keys.len() as u8;
118            account_keys.push(*pubkey);
119            Ok((index, true))
120        }
121    }
122
123    fn increment_readonly_unsigned_accounts(header: &mut MessageHeader) -> Result<(), KoraError> {
124        header.num_readonly_unsigned_accounts =
125            header.num_readonly_unsigned_accounts.checked_add(1).ok_or_else(|| {
126                KoraError::ValidationError(
127                    "num_readonly_unsigned_accounts overflow when appending instruction"
128                        .to_string(),
129                )
130            })?;
131        Ok(())
132    }
133
134    /// Append an instruction to a versioned transaction
135    fn append_instruction_to_transaction(
136        transaction: &mut VersionedTransaction,
137        instruction: Instruction,
138    ) -> Result<(), KoraError> {
139        match &mut transaction.message {
140            VersionedMessage::Legacy(message) => {
141                let (program_id_index, program_added) =
142                    Self::find_or_add_account(&mut message.account_keys, &instruction.program_id)?;
143                if program_added {
144                    Self::increment_readonly_unsigned_accounts(&mut message.header)?;
145                }
146
147                let mut account_indices: Vec<u8> = Vec::with_capacity(instruction.accounts.len());
148                for meta in &instruction.accounts {
149                    let (index, added) =
150                        Self::find_or_add_account(&mut message.account_keys, &meta.pubkey)?;
151                    if added {
152                        if meta.is_signer || meta.is_writable {
153                            return Err(KoraError::ValidationError(
154                                "Appending new signer/writable accounts is not supported"
155                                    .to_string(),
156                            ));
157                        }
158                        Self::increment_readonly_unsigned_accounts(&mut message.header)?;
159                    }
160                    account_indices.push(index);
161                }
162
163                message.instructions.push(CompiledInstruction {
164                    program_id_index,
165                    accounts: account_indices,
166                    data: instruction.data,
167                });
168
169                Ok(())
170            }
171            VersionedMessage::V0(message) => {
172                let (program_id_index, program_added) =
173                    Self::find_or_add_account(&mut message.account_keys, &instruction.program_id)?;
174                if program_added {
175                    Self::increment_readonly_unsigned_accounts(&mut message.header)?;
176                }
177
178                let mut account_indices: Vec<u8> = Vec::with_capacity(instruction.accounts.len());
179                for meta in &instruction.accounts {
180                    let (index, added) =
181                        Self::find_or_add_account(&mut message.account_keys, &meta.pubkey)?;
182                    if added {
183                        if meta.is_signer || meta.is_writable {
184                            return Err(KoraError::ValidationError(
185                                "Appending new signer/writable accounts is not supported"
186                                    .to_string(),
187                            ));
188                        }
189                        Self::increment_readonly_unsigned_accounts(&mut message.header)?;
190                    }
191                    account_indices.push(index);
192                }
193
194                message.instructions.push(CompiledInstruction {
195                    program_id_index,
196                    accounts: account_indices,
197                    data: instruction.data,
198                });
199
200                Ok(())
201            }
202        }
203    }
204
205    /// Append a Lighthouse assertion instruction to a transaction.
206    /// Handles size overflow based on config settings.
207    pub(crate) fn append_lighthouse_assertion(
208        transaction: &mut VersionedTransaction,
209        assertion_ix: Instruction,
210        config: &LighthouseConfig,
211    ) -> Result<(), KoraError> {
212        // Clone and append to get actual size
213        let mut tx_with_assertion = transaction.clone();
214        Self::append_instruction_to_transaction(&mut tx_with_assertion, assertion_ix)?;
215
216        let new_size = bincode::serialize(&tx_with_assertion)
217            .map_err(|e| {
218                KoraError::SerializationError(sanitize_error!(format!(
219                    "Failed to serialize transaction: {e}"
220                )))
221            })?
222            .len();
223
224        if new_size > MAX_TRANSACTION_SIZE {
225            if config.fail_if_transaction_size_overflow {
226                return Err(KoraError::ValidationError(format!(
227                    "Adding Lighthouse assertion would exceed transaction size limit ({} > {})",
228                    new_size, MAX_TRANSACTION_SIZE
229                )));
230            } else {
231                log::warn!(
232                    "Lighthouse assertion would exceed transaction size limit ({} > {}). Skipping.",
233                    new_size,
234                    MAX_TRANSACTION_SIZE
235                );
236                return Ok(());
237            }
238        }
239
240        // Commit the change
241        *transaction = tx_with_assertion;
242        Ok(())
243    }
244}
245
246#[cfg(test)]
247mod tests {
248    use super::*;
249    use solana_message::{v0, Message, VersionedMessage};
250    use solana_sdk::{hash::Hash, instruction::AccountMeta, signature::Keypair, signer::Signer};
251
252    #[test]
253    fn test_build_assert_account_info_data() {
254        let data = LighthouseUtil::build_assert_account_info_data(1_000_000);
255
256        // Verify structure: discriminator(1) + log_level(1) + variant(1) + u64(8) + operator(1) = 12 bytes
257        assert_eq!(data.len(), 12);
258        assert_eq!(data[0], 5); // ASSERT_ACCOUNT_INFO_DISCRIMINATOR
259        assert_eq!(data[1], 0); // LogLevel::Silent
260        assert_eq!(data[2], 0); // ACCOUNT_INFO_ASSERTION_LAMPORTS
261                                // Bytes 3-10: u64 little-endian (1_000_000 = 0x000F4240)
262        assert_eq!(u64::from_le_bytes(data[3..11].try_into().unwrap()), 1_000_000);
263        assert_eq!(data[11], 4); // IntegerOperator::GreaterThanOrEqual
264    }
265
266    #[test]
267    fn test_build_fee_payer_assertion() {
268        let fee_payer = Keypair::new().pubkey();
269        let min_lamports = 1_000_000;
270
271        let ix = LighthouseUtil::build_fee_payer_assertion(&fee_payer, min_lamports);
272
273        assert_eq!(ix.data.len(), 12);
274        assert_eq!(ix.accounts.len(), 1);
275        assert_eq!(ix.accounts[0].pubkey, fee_payer);
276        assert!(!ix.accounts[0].is_signer);
277        assert!(!ix.accounts[0].is_writable);
278    }
279
280    #[test]
281    fn test_append_lighthouse_assertion_legacy() {
282        let keypair = Keypair::new();
283        let program_id = Pubkey::new_unique();
284
285        let instruction = Instruction::new_with_bytes(
286            program_id,
287            &[1, 2, 3],
288            vec![AccountMeta::new(keypair.pubkey(), true)],
289        );
290
291        let message =
292            VersionedMessage::Legacy(Message::new(&[instruction], Some(&keypair.pubkey())));
293        let mut transaction = VersionedTransaction::try_new(message, &[&keypair]).unwrap();
294
295        let original_ix_count = transaction.message.instructions().len();
296        let original_readonly_unsigned =
297            transaction.message.header().num_readonly_unsigned_accounts;
298
299        let assertion_ix = LighthouseUtil::build_fee_payer_assertion(&keypair.pubkey(), 1_000_000);
300        let config = LighthouseConfig { enabled: true, fail_if_transaction_size_overflow: true };
301
302        let result =
303            LighthouseUtil::append_lighthouse_assertion(&mut transaction, assertion_ix, &config);
304        assert!(result.is_ok());
305
306        assert_eq!(transaction.message.instructions().len(), original_ix_count + 1);
307        assert_eq!(
308            transaction.message.header().num_readonly_unsigned_accounts,
309            original_readonly_unsigned + 1
310        );
311        assert!(transaction.message.static_account_keys().contains(&LIGHTHOUSE_PROGRAM_ID));
312    }
313
314    #[test]
315    fn test_append_lighthouse_assertion_v0() {
316        let keypair = Keypair::new();
317        let program_id = Pubkey::new_unique();
318
319        let v0_message = v0::Message {
320            header: solana_message::MessageHeader {
321                num_required_signatures: 1,
322                num_readonly_signed_accounts: 0,
323                num_readonly_unsigned_accounts: 1,
324            },
325            account_keys: vec![keypair.pubkey(), program_id],
326            recent_blockhash: Hash::new_unique(),
327            instructions: vec![solana_message::compiled_instruction::CompiledInstruction {
328                program_id_index: 1,
329                accounts: vec![0],
330                data: vec![1, 2, 3],
331            }],
332            address_table_lookups: vec![],
333        };
334
335        let message = VersionedMessage::V0(v0_message);
336        let mut transaction = VersionedTransaction::try_new(message, &[&keypair]).unwrap();
337
338        let original_ix_count = transaction.message.instructions().len();
339        let original_readonly_unsigned =
340            transaction.message.header().num_readonly_unsigned_accounts;
341
342        let assertion_ix = LighthouseUtil::build_fee_payer_assertion(&keypair.pubkey(), 1_000_000);
343        let config = LighthouseConfig { enabled: true, fail_if_transaction_size_overflow: true };
344
345        let result =
346            LighthouseUtil::append_lighthouse_assertion(&mut transaction, assertion_ix, &config);
347        assert!(result.is_ok());
348
349        assert_eq!(transaction.message.instructions().len(), original_ix_count + 1);
350        assert_eq!(
351            transaction.message.header().num_readonly_unsigned_accounts,
352            original_readonly_unsigned + 1
353        );
354        assert!(transaction.message.static_account_keys().contains(&LIGHTHOUSE_PROGRAM_ID));
355    }
356
357    #[test]
358    fn test_append_lighthouse_assertion_header_unchanged_when_lighthouse_program_exists() {
359        let keypair = Keypair::new();
360
361        let instruction = Instruction::new_with_bytes(
362            LIGHTHOUSE_PROGRAM_ID,
363            &[1, 2, 3],
364            vec![AccountMeta::new(keypair.pubkey(), true)],
365        );
366
367        let message =
368            VersionedMessage::Legacy(Message::new(&[instruction], Some(&keypair.pubkey())));
369        let mut transaction = VersionedTransaction::try_new(message, &[&keypair]).unwrap();
370        let original_readonly_unsigned =
371            transaction.message.header().num_readonly_unsigned_accounts;
372
373        let assertion_ix = LighthouseUtil::build_fee_payer_assertion(&keypair.pubkey(), 1_000_000);
374        let config = LighthouseConfig { enabled: true, fail_if_transaction_size_overflow: true };
375
376        let result =
377            LighthouseUtil::append_lighthouse_assertion(&mut transaction, assertion_ix, &config);
378        assert!(result.is_ok());
379        assert_eq!(
380            transaction.message.header().num_readonly_unsigned_accounts,
381            original_readonly_unsigned
382        );
383    }
384
385    #[test]
386    fn test_overflow_skip_behavior() {
387        let keypair = Keypair::new();
388        let program_id = Pubkey::new_unique();
389
390        let large_data = vec![0u8; 1100];
391        let instruction = Instruction::new_with_bytes(
392            program_id,
393            &large_data,
394            vec![AccountMeta::new(keypair.pubkey(), true)],
395        );
396
397        let message =
398            VersionedMessage::Legacy(Message::new(&[instruction], Some(&keypair.pubkey())));
399        let mut transaction = VersionedTransaction::try_new(message, &[&keypair]).unwrap();
400
401        let original_ix_count = transaction.message.instructions().len();
402
403        let assertion_ix = LighthouseUtil::build_fee_payer_assertion(&keypair.pubkey(), 1_000_000);
404        let config = LighthouseConfig { enabled: true, fail_if_transaction_size_overflow: false };
405
406        let result =
407            LighthouseUtil::append_lighthouse_assertion(&mut transaction, assertion_ix, &config);
408        assert!(result.is_ok());
409
410        assert_eq!(transaction.message.instructions().len(), original_ix_count);
411    }
412
413    #[test]
414    fn test_overflow_fail_behavior() {
415        let keypair = Keypair::new();
416        let program_id = Pubkey::new_unique();
417
418        let large_data = vec![0u8; 1100];
419        let instruction = Instruction::new_with_bytes(
420            program_id,
421            &large_data,
422            vec![AccountMeta::new(keypair.pubkey(), true)],
423        );
424
425        let message =
426            VersionedMessage::Legacy(Message::new(&[instruction], Some(&keypair.pubkey())));
427        let mut transaction = VersionedTransaction::try_new(message, &[&keypair]).unwrap();
428
429        let assertion_ix = LighthouseUtil::build_fee_payer_assertion(&keypair.pubkey(), 1_000_000);
430        let config = LighthouseConfig { enabled: true, fail_if_transaction_size_overflow: true };
431
432        let result =
433            LighthouseUtil::append_lighthouse_assertion(&mut transaction, assertion_ix, &config);
434        assert!(result.is_err());
435
436        if let Err(KoraError::ValidationError(msg)) = result {
437            assert!(msg.contains("exceed transaction size limit"));
438        } else {
439            panic!("Expected ValidationError");
440        }
441    }
442}
kora_lib/lighthouse/assertion.rs

kora_lib/lighthouse/
assertion.rs
