solana_perf/
sigverify.rs

1//! The `sigverify` module provides digital signature verification functions.
2//! By default, signatures are verified in parallel using all available CPU
3//! cores.  When perf-libs are available signature verification is offloaded
4//! to the GPU.
5//!
6use {
7    crate::{
8        cuda_runtime::PinnedVec,
9        packet::{Packet, PacketBatch, PacketFlags, PACKET_DATA_SIZE},
10        perf_libs,
11        recycler::Recycler,
12    },
13    rayon::{prelude::*, ThreadPool},
14    solana_hash::Hash,
15    solana_message::{MESSAGE_HEADER_LENGTH, MESSAGE_VERSION_PREFIX},
16    solana_pubkey::Pubkey,
17    solana_rayon_threadlimit::get_thread_count,
18    solana_short_vec::decode_shortu16_len,
19    solana_signature::Signature,
20    std::{convert::TryFrom, mem::size_of},
21};
22
23// Empirically derived to constrain max verify latency to ~8ms at lower packet counts
24pub const VERIFY_PACKET_CHUNK_SIZE: usize = 128;
25
26lazy_static! {
27    static ref PAR_THREAD_POOL: ThreadPool = rayon::ThreadPoolBuilder::new()
28        .num_threads(get_thread_count())
29        .thread_name(|i| format!("solSigVerify{i:02}"))
30        .build()
31        .unwrap();
32}
33
34pub type TxOffset = PinnedVec<u32>;
35
36type TxOffsets = (TxOffset, TxOffset, TxOffset, TxOffset, Vec<Vec<u32>>);
37
38#[derive(Debug, PartialEq, Eq)]
39struct PacketOffsets {
40    pub sig_len: u32,
41    pub sig_start: u32,
42    pub msg_start: u32,
43    pub pubkey_start: u32,
44    pub pubkey_len: u32,
45}
46
47impl PacketOffsets {
48    pub fn new(
49        sig_len: u32,
50        sig_start: u32,
51        msg_start: u32,
52        pubkey_start: u32,
53        pubkey_len: u32,
54    ) -> Self {
55        Self {
56            sig_len,
57            sig_start,
58            msg_start,
59            pubkey_start,
60            pubkey_len,
61        }
62    }
63}
64
65#[derive(Debug, PartialEq, Eq)]
66pub enum PacketError {
67    InvalidLen,
68    InvalidPubkeyLen,
69    InvalidShortVec,
70    InvalidSignatureLen,
71    MismatchSignatureLen,
72    PayerNotWritable,
73    InvalidProgramIdIndex,
74    InvalidProgramLen,
75    UnsupportedVersion,
76}
77
78impl std::convert::From<std::boxed::Box<bincode::ErrorKind>> for PacketError {
79    fn from(_e: std::boxed::Box<bincode::ErrorKind>) -> PacketError {
80        PacketError::InvalidShortVec
81    }
82}
83
84impl std::convert::From<std::num::TryFromIntError> for PacketError {
85    fn from(_e: std::num::TryFromIntError) -> Self {
86        Self::InvalidLen
87    }
88}
89
90pub fn init() {
91    if let Some(api) = perf_libs::api() {
92        unsafe {
93            (api.ed25519_set_verbose)(true);
94            assert!((api.ed25519_init)(), "ed25519_init() failed");
95            (api.ed25519_set_verbose)(false);
96        }
97    }
98}
99
100/// Returns true if the signatrue on the packet verifies.
101/// Caller must do packet.set_discard(true) if this returns false.
102#[must_use]
103fn verify_packet(packet: &mut Packet, reject_non_vote: bool) -> bool {
104    // If this packet was already marked as discard, drop it
105    if packet.meta().discard() {
106        return false;
107    }
108
109    let packet_offsets = get_packet_offsets(packet, 0, reject_non_vote);
110    let mut sig_start = packet_offsets.sig_start as usize;
111    let mut pubkey_start = packet_offsets.pubkey_start as usize;
112    let msg_start = packet_offsets.msg_start as usize;
113
114    if packet_offsets.sig_len == 0 {
115        return false;
116    }
117
118    if packet.meta().size <= msg_start {
119        return false;
120    }
121
122    for _ in 0..packet_offsets.sig_len {
123        let pubkey_end = pubkey_start.saturating_add(size_of::<Pubkey>());
124        let Some(sig_end) = sig_start.checked_add(size_of::<Signature>()) else {
125            return false;
126        };
127        let Some(Ok(signature)) = packet.data(sig_start..sig_end).map(Signature::try_from) else {
128            return false;
129        };
130        let Some(pubkey) = packet.data(pubkey_start..pubkey_end) else {
131            return false;
132        };
133        let Some(message) = packet.data(msg_start..) else {
134            return false;
135        };
136        if !signature.verify(pubkey, message) {
137            return false;
138        }
139        pubkey_start = pubkey_end;
140        sig_start = sig_end;
141    }
142    true
143}
144
145pub fn count_packets_in_batches(batches: &[PacketBatch]) -> usize {
146    batches.iter().map(|batch| batch.len()).sum()
147}
148
149pub fn count_valid_packets(batches: &[PacketBatch]) -> usize {
150    batches
151        .iter()
152        .map(|batch| batch.iter().filter(|p| !p.meta().discard()).count())
153        .sum()
154}
155
156pub fn count_discarded_packets(batches: &[PacketBatch]) -> usize {
157    batches
158        .iter()
159        .map(|batch| batch.iter().filter(|p| p.meta().discard()).count())
160        .sum()
161}
162
163// internal function to be unit-tested; should be used only by get_packet_offsets
164fn do_get_packet_offsets(
165    packet: &Packet,
166    current_offset: usize,
167) -> Result<PacketOffsets, PacketError> {
168    // should have at least 1 signature and sig lengths
169    let _ = 1usize
170        .checked_add(size_of::<Signature>())
171        .filter(|v| *v <= packet.meta().size)
172        .ok_or(PacketError::InvalidLen)?;
173
174    // read the length of Transaction.signatures (serialized with short_vec)
175    let (sig_len_untrusted, sig_size) = packet
176        .data(..)
177        .and_then(|bytes| decode_shortu16_len(bytes).ok())
178        .ok_or(PacketError::InvalidShortVec)?;
179    // Using msg_start_offset which is based on sig_len_untrusted introduces uncertainty.
180    // Ultimately, the actual sigverify will determine the uncertainty.
181    let msg_start_offset = sig_len_untrusted
182        .checked_mul(size_of::<Signature>())
183        .and_then(|v| v.checked_add(sig_size))
184        .ok_or(PacketError::InvalidLen)?;
185
186    // Determine the start of the message header by checking the message prefix bit.
187    let msg_header_offset = {
188        // Packet should have data for prefix bit
189        if msg_start_offset >= packet.meta().size {
190            return Err(PacketError::InvalidSignatureLen);
191        }
192
193        // next byte indicates if the transaction is versioned. If the top bit
194        // is set, the remaining bits encode a version number. If the top bit is
195        // not set, this byte is the first byte of the message header.
196        let message_prefix = *packet
197            .data(msg_start_offset)
198            .ok_or(PacketError::InvalidSignatureLen)?;
199        if message_prefix & MESSAGE_VERSION_PREFIX != 0 {
200            let version = message_prefix & !MESSAGE_VERSION_PREFIX;
201            match version {
202                0 => {
203                    // header begins immediately after prefix byte
204                    msg_start_offset
205                        .checked_add(1)
206                        .ok_or(PacketError::InvalidLen)?
207                }
208
209                // currently only v0 is supported
210                _ => return Err(PacketError::UnsupportedVersion),
211            }
212        } else {
213            msg_start_offset
214        }
215    };
216
217    let msg_header_offset_plus_one = msg_header_offset
218        .checked_add(1)
219        .ok_or(PacketError::InvalidLen)?;
220
221    // Packet should have data at least for MessageHeader and 1 byte for Message.account_keys.len
222    let _ = msg_header_offset_plus_one
223        .checked_add(MESSAGE_HEADER_LENGTH)
224        .filter(|v| *v <= packet.meta().size)
225        .ok_or(PacketError::InvalidSignatureLen)?;
226
227    // read MessageHeader.num_required_signatures (serialized with u8)
228    let sig_len_maybe_trusted = *packet
229        .data(msg_header_offset)
230        .ok_or(PacketError::InvalidSignatureLen)?;
231    let message_account_keys_len_offset = msg_header_offset
232        .checked_add(MESSAGE_HEADER_LENGTH)
233        .ok_or(PacketError::InvalidSignatureLen)?;
234
235    // This reads and compares the MessageHeader num_required_signatures and
236    // num_readonly_signed_accounts bytes. If num_required_signatures is not larger than
237    // num_readonly_signed_accounts, the first account is not debitable, and cannot be charged
238    // required transaction fees.
239    let readonly_signer_offset = msg_header_offset_plus_one;
240    if sig_len_maybe_trusted
241        <= *packet
242            .data(readonly_signer_offset)
243            .ok_or(PacketError::InvalidSignatureLen)?
244    {
245        return Err(PacketError::PayerNotWritable);
246    }
247
248    if usize::from(sig_len_maybe_trusted) != sig_len_untrusted {
249        return Err(PacketError::MismatchSignatureLen);
250    }
251
252    // read the length of Message.account_keys (serialized with short_vec)
253    let (pubkey_len, pubkey_len_size) = packet
254        .data(message_account_keys_len_offset..)
255        .and_then(|bytes| decode_shortu16_len(bytes).ok())
256        .ok_or(PacketError::InvalidShortVec)?;
257    let pubkey_start = message_account_keys_len_offset
258        .checked_add(pubkey_len_size)
259        .ok_or(PacketError::InvalidPubkeyLen)?;
260
261    let _ = pubkey_len
262        .checked_mul(size_of::<Pubkey>())
263        .and_then(|v| v.checked_add(pubkey_start))
264        .filter(|v| *v <= packet.meta().size)
265        .ok_or(PacketError::InvalidPubkeyLen)?;
266
267    if pubkey_len < sig_len_untrusted {
268        return Err(PacketError::InvalidPubkeyLen);
269    }
270
271    let sig_start = current_offset
272        .checked_add(sig_size)
273        .ok_or(PacketError::InvalidLen)?;
274    let msg_start = current_offset
275        .checked_add(msg_start_offset)
276        .ok_or(PacketError::InvalidLen)?;
277    let pubkey_start = current_offset
278        .checked_add(pubkey_start)
279        .ok_or(PacketError::InvalidLen)?;
280
281    Ok(PacketOffsets::new(
282        u32::try_from(sig_len_untrusted)?,
283        u32::try_from(sig_start)?,
284        u32::try_from(msg_start)?,
285        u32::try_from(pubkey_start)?,
286        u32::try_from(pubkey_len)?,
287    ))
288}
289
290fn get_packet_offsets(
291    packet: &mut Packet,
292    current_offset: usize,
293    reject_non_vote: bool,
294) -> PacketOffsets {
295    let unsanitized_packet_offsets = do_get_packet_offsets(packet, current_offset);
296    if let Ok(offsets) = unsanitized_packet_offsets {
297        check_for_simple_vote_transaction(packet, &offsets, current_offset).ok();
298        if !reject_non_vote || packet.meta().is_simple_vote_tx() {
299            return offsets;
300        }
301    }
302    // force sigverify to fail by returning zeros
303    PacketOffsets::new(0, 0, 0, 0, 0)
304}
305
306fn check_for_simple_vote_transaction(
307    packet: &mut Packet,
308    packet_offsets: &PacketOffsets,
309    current_offset: usize,
310) -> Result<(), PacketError> {
311    // vote could have 1 or 2 sigs; zero sig has already been excluded at
312    // do_get_packet_offsets.
313    if packet_offsets.sig_len > 2 {
314        return Err(PacketError::InvalidSignatureLen);
315    }
316
317    // simple vote should only be legacy message
318    let msg_start = (packet_offsets.msg_start as usize)
319        .checked_sub(current_offset)
320        .ok_or(PacketError::InvalidLen)?;
321    let message_prefix = *packet.data(msg_start).ok_or(PacketError::InvalidLen)?;
322    if message_prefix & MESSAGE_VERSION_PREFIX != 0 {
323        return Ok(());
324    }
325
326    let pubkey_start = (packet_offsets.pubkey_start as usize)
327        .checked_sub(current_offset)
328        .ok_or(PacketError::InvalidLen)?;
329
330    let instructions_len_offset = (packet_offsets.pubkey_len as usize)
331        .checked_mul(size_of::<Pubkey>())
332        .and_then(|v| v.checked_add(pubkey_start))
333        .and_then(|v| v.checked_add(size_of::<Hash>()))
334        .ok_or(PacketError::InvalidLen)?;
335
336    // Packet should have at least 1 more byte for instructions.len
337    let _ = instructions_len_offset
338        .checked_add(1usize)
339        .filter(|v| *v <= packet.meta().size)
340        .ok_or(PacketError::InvalidLen)?;
341
342    let (instruction_len, instruction_len_size) = packet
343        .data(instructions_len_offset..)
344        .and_then(|bytes| decode_shortu16_len(bytes).ok())
345        .ok_or(PacketError::InvalidLen)?;
346    // skip if has more than 1 instruction
347    if instruction_len != 1 {
348        return Err(PacketError::InvalidProgramLen);
349    }
350
351    let instruction_start = instructions_len_offset
352        .checked_add(instruction_len_size)
353        .ok_or(PacketError::InvalidLen)?;
354
355    // Packet should have at least 1 more byte for one instructions_program_id
356    let _ = instruction_start
357        .checked_add(1usize)
358        .filter(|v| *v <= packet.meta().size)
359        .ok_or(PacketError::InvalidLen)?;
360
361    let instruction_program_id_index: usize = usize::from(
362        *packet
363            .data(instruction_start)
364            .ok_or(PacketError::InvalidLen)?,
365    );
366
367    if instruction_program_id_index >= packet_offsets.pubkey_len as usize {
368        return Err(PacketError::InvalidProgramIdIndex);
369    }
370
371    let instruction_program_id_start = instruction_program_id_index
372        .checked_mul(size_of::<Pubkey>())
373        .and_then(|v| v.checked_add(pubkey_start))
374        .ok_or(PacketError::InvalidLen)?;
375    let instruction_program_id_end = instruction_program_id_start
376        .checked_add(size_of::<Pubkey>())
377        .ok_or(PacketError::InvalidLen)?;
378
379    if packet
380        .data(instruction_program_id_start..instruction_program_id_end)
381        .ok_or(PacketError::InvalidLen)?
382        == solana_sdk_ids::vote::id().as_ref()
383    {
384        packet.meta_mut().flags |= PacketFlags::SIMPLE_VOTE_TX;
385    }
386    Ok(())
387}
388
389pub fn generate_offsets(
390    batches: &mut [PacketBatch],
391    recycler: &Recycler<TxOffset>,
392    reject_non_vote: bool,
393) -> TxOffsets {
394    debug!("allocating..");
395    let mut signature_offsets: PinnedVec<_> = recycler.allocate("sig_offsets");
396    signature_offsets.set_pinnable();
397    let mut pubkey_offsets: PinnedVec<_> = recycler.allocate("pubkey_offsets");
398    pubkey_offsets.set_pinnable();
399    let mut msg_start_offsets: PinnedVec<_> = recycler.allocate("msg_start_offsets");
400    msg_start_offsets.set_pinnable();
401    let mut msg_sizes: PinnedVec<_> = recycler.allocate("msg_size_offsets");
402    msg_sizes.set_pinnable();
403    let mut current_offset: usize = 0;
404    let offsets = batches
405        .iter_mut()
406        .map(|batch| {
407            batch
408                .iter_mut()
409                .map(|packet| {
410                    let packet_offsets =
411                        get_packet_offsets(packet, current_offset, reject_non_vote);
412
413                    trace!("pubkey_offset: {}", packet_offsets.pubkey_start);
414
415                    let mut pubkey_offset = packet_offsets.pubkey_start;
416                    let mut sig_offset = packet_offsets.sig_start;
417                    let msg_size = current_offset.saturating_add(packet.meta().size) as u32;
418                    for _ in 0..packet_offsets.sig_len {
419                        signature_offsets.push(sig_offset);
420                        sig_offset = sig_offset.saturating_add(size_of::<Signature>() as u32);
421
422                        pubkey_offsets.push(pubkey_offset);
423                        pubkey_offset = pubkey_offset.saturating_add(size_of::<Pubkey>() as u32);
424
425                        msg_start_offsets.push(packet_offsets.msg_start);
426
427                        let msg_size = msg_size.saturating_sub(packet_offsets.msg_start);
428                        msg_sizes.push(msg_size);
429                    }
430                    current_offset = current_offset.saturating_add(size_of::<Packet>());
431                    packet_offsets.sig_len
432                })
433                .collect()
434        })
435        .collect();
436    (
437        signature_offsets,
438        pubkey_offsets,
439        msg_start_offsets,
440        msg_sizes,
441        offsets,
442    )
443}
444
445//inplace shrink a batch of packets
446pub fn shrink_batches(batches: &mut Vec<PacketBatch>) {
447    let mut valid_batch_ix = 0;
448    let mut valid_packet_ix = 0;
449    let mut last_valid_batch = 0;
450    for batch_ix in 0..batches.len() {
451        for packet_ix in 0..batches[batch_ix].len() {
452            if batches[batch_ix][packet_ix].meta().discard() {
453                continue;
454            }
455            last_valid_batch = batch_ix.saturating_add(1);
456            let mut found_spot = false;
457            while valid_batch_ix < batch_ix && !found_spot {
458                while valid_packet_ix < batches[valid_batch_ix].len() {
459                    if batches[valid_batch_ix][valid_packet_ix].meta().discard() {
460                        batches[valid_batch_ix][valid_packet_ix] =
461                            batches[batch_ix][packet_ix].clone();
462                        batches[batch_ix][packet_ix].meta_mut().set_discard(true);
463                        last_valid_batch = valid_batch_ix.saturating_add(1);
464                        found_spot = true;
465                        break;
466                    }
467                    valid_packet_ix = valid_packet_ix.saturating_add(1);
468                }
469                if valid_packet_ix >= batches[valid_batch_ix].len() {
470                    valid_packet_ix = 0;
471                    valid_batch_ix = valid_batch_ix.saturating_add(1);
472                }
473            }
474        }
475    }
476    batches.truncate(last_valid_batch);
477}
478
479pub fn ed25519_verify_cpu(batches: &mut [PacketBatch], reject_non_vote: bool, packet_count: usize) {
480    debug!("CPU ECDSA for {}", packet_count);
481    PAR_THREAD_POOL.install(|| {
482        batches.par_iter_mut().flatten().for_each(|packet| {
483            if !packet.meta().discard() && !verify_packet(packet, reject_non_vote) {
484                packet.meta_mut().set_discard(true);
485            }
486        });
487    });
488}
489
490pub fn ed25519_verify_disabled(batches: &mut [PacketBatch]) {
491    let packet_count = count_packets_in_batches(batches);
492    debug!("disabled ECDSA for {}", packet_count);
493    PAR_THREAD_POOL.install(|| {
494        batches.par_iter_mut().flatten().for_each(|packet| {
495            packet.meta_mut().set_discard(false);
496        });
497    });
498}
499
500pub fn copy_return_values<I, T>(sig_lens: I, out: &PinnedVec<u8>, rvs: &mut [Vec<u8>])
501where
502    I: IntoIterator<Item = T>,
503    T: IntoIterator<Item = u32>,
504{
505    debug_assert!(rvs.iter().flatten().all(|&rv| rv == 0u8));
506    let mut offset = 0usize;
507    let rvs = rvs.iter_mut().flatten();
508    for (k, rv) in sig_lens.into_iter().flatten().zip(rvs) {
509        let out = out[offset..].iter().take(k as usize).all(|&x| x == 1u8);
510        *rv = u8::from(k != 0u32 && out);
511        offset = offset.saturating_add(k as usize);
512    }
513}
514
515// return true for success, i.e ge unpacks and !ge.is_small_order()
516pub fn check_packed_ge_small_order(ge: &[u8; 32]) -> bool {
517    if let Some(api) = perf_libs::api() {
518        unsafe {
519            // Returns 1 == fail, 0 == success
520            let res = (api.ed25519_check_packed_ge_small_order)(ge.as_ptr());
521
522            return res == 0;
523        }
524    }
525    false
526}
527
528pub fn get_checked_scalar(scalar: &[u8; 32]) -> Result<[u8; 32], PacketError> {
529    let mut out = [0u8; 32];
530    if let Some(api) = perf_libs::api() {
531        unsafe {
532            let res = (api.ed25519_get_checked_scalar)(out.as_mut_ptr(), scalar.as_ptr());
533            if res == 0 {
534                return Ok(out);
535            } else {
536                return Err(PacketError::InvalidLen);
537            }
538        }
539    }
540    Ok(out)
541}
542
543pub fn mark_disabled(batches: &mut [PacketBatch], r: &[Vec<u8>]) {
544    for (batch, v) in batches.iter_mut().zip(r) {
545        for (pkt, f) in batch.iter_mut().zip(v) {
546            if !pkt.meta().discard() {
547                pkt.meta_mut().set_discard(*f == 0);
548            }
549        }
550    }
551}
552
553pub fn ed25519_verify(
554    batches: &mut [PacketBatch],
555    recycler: &Recycler<TxOffset>,
556    recycler_out: &Recycler<PinnedVec<u8>>,
557    reject_non_vote: bool,
558    valid_packet_count: usize,
559) {
560    let Some(api) = perf_libs::api() else {
561        return ed25519_verify_cpu(batches, reject_non_vote, valid_packet_count);
562    };
563    let total_packet_count = count_packets_in_batches(batches);
564    // micro-benchmarks show GPU time for smallest batch around 15-20ms
565    // and CPU speed for 64-128 sigverifies around 10-20ms. 64 is a nice
566    // power-of-two number around that accounting for the fact that the CPU
567    // may be busy doing other things while being a real validator
568    // TODO: dynamically adjust this crossover
569    let maybe_valid_percentage = 100usize
570        .wrapping_mul(valid_packet_count)
571        .checked_div(total_packet_count);
572    let Some(valid_percentage) = maybe_valid_percentage else {
573        return;
574    };
575    if valid_percentage < 90 || valid_packet_count < 64 {
576        ed25519_verify_cpu(batches, reject_non_vote, valid_packet_count);
577        return;
578    }
579
580    let (signature_offsets, pubkey_offsets, msg_start_offsets, msg_sizes, sig_lens) =
581        generate_offsets(batches, recycler, reject_non_vote);
582
583    debug!("CUDA ECDSA for {}", valid_packet_count);
584    debug!("allocating out..");
585    let mut out = recycler_out.allocate("out_buffer");
586    out.set_pinnable();
587    let mut elems = Vec::new();
588    let mut rvs = Vec::new();
589
590    let mut num_packets: usize = 0;
591    for batch in batches.iter() {
592        elems.push(perf_libs::Elems {
593            elems: batch.as_ptr().cast::<u8>(),
594            num: batch.len() as u32,
595        });
596        let v = vec![0u8; batch.len()];
597        rvs.push(v);
598        num_packets = num_packets.saturating_add(batch.len());
599    }
600    out.resize(signature_offsets.len(), 0);
601    trace!("Starting verify num packets: {}", num_packets);
602    trace!("elem len: {}", elems.len() as u32);
603    trace!("packet sizeof: {}", size_of::<Packet>() as u32);
604    trace!("len offset: {}", PACKET_DATA_SIZE as u32);
605    const USE_NON_DEFAULT_STREAM: u8 = 1;
606    unsafe {
607        let res = (api.ed25519_verify_many)(
608            elems.as_ptr(),
609            elems.len() as u32,
610            size_of::<Packet>() as u32,
611            num_packets as u32,
612            signature_offsets.len() as u32,
613            msg_sizes.as_ptr(),
614            pubkey_offsets.as_ptr(),
615            signature_offsets.as_ptr(),
616            msg_start_offsets.as_ptr(),
617            out.as_mut_ptr(),
618            USE_NON_DEFAULT_STREAM,
619        );
620        if res != 0 {
621            trace!("RETURN!!!: {}", res);
622        }
623    }
624    trace!("done verify");
625    copy_return_values(sig_lens, &out, &mut rvs);
626    mark_disabled(batches, &rvs);
627}
628
629#[cfg(test)]
630#[allow(clippy::arithmetic_side_effects)]
631mod tests {
632    use {
633        super::*,
634        crate::{
635            packet::{to_packet_batches, Packet, PacketBatch, PACKETS_PER_BATCH},
636            sigverify::{self, PacketOffsets},
637            test_tx::{new_test_vote_tx, test_multisig_tx, test_tx},
638        },
639        bincode::{deserialize, serialize},
640        curve25519_dalek::{edwards::CompressedEdwardsY, scalar::Scalar},
641        rand::{thread_rng, Rng},
642        solana_keypair::Keypair,
643        solana_message::{compiled_instruction::CompiledInstruction, Message, MessageHeader},
644        solana_signature::Signature,
645        solana_signer::Signer,
646        solana_transaction::Transaction,
647        std::{
648            iter::repeat_with,
649            sync::atomic::{AtomicU64, Ordering},
650        },
651    };
652
653    const SIG_OFFSET: usize = 1;
654
655    pub fn memfind<A: Eq>(a: &[A], b: &[A]) -> Option<usize> {
656        assert!(a.len() >= b.len());
657        let end = a.len() - b.len() + 1;
658        (0..end).find(|&i| a[i..i + b.len()] == b[..])
659    }
660
661    #[test]
662    fn test_copy_return_values() {
663        let mut rng = rand::thread_rng();
664        let sig_lens: Vec<Vec<u32>> = {
665            let size = rng.gen_range(0..64);
666            repeat_with(|| {
667                let size = rng.gen_range(0..16);
668                repeat_with(|| rng.gen_range(0..5)).take(size).collect()
669            })
670            .take(size)
671            .collect()
672        };
673        let out: Vec<Vec<Vec<bool>>> = sig_lens
674            .iter()
675            .map(|sig_lens| {
676                sig_lens
677                    .iter()
678                    .map(|&size| repeat_with(|| rng.gen()).take(size as usize).collect())
679                    .collect()
680            })
681            .collect();
682        let expected: Vec<Vec<u8>> = out
683            .iter()
684            .map(|out| {
685                out.iter()
686                    .map(|out| u8::from(!out.is_empty() && out.iter().all(|&k| k)))
687                    .collect()
688            })
689            .collect();
690        let out =
691            PinnedVec::<u8>::from_vec(out.into_iter().flatten().flatten().map(u8::from).collect());
692        let mut rvs: Vec<Vec<u8>> = sig_lens
693            .iter()
694            .map(|sig_lens| vec![0u8; sig_lens.len()])
695            .collect();
696        copy_return_values(sig_lens, &out, &mut rvs);
697        assert_eq!(rvs, expected);
698    }
699
700    #[test]
701    fn test_mark_disabled() {
702        let batch_size = 1;
703        let mut batch = PacketBatch::with_capacity(batch_size);
704        batch.resize(batch_size, Packet::default());
705        let mut batches: Vec<PacketBatch> = vec![batch];
706        mark_disabled(&mut batches, &[vec![0]]);
707        assert!(batches[0][0].meta().discard());
708        batches[0][0].meta_mut().set_discard(false);
709        mark_disabled(&mut batches, &[vec![1]]);
710        assert!(!batches[0][0].meta().discard());
711    }
712
713    #[test]
714    fn test_layout() {
715        let tx = test_tx();
716        let tx_bytes = serialize(&tx).unwrap();
717        let packet = serialize(&tx).unwrap();
718        assert_matches!(memfind(&packet, &tx_bytes), Some(0));
719        assert_matches!(memfind(&packet, &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]), None);
720    }
721
722    #[test]
723    fn test_system_transaction_layout() {
724        let tx = test_tx();
725        let tx_bytes = serialize(&tx).unwrap();
726        let message_data = tx.message_data();
727        let mut packet = Packet::from_data(None, tx.clone()).unwrap();
728
729        let packet_offsets = sigverify::get_packet_offsets(&mut packet, 0, false);
730
731        assert_eq!(
732            memfind(&tx_bytes, tx.signatures[0].as_ref()),
733            Some(SIG_OFFSET)
734        );
735        assert_eq!(
736            memfind(&tx_bytes, tx.message().account_keys[0].as_ref()),
737            Some(packet_offsets.pubkey_start as usize)
738        );
739        assert_eq!(
740            memfind(&tx_bytes, &message_data),
741            Some(packet_offsets.msg_start as usize)
742        );
743        assert_eq!(
744            memfind(&tx_bytes, tx.signatures[0].as_ref()),
745            Some(packet_offsets.sig_start as usize)
746        );
747        assert_eq!(packet_offsets.sig_len, 1);
748    }
749
750    fn packet_from_num_sigs(required_num_sigs: u8, actual_num_sigs: usize) -> Packet {
751        let message = Message {
752            header: MessageHeader {
753                num_required_signatures: required_num_sigs,
754                num_readonly_signed_accounts: 12,
755                num_readonly_unsigned_accounts: 11,
756            },
757            account_keys: vec![],
758            recent_blockhash: Hash::default(),
759            instructions: vec![],
760        };
761        let mut tx = Transaction::new_unsigned(message);
762        tx.signatures = vec![Signature::default(); actual_num_sigs];
763        Packet::from_data(None, tx).unwrap()
764    }
765
766    #[test]
767    fn test_untrustworthy_sigs() {
768        let required_num_sigs = 14;
769        let actual_num_sigs = 5;
770
771        let packet = packet_from_num_sigs(required_num_sigs, actual_num_sigs);
772
773        let unsanitized_packet_offsets = sigverify::do_get_packet_offsets(&packet, 0);
774
775        assert_eq!(
776            unsanitized_packet_offsets,
777            Err(PacketError::MismatchSignatureLen)
778        );
779    }
780
781    #[test]
782    fn test_small_packet() {
783        let tx = test_tx();
784        let mut packet = Packet::from_data(None, tx).unwrap();
785
786        packet.buffer_mut()[0] = 0xff;
787        packet.buffer_mut()[1] = 0xff;
788        packet.meta_mut().size = 2;
789
790        let res = sigverify::do_get_packet_offsets(&packet, 0);
791        assert_eq!(res, Err(PacketError::InvalidLen));
792    }
793
794    #[test]
795    fn test_pubkey_too_small() {
796        solana_logger::setup();
797        let mut tx = test_tx();
798        let sig = tx.signatures[0];
799        const NUM_SIG: usize = 18;
800        tx.signatures = vec![sig; NUM_SIG];
801        tx.message.account_keys = vec![];
802        tx.message.header.num_required_signatures = NUM_SIG as u8;
803        let mut packet = Packet::from_data(None, tx).unwrap();
804
805        let res = sigverify::do_get_packet_offsets(&packet, 0);
806        assert_eq!(res, Err(PacketError::InvalidPubkeyLen));
807
808        assert!(!verify_packet(&mut packet, false));
809
810        packet.meta_mut().set_discard(false);
811        let mut batches = generate_packet_batches(&packet, 1, 1);
812        ed25519_verify(&mut batches);
813        assert!(batches[0][0].meta().discard());
814    }
815
816    #[test]
817    fn test_pubkey_len() {
818        // See that the verify cannot walk off the end of the packet
819        // trying to index into the account_keys to access pubkey.
820        solana_logger::setup();
821
822        const NUM_SIG: usize = 17;
823        let keypair1 = Keypair::new();
824        let pubkey1 = keypair1.pubkey();
825        let mut message = Message::new(&[], Some(&pubkey1));
826        message.account_keys.push(pubkey1);
827        message.account_keys.push(pubkey1);
828        message.header.num_required_signatures = NUM_SIG as u8;
829        message.recent_blockhash = Hash::new_from_array(pubkey1.to_bytes());
830        let mut tx = Transaction::new_unsigned(message);
831
832        info!("message: {:?}", tx.message_data());
833        info!("tx: {:?}", tx);
834        let sig = keypair1.try_sign_message(&tx.message_data()).unwrap();
835        tx.signatures = vec![sig; NUM_SIG];
836
837        let mut packet = Packet::from_data(None, tx).unwrap();
838
839        let res = sigverify::do_get_packet_offsets(&packet, 0);
840        assert_eq!(res, Err(PacketError::InvalidPubkeyLen));
841
842        assert!(!verify_packet(&mut packet, false));
843
844        packet.meta_mut().set_discard(false);
845        let mut batches = generate_packet_batches(&packet, 1, 1);
846        ed25519_verify(&mut batches);
847        assert!(batches[0][0].meta().discard());
848    }
849
850    #[test]
851    fn test_large_sig_len() {
852        let tx = test_tx();
853        let mut packet = Packet::from_data(None, tx).unwrap();
854
855        // Make the signatures len huge
856        packet.buffer_mut()[0] = 0x7f;
857
858        let res = sigverify::do_get_packet_offsets(&packet, 0);
859        assert_eq!(res, Err(PacketError::InvalidSignatureLen));
860    }
861
862    #[test]
863    fn test_really_large_sig_len() {
864        let tx = test_tx();
865        let mut packet = Packet::from_data(None, tx).unwrap();
866
867        // Make the signatures len huge
868        packet.buffer_mut()[0] = 0xff;
869        packet.buffer_mut()[1] = 0xff;
870        packet.buffer_mut()[2] = 0xff;
871        packet.buffer_mut()[3] = 0xff;
872
873        let res = sigverify::do_get_packet_offsets(&packet, 0);
874        assert_eq!(res, Err(PacketError::InvalidShortVec));
875    }
876
877    #[test]
878    fn test_invalid_pubkey_len() {
879        let tx = test_tx();
880        let mut packet = Packet::from_data(None, tx).unwrap();
881
882        let res = sigverify::do_get_packet_offsets(&packet, 0);
883
884        // make pubkey len huge
885        packet.buffer_mut()[res.unwrap().pubkey_start as usize - 1] = 0x7f;
886
887        let res = sigverify::do_get_packet_offsets(&packet, 0);
888        assert_eq!(res, Err(PacketError::InvalidPubkeyLen));
889    }
890
891    #[test]
892    fn test_fee_payer_is_debitable() {
893        let message = Message {
894            header: MessageHeader {
895                num_required_signatures: 1,
896                num_readonly_signed_accounts: 1,
897                num_readonly_unsigned_accounts: 1,
898            },
899            account_keys: vec![],
900            recent_blockhash: Hash::default(),
901            instructions: vec![],
902        };
903        let mut tx = Transaction::new_unsigned(message);
904        tx.signatures = vec![Signature::default()];
905        let packet = Packet::from_data(None, tx).unwrap();
906        let res = sigverify::do_get_packet_offsets(&packet, 0);
907
908        assert_eq!(res, Err(PacketError::PayerNotWritable));
909    }
910
911    #[test]
912    fn test_unsupported_version() {
913        let tx = test_tx();
914        let mut packet = Packet::from_data(None, tx).unwrap();
915
916        let res = sigverify::do_get_packet_offsets(&packet, 0);
917
918        // set message version to 1
919        packet.buffer_mut()[res.unwrap().msg_start as usize] = MESSAGE_VERSION_PREFIX + 1;
920
921        let res = sigverify::do_get_packet_offsets(&packet, 0);
922        assert_eq!(res, Err(PacketError::UnsupportedVersion));
923    }
924
925    #[test]
926    fn test_versioned_message() {
927        let tx = test_tx();
928        let mut packet = Packet::from_data(None, tx).unwrap();
929
930        let mut legacy_offsets = sigverify::do_get_packet_offsets(&packet, 0).unwrap();
931
932        // set message version to 0
933        let msg_start = legacy_offsets.msg_start as usize;
934        let msg_bytes = packet.data(msg_start..).unwrap().to_vec();
935        packet.buffer_mut()[msg_start] = MESSAGE_VERSION_PREFIX;
936        packet.meta_mut().size += 1;
937        let msg_end = packet.meta().size;
938        packet.buffer_mut()[msg_start + 1..msg_end].copy_from_slice(&msg_bytes);
939
940        let offsets = sigverify::do_get_packet_offsets(&packet, 0).unwrap();
941        let expected_offsets = {
942            legacy_offsets.pubkey_start += 1;
943            legacy_offsets
944        };
945
946        assert_eq!(expected_offsets, offsets);
947    }
948
949    #[test]
950    fn test_system_transaction_data_layout() {
951        let mut tx0 = test_tx();
952        tx0.message.instructions[0].data = vec![1, 2, 3];
953        let message0a = tx0.message_data();
954        let tx_bytes = serialize(&tx0).unwrap();
955        assert!(tx_bytes.len() <= PACKET_DATA_SIZE);
956        assert_eq!(
957            memfind(&tx_bytes, tx0.signatures[0].as_ref()),
958            Some(SIG_OFFSET)
959        );
960        let tx1 = deserialize(&tx_bytes).unwrap();
961        assert_eq!(tx0, tx1);
962        assert_eq!(tx1.message().instructions[0].data, vec![1, 2, 3]);
963
964        tx0.message.instructions[0].data = vec![1, 2, 4];
965        let message0b = tx0.message_data();
966        assert_ne!(message0a, message0b);
967    }
968
969    // Just like get_packet_offsets, but not returning redundant information.
970    fn get_packet_offsets_from_tx(tx: Transaction, current_offset: u32) -> PacketOffsets {
971        let mut packet = Packet::from_data(None, tx).unwrap();
972        let packet_offsets =
973            sigverify::get_packet_offsets(&mut packet, current_offset as usize, false);
974        PacketOffsets::new(
975            packet_offsets.sig_len,
976            packet_offsets.sig_start - current_offset,
977            packet_offsets.msg_start - packet_offsets.sig_start,
978            packet_offsets.pubkey_start - packet_offsets.msg_start,
979            packet_offsets.pubkey_len,
980        )
981    }
982
983    #[test]
984    fn test_get_packet_offsets() {
985        assert_eq!(
986            get_packet_offsets_from_tx(test_tx(), 0),
987            PacketOffsets::new(1, 1, 64, 4, 2)
988        );
989        assert_eq!(
990            get_packet_offsets_from_tx(test_tx(), 100),
991            PacketOffsets::new(1, 1, 64, 4, 2)
992        );
993
994        // Ensure we're not indexing packet by the `current_offset` parameter.
995        assert_eq!(
996            get_packet_offsets_from_tx(test_tx(), 1_000_000),
997            PacketOffsets::new(1, 1, 64, 4, 2)
998        );
999
1000        // Ensure we're returning sig_len, not sig_size.
1001        assert_eq!(
1002            get_packet_offsets_from_tx(test_multisig_tx(), 0),
1003            PacketOffsets::new(2, 1, 128, 4, 4)
1004        );
1005    }
1006
1007    fn generate_packet_batches_random_size(
1008        packet: &Packet,
1009        max_packets_per_batch: usize,
1010        num_batches: usize,
1011    ) -> Vec<PacketBatch> {
1012        // generate packet vector
1013        let batches: Vec<_> = (0..num_batches)
1014            .map(|_| {
1015                let num_packets_per_batch = thread_rng().gen_range(1..max_packets_per_batch);
1016                let mut packet_batch = PacketBatch::with_capacity(num_packets_per_batch);
1017                for _ in 0..num_packets_per_batch {
1018                    packet_batch.push(packet.clone());
1019                }
1020                assert_eq!(packet_batch.len(), num_packets_per_batch);
1021                packet_batch
1022            })
1023            .collect();
1024        assert_eq!(batches.len(), num_batches);
1025
1026        batches
1027    }
1028
1029    fn generate_packet_batches(
1030        packet: &Packet,
1031        num_packets_per_batch: usize,
1032        num_batches: usize,
1033    ) -> Vec<PacketBatch> {
1034        // generate packet vector
1035        let batches: Vec<_> = (0..num_batches)
1036            .map(|_| {
1037                let mut packet_batch = PacketBatch::with_capacity(num_packets_per_batch);
1038                for _ in 0..num_packets_per_batch {
1039                    packet_batch.push(packet.clone());
1040                }
1041                assert_eq!(packet_batch.len(), num_packets_per_batch);
1042                packet_batch
1043            })
1044            .collect();
1045        assert_eq!(batches.len(), num_batches);
1046
1047        batches
1048    }
1049
1050    fn test_verify_n(n: usize, modify_data: bool) {
1051        let tx = test_tx();
1052        let mut packet = Packet::from_data(None, tx).unwrap();
1053
1054        // jumble some data to test failure
1055        if modify_data {
1056            packet.buffer_mut()[20] = packet.data(20).unwrap().wrapping_add(10);
1057        }
1058
1059        let mut batches = generate_packet_batches(&packet, n, 2);
1060
1061        // verify packets
1062        ed25519_verify(&mut batches);
1063
1064        // check result
1065        let should_discard = modify_data;
1066        assert!(batches
1067            .iter()
1068            .flat_map(|batch| batch.iter())
1069            .all(|p| p.meta().discard() == should_discard));
1070    }
1071
1072    fn ed25519_verify(batches: &mut [PacketBatch]) {
1073        let recycler = Recycler::default();
1074        let recycler_out = Recycler::default();
1075        let packet_count = sigverify::count_packets_in_batches(batches);
1076        sigverify::ed25519_verify(batches, &recycler, &recycler_out, false, packet_count);
1077    }
1078
1079    #[test]
1080    fn test_verify_tampered_sig_len() {
1081        let mut tx = test_tx();
1082        // pretend malicious leader dropped a signature...
1083        tx.signatures.pop();
1084        let packet = Packet::from_data(None, tx).unwrap();
1085
1086        let mut batches = generate_packet_batches(&packet, 1, 1);
1087
1088        // verify packets
1089        ed25519_verify(&mut batches);
1090        assert!(batches
1091            .iter()
1092            .flat_map(|batch| batch.iter())
1093            .all(|p| p.meta().discard()));
1094    }
1095
1096    #[test]
1097    fn test_verify_zero() {
1098        test_verify_n(0, false);
1099    }
1100
1101    #[test]
1102    fn test_verify_one() {
1103        test_verify_n(1, false);
1104    }
1105
1106    #[test]
1107    fn test_verify_seventy_one() {
1108        test_verify_n(71, false);
1109    }
1110
1111    #[test]
1112    fn test_verify_medium_pass() {
1113        test_verify_n(VERIFY_PACKET_CHUNK_SIZE, false);
1114    }
1115
1116    #[test]
1117    fn test_verify_large_pass() {
1118        test_verify_n(VERIFY_PACKET_CHUNK_SIZE * get_thread_count(), false);
1119    }
1120
1121    #[test]
1122    fn test_verify_medium_fail() {
1123        test_verify_n(VERIFY_PACKET_CHUNK_SIZE, true);
1124    }
1125
1126    #[test]
1127    fn test_verify_large_fail() {
1128        test_verify_n(VERIFY_PACKET_CHUNK_SIZE * get_thread_count(), true);
1129    }
1130
1131    #[test]
1132    fn test_verify_multisig() {
1133        solana_logger::setup();
1134
1135        let tx = test_multisig_tx();
1136        let mut packet = Packet::from_data(None, tx).unwrap();
1137
1138        let n = 4;
1139        let num_batches = 3;
1140        let mut batches = generate_packet_batches(&packet, n, num_batches);
1141
1142        packet.buffer_mut()[40] = packet.data(40).unwrap().wrapping_add(8);
1143
1144        batches[0].push(packet);
1145
1146        // verify packets
1147        ed25519_verify(&mut batches);
1148
1149        // check result
1150        let ref_ans = 1u8;
1151        let mut ref_vec = vec![vec![ref_ans; n]; num_batches];
1152        ref_vec[0].push(0u8);
1153        assert!(batches
1154            .iter()
1155            .flat_map(|batch| batch.iter())
1156            .zip(ref_vec.into_iter().flatten())
1157            .all(|(p, discard)| {
1158                if discard == 0 {
1159                    p.meta().discard()
1160                } else {
1161                    !p.meta().discard()
1162                }
1163            }));
1164    }
1165
1166    #[test]
1167    fn test_verify_fuzz() {
1168        solana_logger::setup();
1169
1170        let tx = test_multisig_tx();
1171        let packet = Packet::from_data(None, tx).unwrap();
1172
1173        let recycler = Recycler::default();
1174        let recycler_out = Recycler::default();
1175        for _ in 0..50 {
1176            let num_batches = thread_rng().gen_range(2..30);
1177            let mut batches = generate_packet_batches_random_size(&packet, 128, num_batches);
1178
1179            let num_modifications = thread_rng().gen_range(0..5);
1180            for _ in 0..num_modifications {
1181                let batch = thread_rng().gen_range(0..batches.len());
1182                let packet = thread_rng().gen_range(0..batches[batch].len());
1183                let offset = thread_rng().gen_range(0..batches[batch][packet].meta().size);
1184                let add = thread_rng().gen_range(0..255);
1185                batches[batch][packet].buffer_mut()[offset] = batches[batch][packet]
1186                    .data(offset)
1187                    .unwrap()
1188                    .wrapping_add(add);
1189            }
1190
1191            let batch_to_disable = thread_rng().gen_range(0..batches.len());
1192            for p in batches[batch_to_disable].iter_mut() {
1193                p.meta_mut().set_discard(true);
1194            }
1195
1196            // verify from GPU verification pipeline (when GPU verification is enabled) are
1197            // equivalent to the CPU verification pipeline.
1198            let mut batches_cpu = batches.clone();
1199            let packet_count = sigverify::count_packets_in_batches(&batches);
1200            sigverify::ed25519_verify(&mut batches, &recycler, &recycler_out, false, packet_count);
1201            ed25519_verify_cpu(&mut batches_cpu, false, packet_count);
1202
1203            // check result
1204            batches
1205                .iter()
1206                .flat_map(|batch| batch.iter())
1207                .zip(batches_cpu.iter().flat_map(|batch| batch.iter()))
1208                .for_each(|(p1, p2)| assert_eq!(p1, p2));
1209        }
1210    }
1211
1212    #[test]
1213    fn test_verify_fail() {
1214        test_verify_n(5, true);
1215    }
1216
1217    #[test]
1218    fn test_get_checked_scalar() {
1219        solana_logger::setup();
1220        if perf_libs::api().is_none() {
1221            return;
1222        }
1223
1224        let passed_g = AtomicU64::new(0);
1225        let failed_g = AtomicU64::new(0);
1226        (0..4).into_par_iter().for_each(|_| {
1227            let mut input = [0u8; 32];
1228            let mut passed = 0;
1229            let mut failed = 0;
1230            for _ in 0..1_000_000 {
1231                thread_rng().fill(&mut input);
1232                let ans = get_checked_scalar(&input);
1233                let ref_ans = Scalar::from_canonical_bytes(input).into_option();
1234                if let Some(ref_ans) = ref_ans {
1235                    passed += 1;
1236                    assert_eq!(ans.unwrap(), ref_ans.to_bytes());
1237                } else {
1238                    failed += 1;
1239                    assert!(ans.is_err());
1240                }
1241            }
1242            passed_g.fetch_add(passed, Ordering::Relaxed);
1243            failed_g.fetch_add(failed, Ordering::Relaxed);
1244        });
1245        info!(
1246            "passed: {} failed: {}",
1247            passed_g.load(Ordering::Relaxed),
1248            failed_g.load(Ordering::Relaxed)
1249        );
1250    }
1251
1252    #[test]
1253    fn test_ge_small_order() {
1254        solana_logger::setup();
1255        if perf_libs::api().is_none() {
1256            return;
1257        }
1258
1259        let passed_g = AtomicU64::new(0);
1260        let failed_g = AtomicU64::new(0);
1261        (0..4).into_par_iter().for_each(|_| {
1262            let mut input = [0u8; 32];
1263            let mut passed = 0;
1264            let mut failed = 0;
1265            for _ in 0..1_000_000 {
1266                thread_rng().fill(&mut input);
1267                let ans = check_packed_ge_small_order(&input);
1268                let ref_ge = CompressedEdwardsY::from_slice(&input).unwrap();
1269                if let Some(ref_element) = ref_ge.decompress() {
1270                    if ref_element.is_small_order() {
1271                        assert!(!ans);
1272                    } else {
1273                        assert!(ans);
1274                    }
1275                } else {
1276                    assert!(!ans);
1277                }
1278                if ans {
1279                    passed += 1;
1280                } else {
1281                    failed += 1;
1282                }
1283            }
1284            passed_g.fetch_add(passed, Ordering::Relaxed);
1285            failed_g.fetch_add(failed, Ordering::Relaxed);
1286        });
1287        info!(
1288            "passed: {} failed: {}",
1289            passed_g.load(Ordering::Relaxed),
1290            failed_g.load(Ordering::Relaxed)
1291        );
1292    }
1293
1294    #[test]
1295    fn test_is_simple_vote_transaction() {
1296        solana_logger::setup();
1297        let mut rng = rand::thread_rng();
1298
1299        // tansfer tx is not
1300        {
1301            let mut tx = test_tx();
1302            tx.message.instructions[0].data = vec![1, 2, 3];
1303            let mut packet = Packet::from_data(None, tx).unwrap();
1304            let packet_offsets = do_get_packet_offsets(&packet, 0).unwrap();
1305            check_for_simple_vote_transaction(&mut packet, &packet_offsets, 0).ok();
1306            assert!(!packet.meta().is_simple_vote_tx());
1307        }
1308
1309        // single legacy vote tx is
1310        {
1311            let mut tx = new_test_vote_tx(&mut rng);
1312            tx.message.instructions[0].data = vec![1, 2, 3];
1313            let mut packet = Packet::from_data(None, tx).unwrap();
1314            let packet_offsets = do_get_packet_offsets(&packet, 0).unwrap();
1315            check_for_simple_vote_transaction(&mut packet, &packet_offsets, 0).ok();
1316            assert!(packet.meta().is_simple_vote_tx());
1317        }
1318
1319        // single versioned vote tx is not
1320        {
1321            let mut tx = new_test_vote_tx(&mut rng);
1322            tx.message.instructions[0].data = vec![1, 2, 3];
1323            let mut packet = Packet::from_data(None, tx).unwrap();
1324
1325            // set messager version to v0
1326            let mut packet_offsets = do_get_packet_offsets(&packet, 0).unwrap();
1327            let msg_start = packet_offsets.msg_start as usize;
1328            let msg_bytes = packet.data(msg_start..).unwrap().to_vec();
1329            packet.buffer_mut()[msg_start] = MESSAGE_VERSION_PREFIX;
1330            packet.meta_mut().size += 1;
1331            let msg_end = packet.meta().size;
1332            packet.buffer_mut()[msg_start + 1..msg_end].copy_from_slice(&msg_bytes);
1333
1334            packet_offsets = do_get_packet_offsets(&packet, 0).unwrap();
1335            check_for_simple_vote_transaction(&mut packet, &packet_offsets, 0).ok();
1336            assert!(!packet.meta().is_simple_vote_tx());
1337        }
1338
1339        // multiple mixed tx is not
1340        {
1341            let key = Keypair::new();
1342            let key1 = Pubkey::new_unique();
1343            let key2 = Pubkey::new_unique();
1344            let tx = Transaction::new_with_compiled_instructions(
1345                &[&key],
1346                &[key1, key2],
1347                Hash::default(),
1348                vec![solana_vote_program::id(), Pubkey::new_unique()],
1349                vec![
1350                    CompiledInstruction::new(3, &(), vec![0, 1]),
1351                    CompiledInstruction::new(4, &(), vec![0, 2]),
1352                ],
1353            );
1354            let mut packet = Packet::from_data(None, tx).unwrap();
1355            let packet_offsets = do_get_packet_offsets(&packet, 0).unwrap();
1356            check_for_simple_vote_transaction(&mut packet, &packet_offsets, 0).ok();
1357            assert!(!packet.meta().is_simple_vote_tx());
1358        }
1359
1360        // single legacy vote tx with extra (invalid) signature is not
1361        {
1362            let mut tx = new_test_vote_tx(&mut rng);
1363            tx.signatures.push(Signature::default());
1364            tx.message.header.num_required_signatures = 3;
1365            tx.message.instructions[0].data = vec![1, 2, 3];
1366            let mut packet = Packet::from_data(None, tx).unwrap();
1367            let packet_offsets = do_get_packet_offsets(&packet, 0).unwrap();
1368            assert_eq!(
1369                Err(PacketError::InvalidSignatureLen),
1370                check_for_simple_vote_transaction(&mut packet, &packet_offsets, 0)
1371            );
1372            assert!(!packet.meta().is_simple_vote_tx());
1373        }
1374    }
1375
1376    #[test]
1377    fn test_is_simple_vote_transaction_with_offsets() {
1378        solana_logger::setup();
1379        let mut rng = rand::thread_rng();
1380
1381        // batch of legacy messages
1382        {
1383            let mut current_offset = 0usize;
1384            let mut batch = PacketBatch::default();
1385            batch.push(Packet::from_data(None, test_tx()).unwrap());
1386            let tx = new_test_vote_tx(&mut rng);
1387            batch.push(Packet::from_data(None, tx).unwrap());
1388            batch.iter_mut().enumerate().for_each(|(index, packet)| {
1389                let packet_offsets = do_get_packet_offsets(packet, current_offset).unwrap();
1390                check_for_simple_vote_transaction(packet, &packet_offsets, current_offset).ok();
1391                if index == 1 {
1392                    assert!(packet.meta().is_simple_vote_tx());
1393                } else {
1394                    assert!(!packet.meta().is_simple_vote_tx());
1395                }
1396
1397                current_offset = current_offset.saturating_add(size_of::<Packet>());
1398            });
1399        }
1400
1401        // batch of mixed legacy messages and versioned vote tx, which won't be flagged as
1402        // simple_vote_tx
1403        {
1404            let mut current_offset = 0usize;
1405            let mut batch = PacketBatch::default();
1406            batch.push(Packet::from_data(None, test_tx()).unwrap());
1407            // versioned vote tx
1408            let tx = new_test_vote_tx(&mut rng);
1409            let mut packet = Packet::from_data(None, tx).unwrap();
1410            let packet_offsets = do_get_packet_offsets(&packet, 0).unwrap();
1411            let msg_start = packet_offsets.msg_start as usize;
1412            let msg_bytes = packet.data(msg_start..).unwrap().to_vec();
1413            packet.buffer_mut()[msg_start] = MESSAGE_VERSION_PREFIX;
1414            packet.meta_mut().size += 1;
1415            let msg_end = packet.meta().size;
1416            packet.buffer_mut()[msg_start + 1..msg_end].copy_from_slice(&msg_bytes);
1417            batch.push(packet);
1418
1419            batch.iter_mut().for_each(|packet| {
1420                let packet_offsets = do_get_packet_offsets(packet, current_offset).unwrap();
1421                check_for_simple_vote_transaction(packet, &packet_offsets, current_offset).ok();
1422                assert!(!packet.meta().is_simple_vote_tx());
1423
1424                current_offset = current_offset.saturating_add(size_of::<Packet>());
1425            });
1426        }
1427    }
1428
1429    #[test]
1430    fn test_shrink_fuzz() {
1431        for _ in 0..5 {
1432            let mut batches = to_packet_batches(
1433                &(0..PACKETS_PER_BATCH * 3)
1434                    .map(|_| test_tx())
1435                    .collect::<Vec<_>>(),
1436                PACKETS_PER_BATCH,
1437            );
1438            batches.iter_mut().for_each(|b| {
1439                b.iter_mut()
1440                    .for_each(|p| p.meta_mut().set_discard(thread_rng().gen()))
1441            });
1442            //find all the non discarded packets
1443            let mut start = vec![];
1444            batches.iter_mut().for_each(|b| {
1445                b.iter_mut()
1446                    .filter(|p| !p.meta().discard())
1447                    .for_each(|p| start.push(p.clone()))
1448            });
1449            start.sort_by(|a, b| a.data(..).cmp(&b.data(..)));
1450
1451            let packet_count = count_valid_packets(&batches);
1452            shrink_batches(&mut batches);
1453
1454            //make sure all the non discarded packets are the same
1455            let mut end = vec![];
1456            batches.iter_mut().for_each(|b| {
1457                b.iter_mut()
1458                    .filter(|p| !p.meta().discard())
1459                    .for_each(|p| end.push(p.clone()))
1460            });
1461            end.sort_by(|a, b| a.data(..).cmp(&b.data(..)));
1462            let packet_count2 = count_valid_packets(&batches);
1463            assert_eq!(packet_count, packet_count2);
1464            assert_eq!(start, end);
1465        }
1466    }
1467
1468    #[test]
1469    fn test_shrink_empty() {
1470        const PACKET_COUNT: usize = 1024;
1471        const BATCH_COUNT: usize = PACKET_COUNT / PACKETS_PER_BATCH;
1472
1473        // No batches
1474        // truncate of 1 on len 0 is a noop
1475        shrink_batches(&mut Vec::new());
1476        // One empty batch
1477        {
1478            let mut batches = vec![PacketBatch::with_capacity(0)];
1479            shrink_batches(&mut batches);
1480            assert_eq!(batches.len(), 0);
1481        }
1482        // Many empty batches
1483        {
1484            let mut batches = (0..BATCH_COUNT)
1485                .map(|_| PacketBatch::with_capacity(0))
1486                .collect::<Vec<_>>();
1487            shrink_batches(&mut batches);
1488            assert_eq!(batches.len(), 0);
1489        }
1490    }
1491
1492    #[test]
1493    fn test_shrink_vectors() {
1494        const PACKET_COUNT: usize = 1024;
1495        const BATCH_COUNT: usize = PACKET_COUNT / PACKETS_PER_BATCH;
1496
1497        let set_discards = [
1498            // contiguous
1499            // 0
1500            // No discards
1501            |_, _| false,
1502            // All discards
1503            |_, _| true,
1504            // single partitions
1505            // discard last half of packets
1506            |b, p| ((b * PACKETS_PER_BATCH) + p) >= (PACKET_COUNT / 2),
1507            // discard first half of packets
1508            |b, p| ((b * PACKETS_PER_BATCH) + p) < (PACKET_COUNT / 2),
1509            // discard last half of each batch
1510            |_, p| p >= (PACKETS_PER_BATCH / 2),
1511            // 5
1512            // discard first half of each batch
1513            |_, p| p < (PACKETS_PER_BATCH / 2),
1514            // uniform sparse
1515            // discard even packets
1516            |b, p| ((b * PACKETS_PER_BATCH) + p) % 2 == 0,
1517            // discard odd packets
1518            |b, p| ((b * PACKETS_PER_BATCH) + p) % 2 == 1,
1519            // discard even batches
1520            |b, _| b % 2 == 0,
1521            // discard odd batches
1522            |b, _| b % 2 == 1,
1523            // edges
1524            // 10
1525            // discard first batch
1526            |b, _| b == 0,
1527            // discard last batch
1528            |b, _| b == BATCH_COUNT - 1,
1529            // discard first and last batches
1530            |b, _| b == 0 || b == BATCH_COUNT - 1,
1531            // discard all but first and last batches
1532            |b, _| b != 0 && b != BATCH_COUNT - 1,
1533            // discard first packet
1534            |b, p| ((b * PACKETS_PER_BATCH) + p) == 0,
1535            // 15
1536            // discard all but first packet
1537            |b, p| ((b * PACKETS_PER_BATCH) + p) != 0,
1538            // discard last packet
1539            |b, p| ((b * PACKETS_PER_BATCH) + p) == PACKET_COUNT - 1,
1540            // discard all but last packet
1541            |b, p| ((b * PACKETS_PER_BATCH) + p) != PACKET_COUNT - 1,
1542            // discard first packet of each batch
1543            |_, p| p == 0,
1544            // discard all but first packet of each batch
1545            |_, p| p != 0,
1546            // 20
1547            // discard last packet of each batch
1548            |_, p| p == PACKETS_PER_BATCH - 1,
1549            // discard all but last packet of each batch
1550            |_, p| p != PACKETS_PER_BATCH - 1,
1551            // discard first and last packet of each batch
1552            |_, p| p == 0 || p == PACKETS_PER_BATCH - 1,
1553            // discard all but first and last packet of each batch
1554            |_, p| p != 0 && p != PACKETS_PER_BATCH - 1,
1555            // discard all after first packet in second to last batch
1556            |b, p| (b == BATCH_COUNT - 2 && p > 0) || b == BATCH_COUNT - 1,
1557            // 25
1558        ];
1559
1560        let expect_valids = [
1561            // (expected_batches, expected_valid_packets)
1562            //
1563            // contiguous
1564            // 0
1565            (BATCH_COUNT, PACKET_COUNT),
1566            (0, 0),
1567            // single partitions
1568            (BATCH_COUNT / 2, PACKET_COUNT / 2),
1569            (BATCH_COUNT / 2, PACKET_COUNT / 2),
1570            (BATCH_COUNT / 2, PACKET_COUNT / 2),
1571            // 5
1572            (BATCH_COUNT / 2, PACKET_COUNT / 2),
1573            // uniform sparse
1574            (BATCH_COUNT / 2, PACKET_COUNT / 2),
1575            (BATCH_COUNT / 2, PACKET_COUNT / 2),
1576            (BATCH_COUNT / 2, PACKET_COUNT / 2),
1577            (BATCH_COUNT / 2, PACKET_COUNT / 2),
1578            // edges
1579            // 10
1580            (BATCH_COUNT - 1, PACKET_COUNT - PACKETS_PER_BATCH),
1581            (BATCH_COUNT - 1, PACKET_COUNT - PACKETS_PER_BATCH),
1582            (BATCH_COUNT - 2, PACKET_COUNT - 2 * PACKETS_PER_BATCH),
1583            (2, 2 * PACKETS_PER_BATCH),
1584            (BATCH_COUNT, PACKET_COUNT - 1),
1585            // 15
1586            (1, 1),
1587            (BATCH_COUNT, PACKET_COUNT - 1),
1588            (1, 1),
1589            (
1590                (BATCH_COUNT * (PACKETS_PER_BATCH - 1) + PACKETS_PER_BATCH) / PACKETS_PER_BATCH,
1591                (PACKETS_PER_BATCH - 1) * BATCH_COUNT,
1592            ),
1593            (
1594                (BATCH_COUNT + PACKETS_PER_BATCH) / PACKETS_PER_BATCH,
1595                BATCH_COUNT,
1596            ),
1597            // 20
1598            (
1599                (BATCH_COUNT * (PACKETS_PER_BATCH - 1) + PACKETS_PER_BATCH) / PACKETS_PER_BATCH,
1600                (PACKETS_PER_BATCH - 1) * BATCH_COUNT,
1601            ),
1602            (
1603                (BATCH_COUNT + PACKETS_PER_BATCH) / PACKETS_PER_BATCH,
1604                BATCH_COUNT,
1605            ),
1606            (
1607                (BATCH_COUNT * (PACKETS_PER_BATCH - 2) + PACKETS_PER_BATCH) / PACKETS_PER_BATCH,
1608                (PACKETS_PER_BATCH - 2) * BATCH_COUNT,
1609            ),
1610            (
1611                (2 * BATCH_COUNT + PACKETS_PER_BATCH) / PACKETS_PER_BATCH,
1612                PACKET_COUNT - (PACKETS_PER_BATCH - 2) * BATCH_COUNT,
1613            ),
1614            (BATCH_COUNT - 1, PACKET_COUNT - 2 * PACKETS_PER_BATCH + 1),
1615            // 25
1616        ];
1617
1618        let test_cases = set_discards.iter().zip(&expect_valids).enumerate();
1619        for (i, (set_discard, (expect_batch_count, expect_valid_packets))) in test_cases {
1620            debug!("test_shrink case: {}", i);
1621            let mut batches = to_packet_batches(
1622                &(0..PACKET_COUNT).map(|_| test_tx()).collect::<Vec<_>>(),
1623                PACKETS_PER_BATCH,
1624            );
1625            assert_eq!(batches.len(), BATCH_COUNT);
1626            assert_eq!(count_valid_packets(&batches), PACKET_COUNT);
1627            batches.iter_mut().enumerate().for_each(|(i, b)| {
1628                b.iter_mut()
1629                    .enumerate()
1630                    .for_each(|(j, p)| p.meta_mut().set_discard(set_discard(i, j)))
1631            });
1632            assert_eq!(count_valid_packets(&batches), *expect_valid_packets);
1633            debug!("show valid packets for case {}", i);
1634            batches.iter_mut().enumerate().for_each(|(i, b)| {
1635                b.iter_mut().enumerate().for_each(|(j, p)| {
1636                    if !p.meta().discard() {
1637                        trace!("{} {}", i, j)
1638                    }
1639                })
1640            });
1641            debug!("done show valid packets for case {}", i);
1642            shrink_batches(&mut batches);
1643            let shrunken_batch_count = batches.len();
1644            debug!("shrunk batch test {} count: {}", i, shrunken_batch_count);
1645            assert_eq!(shrunken_batch_count, *expect_batch_count);
1646            assert_eq!(count_valid_packets(&batches), *expect_valid_packets);
1647        }
1648    }
1649}