Skip to main content

blvm_protocol/
p2p_frame.rs

1//! Bitcoin P2P message framing: magic, 12-byte command, LE length, 4-byte checksum, payload.
2//!
3//! Shared between crates so [`crate::p2p_framing::MAX_PROTOCOL_MESSAGE_LENGTH`] and checksum
4//! rules stay single-sourced.
5
6use crate::error::{ProtocolError, Result};
7use crate::p2p_framing::MAX_PROTOCOL_MESSAGE_LENGTH;
8use sha2::{Digest, Sha256};
9use std::borrow::Cow;
10
11/// First four bytes of double-SHA256(payload), Bitcoin P2P checksum.
12#[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
21/// Parse the 24-byte header and verify checksum; returns command name and payload slice.
22///
23/// `command_allowed` should return true for commands this node/process accepts (e.g. allowlist).
24pub 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
85/// Build a full P2P frame: magic + command (null-padded) + length + checksum + payload.
86pub 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}