Files
aria2-rust-pro/crates/aria2-rust-pro-protocol/src/torrent/peer_wire/handshake.rs
T

123 lines
4.0 KiB
Rust

use super::{
PEER_WIRE_HANDSHAKE_PREFIX_LEN, PEER_WIRE_PROTOCOL_LEN, PEER_WIRE_PROTOCOL_NAME,
PeerWireHandshakeModel,
};
impl PeerWireHandshakeModel {
#[must_use]
/// Builds a handshake with all reserved bits cleared.
pub fn new(info_hash: [u8; 20], peer_id: [u8; 20]) -> Self {
Self {
reserved: [0; 8],
info_hash,
peer_id,
}
}
/// Returns a handshake with the extension-protocol bit enabled.
#[must_use]
pub fn with_extension_protocol_enabled(mut self) -> Self {
self.reserved[5] |= 0x10;
self
}
/// Returns a handshake with the DHT bit enabled.
#[must_use]
pub fn with_dht_enabled(mut self) -> Self {
self.reserved[7] |= 0x01;
self
}
/// Serializes the handshake to its peer-wire byte representation.
///
/// # Panics
///
/// Panics if the fixed peer-wire protocol name no longer fits into a single-byte
/// length prefix.
#[must_use]
pub fn serialize(&self) -> Vec<u8> {
let mut bytes = Vec::with_capacity(PEER_WIRE_HANDSHAKE_PREFIX_LEN + 19);
bytes.push(PEER_WIRE_PROTOCOL_LEN);
bytes.extend_from_slice(PEER_WIRE_PROTOCOL_NAME.as_bytes());
bytes.extend_from_slice(&self.reserved);
bytes.extend_from_slice(&self.info_hash);
bytes.extend_from_slice(&self.peer_id);
bytes
}
/// Parses one complete peer-wire handshake from `input`.
///
/// # Errors
///
/// Returns an error when the frame is truncated, malformed, or contains trailing bytes.
pub fn parse(input: &[u8]) -> Result<Self, String> {
let (handshake, consumed) = Self::parse_prefix(input)?;
if consumed != input.len() {
return Err("trailing bytes after peer-wire handshake".to_owned());
}
Ok(handshake)
}
/// Parses a peer-wire handshake prefix and returns the consumed byte count.
///
/// # Errors
///
/// Returns an error when the frame is truncated or malformed.
pub fn parse_prefix(input: &[u8]) -> Result<(Self, usize), String> {
let Some(&protocol_len_byte) = input.first() else {
return Err("truncated peer-wire handshake: missing protocol length".to_owned());
};
let protocol_len = usize::from(protocol_len_byte);
let total_len = PEER_WIRE_HANDSHAKE_PREFIX_LEN + protocol_len;
if input.len() < total_len {
return Err(format!(
"truncated peer-wire handshake: expected {total_len} bytes, got {}",
input.len()
));
}
let protocol = &input[1..=protocol_len];
if protocol_len != PEER_WIRE_PROTOCOL_NAME.len() {
return Err(format!(
"invalid peer-wire protocol length: expected {}, got {protocol_len}",
PEER_WIRE_PROTOCOL_NAME.len()
));
}
if protocol != PEER_WIRE_PROTOCOL_NAME.as_bytes() {
return Err("invalid peer-wire protocol header".to_owned());
}
let reserved_start = 1 + protocol_len;
let mut reserved = [0_u8; 8];
reserved.copy_from_slice(&input[reserved_start..reserved_start + 8]);
let info_hash_start = reserved_start + 8;
let mut info_hash = [0_u8; 20];
info_hash.copy_from_slice(&input[info_hash_start..info_hash_start + 20]);
let peer_id_start = info_hash_start + 20;
let mut peer_id = [0_u8; 20];
peer_id.copy_from_slice(&input[peer_id_start..peer_id_start + 20]);
Ok((
Self {
reserved,
info_hash,
peer_id,
},
total_len,
))
}
#[must_use]
/// Returns whether the extension-protocol reserved bit is enabled.
pub fn extension_protocol_enabled(&self) -> bool {
self.reserved[5] & 0x10 != 0
}
#[must_use]
/// Returns whether the DHT reserved bit is enabled.
pub fn dht_enabled(&self) -> bool {
self.reserved[7] & 0x01 != 0
}
}