solana/core/src/sigverify_shreds.rs

196 lines
6.3 KiB
Rust
Raw Normal View History

#![allow(clippy::implicit_hasher)]
2022-05-24 14:01:41 -07:00
use {
2022-05-24 14:01:41 -07:00
crate::{
sigverify,
sigverify_stage::{SigVerifier, SigVerifyServiceError},
},
crossbeam_channel::Sender,
solana_ledger::{
leader_schedule_cache::LeaderScheduleCache, shred, sigverify_shreds::verify_shreds_gpu,
},
2021-12-11 06:44:15 -08:00
solana_perf::{self, packet::PacketBatch, recycler_cache::RecyclerCache},
solana_runtime::{bank::Bank, bank_forks::BankForks},
solana_sdk::{clock::Slot, pubkey::Pubkey},
std::{
collections::HashMap,
sync::{Arc, RwLock},
},
};
2019-10-28 16:07:51 -07:00
#[derive(Clone)]
pub struct ShredSigVerifier {
pubkey: Pubkey, // TODO: Hot swap will change pubkey.
2019-10-28 16:07:51 -07:00
bank_forks: Arc<RwLock<BankForks>>,
leader_schedule_cache: Arc<LeaderScheduleCache>,
recycler_cache: RecyclerCache,
2022-05-24 14:01:41 -07:00
packet_sender: Sender<Vec<PacketBatch>>,
2019-10-28 16:07:51 -07:00
}
impl ShredSigVerifier {
pub fn new(
pubkey: Pubkey,
2019-10-28 16:07:51 -07:00
bank_forks: Arc<RwLock<BankForks>>,
leader_schedule_cache: Arc<LeaderScheduleCache>,
2022-05-24 14:01:41 -07:00
packet_sender: Sender<Vec<PacketBatch>>,
2019-10-28 16:07:51 -07:00
) -> Self {
sigverify::init();
Self {
pubkey,
2019-10-28 16:07:51 -07:00
bank_forks,
leader_schedule_cache,
recycler_cache: RecyclerCache::warmed(),
2022-05-24 14:01:41 -07:00
packet_sender,
2019-10-28 16:07:51 -07:00
}
}
}
impl SigVerifier for ShredSigVerifier {
2022-05-24 14:01:41 -07:00
type SendType = Vec<PacketBatch>;
fn send_packets(
&mut self,
packet_batches: Vec<PacketBatch>,
) -> Result<(), SigVerifyServiceError<Self::SendType>> {
self.packet_sender.send(packet_batches)?;
Ok(())
}
fn verify_batches(
&self,
mut batches: Vec<PacketBatch>,
_valid_packets: usize,
) -> Vec<PacketBatch> {
let working_bank = self.bank_forks.read().unwrap().working_bank();
let leader_slots: HashMap<Slot, [u8; 32]> = get_slot_leaders(
&self.pubkey,
&mut batches,
&self.leader_schedule_cache,
&working_bank,
)
.into_iter()
.filter_map(|(slot, pubkey)| Some((slot, pubkey?.to_bytes())))
.chain(std::iter::once((Slot::MAX, [0u8; 32])))
.collect();
let r = verify_shreds_gpu(&batches, &leader_slots, &self.recycler_cache);
solana_perf::sigverify::mark_disabled(&mut batches, &r);
batches
2019-10-28 16:07:51 -07:00
}
}
// Returns pubkey of leaders for shred slots refrenced in the packets.
// Marks packets as discard if:
// - fails to deserialize the shred slot.
// - slot leader is unknown.
// - slot leader is the node itself (circular transmission).
fn get_slot_leaders(
self_pubkey: &Pubkey,
batches: &mut [PacketBatch],
leader_schedule_cache: &LeaderScheduleCache,
bank: &Bank,
) -> HashMap<Slot, Option<Pubkey>> {
let mut leaders = HashMap::<Slot, Option<Pubkey>>::new();
for batch in batches {
for packet in batch.iter_mut() {
if packet.meta.discard() {
continue;
}
let shred = shred::layout::get_shred(packet);
let slot = match shred.and_then(shred::layout::get_slot) {
None => {
packet.meta.set_discard(true);
continue;
}
Some(slot) => slot,
};
let leader = leaders.entry(slot).or_insert_with(|| {
let leader = leader_schedule_cache.slot_leader_at(slot, Some(bank))?;
// Discard the shred if the slot leader is the node itself.
(&leader != self_pubkey).then(|| leader)
});
if leader.is_none() {
packet.meta.set_discard(true);
}
}
}
leaders
}
#[cfg(test)]
pub mod tests {
use {
super::*,
2022-05-24 14:01:41 -07:00
crossbeam_channel::unbounded,
solana_ledger::{
genesis_utils::create_genesis_config_with_leader,
shred::{Shred, ShredFlags},
},
solana_perf::packet::Packet,
solana_runtime::bank::Bank,
solana_sdk::signature::{Keypair, Signer},
};
2019-10-28 16:07:51 -07:00
#[test]
2021-12-11 06:44:15 -08:00
fn test_sigverify_shreds_verify_batches() {
2019-10-28 16:07:51 -07:00
let leader_keypair = Arc::new(Keypair::new());
let leader_pubkey = leader_keypair.pubkey();
let bank = Bank::new_for_tests(
&create_genesis_config_with_leader(100, &leader_pubkey, 10).genesis_config,
);
2019-10-28 16:07:51 -07:00
let cache = Arc::new(LeaderScheduleCache::new_from_bank(&bank));
let bf = Arc::new(RwLock::new(BankForks::new(bank)));
2022-05-24 14:01:41 -07:00
let (sender, receiver) = unbounded();
let mut verifier = ShredSigVerifier::new(Pubkey::new_unique(), bf, cache, sender);
2019-10-28 16:07:51 -07:00
let batch_size = 2;
let mut batch = PacketBatch::with_capacity(batch_size);
batch.resize(batch_size, Packet::default());
let mut batches = vec![batch];
2019-10-28 16:07:51 -07:00
let mut shred = Shred::new_from_data(
0,
0xc0de,
0xdead,
&[1, 2, 3, 4],
ShredFlags::LAST_SHRED_IN_SLOT,
0,
0,
0xc0de,
);
shred.sign(&leader_keypair);
batches[0][0].buffer_mut()[..shred.payload().len()].copy_from_slice(shred.payload());
batches[0][0].meta.size = shred.payload().len();
2019-10-28 16:07:51 -07:00
let mut shred = Shred::new_from_data(
0,
0xbeef,
0xc0de,
&[1, 2, 3, 4],
ShredFlags::LAST_SHRED_IN_SLOT,
0,
0,
0xc0de,
);
2019-10-28 16:07:51 -07:00
let wrong_keypair = Keypair::new();
shred.sign(&wrong_keypair);
batches[0][1].buffer_mut()[..shred.payload().len()].copy_from_slice(shred.payload());
batches[0][1].meta.size = shred.payload().len();
2019-10-28 16:07:51 -07:00
let num_packets = solana_perf::sigverify::count_packets_in_batches(&batches);
let rv = verifier.verify_batches(batches, num_packets);
assert!(!rv[0][0].meta.discard());
assert!(rv[0][1].meta.discard());
2022-05-24 14:01:41 -07:00
verifier.send_packets(rv.clone()).unwrap();
let received_packets = receiver.recv().unwrap();
assert_eq!(received_packets.len(), rv.len());
for (received_packet_batch, original_packet_batch) in received_packets.iter().zip(rv.iter())
{
assert_eq!(
received_packet_batch.iter().collect::<Vec<_>>(),
original_packet_batch.iter().collect::<Vec<_>>()
);
}
2019-10-28 16:07:51 -07:00
}
}