blvm_protocol/
p2p_frame.rs1use crate::error::{ProtocolError, Result};
7use crate::p2p_framing::MAX_PROTOCOL_MESSAGE_LENGTH;
8use sha2::{Digest, Sha256};
9use std::borrow::Cow;
10
11#[inline]
13pub fn bitcoin_p2p_payload_checksum(payload: &[u8]) -> [u8; 4] {
14 let hash1 = Sha256::digest(payload);
15 let hash2 = Sha256::digest(hash1);
16 let mut out = [0u8; 4];
17 out.copy_from_slice(&hash2[..4]);
18 out
19}
20
21pub fn parse_p2p_frame(
25 data: &[u8],
26 expected_magic_le: u32,
27 command_allowed: impl Fn(&str) -> bool,
28) -> Result<(&str, &[u8])> {
29 if data.len() < 24 {
30 return Err(ProtocolError::InvalidMessage(Cow::Owned(format!(
31 "Message too short: {} bytes",
32 data.len()
33 ))));
34 }
35 if data.len() > MAX_PROTOCOL_MESSAGE_LENGTH {
36 return Err(ProtocolError::MessageTooLarge {
37 size: data.len(),
38 max: MAX_PROTOCOL_MESSAGE_LENGTH,
39 });
40 }
41
42 let magic = u32::from_le_bytes([data[0], data[1], data[2], data[3]]);
43 if magic != expected_magic_le {
44 return Err(ProtocolError::InvalidMessage(Cow::Owned(format!(
45 "Invalid magic number 0x{magic:08x}"
46 ))));
47 }
48
49 let cmd_bytes = &data[4..16];
50 let end = cmd_bytes.iter().position(|&b| b == 0).unwrap_or(12);
51 let command = std::str::from_utf8(&cmd_bytes[..end]).map_err(|_| {
52 ProtocolError::InvalidMessage(Cow::Borrowed("Invalid UTF-8 in P2P command"))
53 })?;
54
55 if !command_allowed(command) {
56 return Err(ProtocolError::InvalidMessage(Cow::Owned(format!(
57 "Unknown command: {command}"
58 ))));
59 }
60
61 let payload_length = u32::from_le_bytes([data[16], data[17], data[18], data[19]]) as usize;
62 if payload_length > MAX_PROTOCOL_MESSAGE_LENGTH.saturating_sub(24) {
63 return Err(ProtocolError::InvalidMessage(Cow::Borrowed(
64 "Payload too large",
65 )));
66 }
67 if data.len() < 24 + payload_length {
68 return Err(ProtocolError::InvalidMessage(Cow::Borrowed(
69 "Incomplete message",
70 )));
71 }
72
73 let payload = &data[24..24 + payload_length];
74 let checksum = &data[20..24];
75 let expected = bitcoin_p2p_payload_checksum(payload);
76 if checksum != expected {
77 return Err(ProtocolError::InvalidMessage(Cow::Borrowed(
78 "Invalid checksum",
79 )));
80 }
81
82 Ok((command, payload))
83}
84
85pub fn build_p2p_frame(magic: [u8; 4], command: &str, payload: &[u8]) -> Result<Vec<u8>> {
87 if command.len() > 12 {
88 return Err(ProtocolError::InvalidMessage(Cow::Borrowed(
89 "P2P command longer than 12 bytes",
90 )));
91 }
92
93 let mut message = Vec::with_capacity(24 + payload.len());
94 message.extend_from_slice(&magic);
95 let mut command_bytes = [0u8; 12];
96 command_bytes[..command.len()].copy_from_slice(command.as_bytes());
97 message.extend_from_slice(&command_bytes);
98 message.extend_from_slice(&(payload.len() as u32).to_le_bytes());
99 message.extend_from_slice(&bitcoin_p2p_payload_checksum(payload));
100 message.extend_from_slice(payload);
101 Ok(message)
102}
103
104#[cfg(test)]
105mod tests {
106 use super::*;
107 use crate::p2p_framing::BITCOIN_MAGIC_MAINNET;
108
109 #[test]
110 fn checksum_matches_double_sha256_prefix() {
111 let p = [1u8, 2, 3];
112 let c = bitcoin_p2p_payload_checksum(&p);
113 let h1 = Sha256::digest(p);
114 let h2 = Sha256::digest(h1);
115 assert_eq!(c, h2[..4]);
116 }
117
118 #[test]
119 fn build_and_parse_roundtrip() {
120 let payload = vec![0xab, 0xcd];
121 let frame = build_p2p_frame(BITCOIN_MAGIC_MAINNET, "ping", &payload).unwrap();
122 let magic_le = u32::from_le_bytes(BITCOIN_MAGIC_MAINNET);
123 let (cmd, pl) = parse_p2p_frame(&frame, magic_le, |c| c == "ping" || c == "pong").unwrap();
124 assert_eq!(cmd, "ping");
125 assert_eq!(pl, payload.as_slice());
126 }
127
128 #[test]
129 fn unknown_command_rejected() {
130 let frame = build_p2p_frame(BITCOIN_MAGIC_MAINNET, "weird", &[]).unwrap();
131 let magic_le = u32::from_le_bytes(BITCOIN_MAGIC_MAINNET);
132 let err = parse_p2p_frame(&frame, magic_le, |c| c == "ping").unwrap_err();
133 assert!(format!("{err}").contains("Unknown command"));
134 }
135}