solana_transaction_metrics_tracker/
lib.rs

1use {
2    lazy_static::lazy_static, log::*, rand::Rng, solana_packet::Packet,
3    solana_perf::sigverify::PacketError, solana_short_vec::decode_shortu16_len,
4    solana_signature::SIGNATURE_BYTES,
5};
6
7// The mask is 12 bits long (1<<12 = 4096), it means the probability of matching
8// the transaction is 1/4096 assuming the portion being matched is random.
9lazy_static! {
10    static ref TXN_MASK: u16 = rand::thread_rng().gen_range(0..4096);
11}
12
13/// Check if a transaction given its signature matches the randomly selected mask.
14/// The signaure should be from the reference of Signature
15pub fn should_track_transaction(signature: &[u8; SIGNATURE_BYTES]) -> bool {
16    // We do not use the highest signature byte as it is not really random
17    let match_portion: u16 = u16::from_le_bytes([signature[61], signature[62]]) >> 4;
18    trace!("Matching txn: {match_portion:016b} {:016b}", *TXN_MASK);
19    *TXN_MASK == match_portion
20}
21
22/// Check if a transaction packet's signature matches the mask.
23/// This does a rudimentry verification to make sure the packet at least
24/// contains the signature data and it returns the reference to the signature.
25pub fn signature_if_should_track_packet(
26    packet: &Packet,
27) -> Result<Option<&[u8; SIGNATURE_BYTES]>, PacketError> {
28    let signature = get_signature_from_packet(packet)?;
29    Ok(should_track_transaction(signature).then_some(signature))
30}
31
32/// Get the signature of the transaction packet
33/// This does a rudimentry verification to make sure the packet at least
34/// contains the signature data and it returns the reference to the signature.
35pub fn get_signature_from_packet(packet: &Packet) -> Result<&[u8; SIGNATURE_BYTES], PacketError> {
36    let (sig_len_untrusted, sig_start) = packet
37        .data(..)
38        .and_then(|bytes| decode_shortu16_len(bytes).ok())
39        .ok_or(PacketError::InvalidShortVec)?;
40
41    if sig_len_untrusted < 1 {
42        return Err(PacketError::InvalidSignatureLen);
43    }
44
45    let signature = packet
46        .data(sig_start..sig_start.saturating_add(SIGNATURE_BYTES))
47        .ok_or(PacketError::InvalidSignatureLen)?;
48    let signature = signature
49        .try_into()
50        .map_err(|_| PacketError::InvalidSignatureLen)?;
51    Ok(signature)
52}
53
54#[cfg(test)]
55mod tests {
56    use {
57        super::*, solana_hash::Hash, solana_keypair::Keypair, solana_signature::Signature,
58        solana_system_transaction as system_transaction,
59    };
60
61    #[test]
62    fn test_get_signature_from_packet() {
63        // Default invalid txn packet
64        let packet = Packet::default();
65        let sig = get_signature_from_packet(&packet);
66        assert_eq!(sig, Err(PacketError::InvalidShortVec));
67
68        // Use a valid transaction, it should succeed
69        let tx = system_transaction::transfer(
70            &Keypair::new(),
71            &solana_pubkey::new_rand(),
72            1,
73            Hash::new_unique(),
74        );
75        let mut packet = Packet::from_data(None, tx).unwrap();
76
77        let sig = get_signature_from_packet(&packet);
78        assert!(sig.is_ok());
79
80        // Invalid signature length
81        packet.buffer_mut()[0] = 0x0;
82        let sig = get_signature_from_packet(&packet);
83        assert_eq!(sig, Err(PacketError::InvalidSignatureLen));
84    }
85
86    #[test]
87    fn test_should_track_transaction() {
88        let mut sig = [0x0; SIGNATURE_BYTES];
89        let track = should_track_transaction(&sig);
90        assert!(!track);
91
92        // Intentionally matching the randomly generated mask
93        // The lower four bits are ignored as only 12 highest bits from
94        // signature's 61 and 62 u8 are used for matching.
95        // We generate a random one
96        let mut rng = rand::thread_rng();
97        let random_number: u8 = rng.gen_range(0..=15);
98        sig[61] = ((*TXN_MASK & 0xf_u16) << 4) as u8 | random_number;
99        sig[62] = (*TXN_MASK >> 4) as u8;
100
101        let track = should_track_transaction(&sig);
102        assert!(track);
103    }
104
105    #[test]
106    fn test_signature_if_should_track_packet() {
107        // Default invalid txn packet
108        let packet = Packet::default();
109        let sig = signature_if_should_track_packet(&packet);
110        assert_eq!(sig, Err(PacketError::InvalidShortVec));
111
112        // Use a valid transaction which is not matched
113        let tx = system_transaction::transfer(
114            &Keypair::new(),
115            &solana_pubkey::new_rand(),
116            1,
117            Hash::new_unique(),
118        );
119        let packet = Packet::from_data(None, tx).unwrap();
120        let sig = signature_if_should_track_packet(&packet);
121        assert_eq!(Ok(None), sig);
122
123        // Now simulate a txn matching the signature mask
124        let mut tx = system_transaction::transfer(
125            &Keypair::new(),
126            &solana_pubkey::new_rand(),
127            1,
128            Hash::new_unique(),
129        );
130        let mut sig = [0x0; SIGNATURE_BYTES];
131        sig[61] = ((*TXN_MASK & 0xf_u16) << 4) as u8;
132        sig[62] = (*TXN_MASK >> 4) as u8;
133
134        let sig = Signature::from(sig);
135        tx.signatures[0] = sig;
136        let mut packet = Packet::from_data(None, tx).unwrap();
137        let sig2 = signature_if_should_track_packet(&packet);
138
139        match sig2 {
140            Ok(sig) => {
141                assert!(sig.is_some());
142            }
143            Err(_) => panic!("Expected to get a matching signature!"),
144        }
145
146        // Invalid signature length
147        packet.buffer_mut()[0] = 0x0;
148        let sig = signature_if_should_track_packet(&packet);
149        assert_eq!(sig, Err(PacketError::InvalidSignatureLen));
150    }
151}
solana_transaction_metrics_tracker/lib.rs

solana_transaction_metrics_tracker/
lib.rs
