507 lines
19 KiB
Rust
507 lines
19 KiB
Rust
use super::{
|
|
PEER_WIRE_BITFIELD_ID, PEER_WIRE_CANCEL_ID, PEER_WIRE_CHOKE_ID, PEER_WIRE_EXTENSION_ID,
|
|
PEER_WIRE_HAVE_ID, PEER_WIRE_INTERESTED_ID, PEER_WIRE_NOT_INTERESTED_ID, PEER_WIRE_PIECE_ID,
|
|
PEER_WIRE_PORT_ID, PEER_WIRE_REQUEST_ID, PEER_WIRE_UNCHOKE_ID, PeerWireBitfieldModel,
|
|
PeerWireBlockRequestModel, PeerWireExtensionMessageModel, PeerWireFrameHeaderModel,
|
|
PeerWireMessageKind, PeerWireMessageModel, PeerWirePieceBlockModel,
|
|
PeerWireUnknownMessageModel, TorrentMessageModel,
|
|
};
|
|
|
|
impl PeerWireMessageModel {
|
|
#[must_use]
|
|
/// Wraps one parsed torrent message with an optional peer id.
|
|
pub fn new(peer_id: Option<[u8; 20]>, message: TorrentMessageModel) -> Self {
|
|
Self { peer_id, message }
|
|
}
|
|
|
|
/// Parses one framed peer-wire message and reports the number of bytes consumed.
|
|
///
|
|
/// # Errors
|
|
///
|
|
/// Returns an error when the frame header or payload is malformed.
|
|
pub fn parse_frame(input: &[u8]) -> Result<(Self, usize), String> {
|
|
let (message, consumed) = TorrentMessageModel::parse_peer_wire_frame(input)?;
|
|
Ok((
|
|
Self {
|
|
peer_id: None,
|
|
message,
|
|
},
|
|
consumed,
|
|
))
|
|
}
|
|
|
|
/// Parses one complete framed peer-wire message.
|
|
///
|
|
/// # Errors
|
|
///
|
|
/// Returns an error when the frame is malformed or contains trailing bytes.
|
|
pub fn parse_frame_exact(input: &[u8]) -> Result<Self, String> {
|
|
let (message, consumed) = Self::parse_frame(input)?;
|
|
if consumed != input.len() {
|
|
return Err("trailing bytes after peer-wire frame".to_owned());
|
|
}
|
|
Ok(message)
|
|
}
|
|
|
|
/// Serializes the wrapped message as a framed peer-wire payload.
|
|
///
|
|
/// # Errors
|
|
///
|
|
/// Returns an error when the inner message cannot be represented as peer-wire bytes.
|
|
pub fn serialize_frame(&self) -> Result<Vec<u8>, String> {
|
|
self.message.serialize_peer_wire_frame()
|
|
}
|
|
}
|
|
|
|
impl PeerWireBitfieldModel {
|
|
#[must_use]
|
|
/// Builds a bitfield from per-piece completion flags.
|
|
pub fn from_piece_flags(flags: &[bool]) -> Self {
|
|
let mut bytes = vec![0_u8; flags.len().div_ceil(8)];
|
|
for (index, &present) in flags.iter().enumerate() {
|
|
if present {
|
|
bytes[index / 8] |= 1 << (7 - (index % 8));
|
|
}
|
|
}
|
|
Self { bytes }
|
|
}
|
|
|
|
#[must_use]
|
|
/// Returns the maximum number of pieces represented by this bitfield.
|
|
pub fn piece_capacity(&self) -> usize {
|
|
self.bytes.len() * 8
|
|
}
|
|
|
|
#[must_use]
|
|
/// Returns whether the bitfield marks `piece_index` as present.
|
|
pub fn has_piece(&self, piece_index: usize) -> bool {
|
|
let byte = piece_index / 8;
|
|
let bit = piece_index % 8;
|
|
self.bytes
|
|
.get(byte)
|
|
.is_some_and(|value| value & (1 << (7 - bit)) != 0)
|
|
}
|
|
|
|
#[must_use]
|
|
/// Expands the bitfield into per-piece completion flags.
|
|
pub fn to_piece_flags(&self, piece_count: usize) -> Vec<bool> {
|
|
(0..piece_count)
|
|
.map(|index| self.has_piece(index))
|
|
.collect()
|
|
}
|
|
}
|
|
|
|
impl TorrentMessageModel {
|
|
#[must_use]
|
|
/// Builds an internal torrent message from a peer-wire message variant.
|
|
pub fn from_peer_wire_kind(kind: PeerWireMessageKind) -> Self {
|
|
match kind {
|
|
PeerWireMessageKind::KeepAlive => Self {
|
|
message_type: "keepalive".to_owned(),
|
|
payload: Vec::new(),
|
|
},
|
|
PeerWireMessageKind::Choke => Self {
|
|
message_type: "choke".to_owned(),
|
|
payload: Vec::new(),
|
|
},
|
|
PeerWireMessageKind::Unchoke => Self {
|
|
message_type: "unchoke".to_owned(),
|
|
payload: Vec::new(),
|
|
},
|
|
PeerWireMessageKind::Interested => Self {
|
|
message_type: "interested".to_owned(),
|
|
payload: Vec::new(),
|
|
},
|
|
PeerWireMessageKind::NotInterested => Self {
|
|
message_type: "not_interested".to_owned(),
|
|
payload: Vec::new(),
|
|
},
|
|
PeerWireMessageKind::Have(piece_index) => Self {
|
|
message_type: "have".to_owned(),
|
|
payload: piece_index.to_be_bytes().to_vec(),
|
|
},
|
|
PeerWireMessageKind::Bitfield(bitfield) => Self {
|
|
message_type: "bitfield".to_owned(),
|
|
payload: bitfield.bytes,
|
|
},
|
|
PeerWireMessageKind::Request(request) => Self {
|
|
message_type: "request".to_owned(),
|
|
payload: encode_block_request_payload(&request),
|
|
},
|
|
PeerWireMessageKind::Piece(piece) => Self {
|
|
message_type: "piece".to_owned(),
|
|
payload: encode_piece_payload(&piece),
|
|
},
|
|
PeerWireMessageKind::Cancel(request) => Self {
|
|
message_type: "cancel".to_owned(),
|
|
payload: encode_block_request_payload(&request),
|
|
},
|
|
PeerWireMessageKind::Port(port) => Self {
|
|
message_type: "port".to_owned(),
|
|
payload: port.to_be_bytes().to_vec(),
|
|
},
|
|
PeerWireMessageKind::Extension(extension) => {
|
|
let mut payload = Vec::with_capacity(1 + extension.payload.len());
|
|
payload.push(extension.extension_message_id);
|
|
payload.extend_from_slice(&extension.payload);
|
|
Self {
|
|
message_type: "extension".to_owned(),
|
|
payload,
|
|
}
|
|
}
|
|
PeerWireMessageKind::Unknown(message) => Self {
|
|
message_type: format!("unknown:{}", message.message_id),
|
|
payload: message.payload,
|
|
},
|
|
}
|
|
}
|
|
|
|
/// Reconstructs the typed peer-wire message kind from the internal message payload.
|
|
///
|
|
/// # Errors
|
|
///
|
|
/// Returns an error when the message type or payload shape is unsupported.
|
|
pub fn peer_wire_kind(&self) -> Result<PeerWireMessageKind, String> {
|
|
peer_wire_kind_from_raw(&self.message_type, &self.payload)
|
|
}
|
|
|
|
/// Inspects a framed peer-wire message header without fully decoding the payload.
|
|
///
|
|
/// # Errors
|
|
///
|
|
/// Returns an error when the frame is truncated or malformed.
|
|
pub fn inspect_peer_wire_frame(
|
|
input: &[u8],
|
|
) -> Result<(PeerWireFrameHeaderModel, usize), String> {
|
|
if input.len() < 4 {
|
|
return Err("truncated peer-wire frame: missing length prefix".to_owned());
|
|
}
|
|
|
|
let frame_len =
|
|
usize::try_from(u32::from_be_bytes([input[0], input[1], input[2], input[3]]))
|
|
.map_err(|_| "peer-wire frame length does not fit usize".to_owned())?;
|
|
let total_len = 4_usize
|
|
.checked_add(frame_len)
|
|
.ok_or_else(|| "peer-wire frame length overflow".to_owned())?;
|
|
if input.len() < total_len {
|
|
return Err(format!(
|
|
"truncated peer-wire frame: expected {total_len} bytes, got {}",
|
|
input.len()
|
|
));
|
|
}
|
|
|
|
if frame_len == 0 {
|
|
return Ok((
|
|
PeerWireFrameHeaderModel {
|
|
message_id: None,
|
|
payload_len: 0,
|
|
},
|
|
total_len,
|
|
));
|
|
}
|
|
|
|
let message_id = input[4];
|
|
Ok((
|
|
PeerWireFrameHeaderModel {
|
|
message_id: Some(message_id),
|
|
payload_len: frame_len - 1,
|
|
},
|
|
total_len,
|
|
))
|
|
}
|
|
|
|
/// Parses a framed peer-wire message and reports the number of consumed bytes.
|
|
///
|
|
/// # Errors
|
|
///
|
|
/// Returns an error when the frame is truncated or malformed.
|
|
pub fn parse_peer_wire_frame(input: &[u8]) -> Result<(Self, usize), String> {
|
|
let (header, consumed) = Self::inspect_peer_wire_frame(input)?;
|
|
let Some(message_id) = header.message_id else {
|
|
return Ok((
|
|
Self::from_peer_wire_kind(PeerWireMessageKind::KeepAlive),
|
|
consumed,
|
|
));
|
|
};
|
|
|
|
let payload = &input[5..consumed];
|
|
let kind = peer_wire_kind_from_message_id(message_id, payload)?;
|
|
Ok((Self::from_peer_wire_kind(kind), consumed))
|
|
}
|
|
|
|
/// Parses one complete framed peer-wire message.
|
|
///
|
|
/// # Errors
|
|
///
|
|
/// Returns an error when the frame is truncated, malformed, or has trailing bytes.
|
|
pub fn parse_peer_wire_frame_exact(input: &[u8]) -> Result<Self, String> {
|
|
let (message, consumed) = Self::parse_peer_wire_frame(input)?;
|
|
if consumed != input.len() {
|
|
return Err("trailing bytes after peer-wire frame".to_owned());
|
|
}
|
|
Ok(message)
|
|
}
|
|
|
|
/// Serializes the message as a framed peer-wire payload.
|
|
///
|
|
/// # Errors
|
|
///
|
|
/// Returns an error when the message cannot be represented as a supported peer-wire frame.
|
|
pub fn serialize_peer_wire_frame(&self) -> Result<Vec<u8>, String> {
|
|
let kind = self.peer_wire_kind()?;
|
|
serialize_peer_wire_kind(&kind)
|
|
}
|
|
}
|
|
|
|
/// Serializes one peer-wire message kind into a framed peer-wire payload.
|
|
fn serialize_peer_wire_kind(kind: &PeerWireMessageKind) -> Result<Vec<u8>, String> {
|
|
let mut payload = Vec::new();
|
|
let message_id = match kind {
|
|
PeerWireMessageKind::KeepAlive => None,
|
|
PeerWireMessageKind::Choke => Some(PEER_WIRE_CHOKE_ID),
|
|
PeerWireMessageKind::Unchoke => Some(PEER_WIRE_UNCHOKE_ID),
|
|
PeerWireMessageKind::Interested => Some(PEER_WIRE_INTERESTED_ID),
|
|
PeerWireMessageKind::NotInterested => Some(PEER_WIRE_NOT_INTERESTED_ID),
|
|
PeerWireMessageKind::Have(piece_index) => {
|
|
payload.extend_from_slice(&piece_index.to_be_bytes());
|
|
Some(PEER_WIRE_HAVE_ID)
|
|
}
|
|
PeerWireMessageKind::Bitfield(bitfield) => {
|
|
payload.extend_from_slice(&bitfield.bytes);
|
|
Some(PEER_WIRE_BITFIELD_ID)
|
|
}
|
|
PeerWireMessageKind::Request(request) => {
|
|
payload.extend_from_slice(&encode_block_request_payload(request));
|
|
Some(PEER_WIRE_REQUEST_ID)
|
|
}
|
|
PeerWireMessageKind::Piece(piece) => {
|
|
payload.extend_from_slice(&encode_piece_payload(piece));
|
|
Some(PEER_WIRE_PIECE_ID)
|
|
}
|
|
PeerWireMessageKind::Cancel(request) => {
|
|
payload.extend_from_slice(&encode_block_request_payload(request));
|
|
Some(PEER_WIRE_CANCEL_ID)
|
|
}
|
|
PeerWireMessageKind::Port(port) => {
|
|
payload.extend_from_slice(&port.to_be_bytes());
|
|
Some(PEER_WIRE_PORT_ID)
|
|
}
|
|
PeerWireMessageKind::Extension(extension) => {
|
|
payload.push(extension.extension_message_id);
|
|
payload.extend_from_slice(&extension.payload);
|
|
Some(PEER_WIRE_EXTENSION_ID)
|
|
}
|
|
PeerWireMessageKind::Unknown(message) => {
|
|
payload.extend_from_slice(&message.payload);
|
|
Some(message.message_id)
|
|
}
|
|
};
|
|
|
|
let Some(message_id) = message_id else {
|
|
return Ok(vec![0, 0, 0, 0]);
|
|
};
|
|
|
|
let frame_len = 1 + payload.len();
|
|
let frame_len_u32 = u32::try_from(frame_len)
|
|
.map_err(|_| "peer-wire frame exceeds u32 length prefix".to_owned())?;
|
|
|
|
let mut bytes = Vec::with_capacity(4 + frame_len);
|
|
bytes.extend_from_slice(&frame_len_u32.to_be_bytes());
|
|
bytes.push(message_id);
|
|
bytes.extend_from_slice(&payload);
|
|
Ok(bytes)
|
|
}
|
|
|
|
/// Interprets one internal message-type label and payload as a peer-wire message kind.
|
|
fn peer_wire_kind_from_raw(
|
|
message_type: &str,
|
|
payload: &[u8],
|
|
) -> Result<PeerWireMessageKind, String> {
|
|
match message_type {
|
|
"keepalive" => {
|
|
expect_empty_payload("keepalive", payload).map(|()| PeerWireMessageKind::KeepAlive)
|
|
}
|
|
"choke" => expect_empty_payload("choke", payload).map(|()| PeerWireMessageKind::Choke),
|
|
"unchoke" => {
|
|
expect_empty_payload("unchoke", payload).map(|()| PeerWireMessageKind::Unchoke)
|
|
}
|
|
"interested" => {
|
|
expect_empty_payload("interested", payload).map(|()| PeerWireMessageKind::Interested)
|
|
}
|
|
"not_interested" | "not-interested" => expect_empty_payload("not_interested", payload)
|
|
.map(|()| PeerWireMessageKind::NotInterested),
|
|
"have" => parse_have_payload(payload),
|
|
"bitfield" => Ok(PeerWireMessageKind::Bitfield(PeerWireBitfieldModel {
|
|
bytes: payload.to_vec(),
|
|
})),
|
|
"request" => {
|
|
parse_block_request_payload("request", payload).map(PeerWireMessageKind::Request)
|
|
}
|
|
"piece" => parse_piece_payload(payload).map(PeerWireMessageKind::Piece),
|
|
"cancel" => parse_block_request_payload("cancel", payload).map(PeerWireMessageKind::Cancel),
|
|
"port" => parse_port_payload(payload),
|
|
"extension" => parse_extension_payload(payload),
|
|
_ => parse_unknown_message_id(message_type).map_or_else(
|
|
|| {
|
|
Err(format!(
|
|
"unsupported peer-wire message type: {message_type}"
|
|
))
|
|
},
|
|
|message_id| {
|
|
Ok(PeerWireMessageKind::Unknown(PeerWireUnknownMessageModel {
|
|
message_id,
|
|
payload: payload.to_vec(),
|
|
}))
|
|
},
|
|
),
|
|
}
|
|
}
|
|
|
|
/// Interprets a peer-wire message id and payload as a typed peer-wire message kind.
|
|
fn peer_wire_kind_from_message_id(
|
|
message_id: u8,
|
|
payload: &[u8],
|
|
) -> Result<PeerWireMessageKind, String> {
|
|
match message_id {
|
|
PEER_WIRE_CHOKE_ID => {
|
|
expect_empty_payload("choke", payload).map(|()| PeerWireMessageKind::Choke)
|
|
}
|
|
PEER_WIRE_UNCHOKE_ID => {
|
|
expect_empty_payload("unchoke", payload).map(|()| PeerWireMessageKind::Unchoke)
|
|
}
|
|
PEER_WIRE_INTERESTED_ID => {
|
|
expect_empty_payload("interested", payload).map(|()| PeerWireMessageKind::Interested)
|
|
}
|
|
PEER_WIRE_NOT_INTERESTED_ID => expect_empty_payload("not_interested", payload)
|
|
.map(|()| PeerWireMessageKind::NotInterested),
|
|
PEER_WIRE_HAVE_ID => parse_have_payload(payload),
|
|
PEER_WIRE_BITFIELD_ID => Ok(PeerWireMessageKind::Bitfield(PeerWireBitfieldModel {
|
|
bytes: payload.to_vec(),
|
|
})),
|
|
PEER_WIRE_REQUEST_ID => {
|
|
parse_block_request_payload("request", payload).map(PeerWireMessageKind::Request)
|
|
}
|
|
PEER_WIRE_PIECE_ID => parse_piece_payload(payload).map(PeerWireMessageKind::Piece),
|
|
PEER_WIRE_CANCEL_ID => {
|
|
parse_block_request_payload("cancel", payload).map(PeerWireMessageKind::Cancel)
|
|
}
|
|
PEER_WIRE_PORT_ID => parse_port_payload(payload),
|
|
PEER_WIRE_EXTENSION_ID => parse_extension_payload(payload),
|
|
_ => Ok(PeerWireMessageKind::Unknown(PeerWireUnknownMessageModel {
|
|
message_id,
|
|
payload: payload.to_vec(),
|
|
})),
|
|
}
|
|
}
|
|
|
|
/// Verifies that a peer-wire control payload is empty.
|
|
fn expect_empty_payload(name: &str, payload: &[u8]) -> Result<(), String> {
|
|
if payload.is_empty() {
|
|
Ok(())
|
|
} else {
|
|
Err(format!(
|
|
"peer-wire {name} payload must be empty, got {} bytes",
|
|
payload.len()
|
|
))
|
|
}
|
|
}
|
|
|
|
/// Parses a `have` payload into its piece index variant.
|
|
fn parse_have_payload(payload: &[u8]) -> Result<PeerWireMessageKind, String> {
|
|
let piece_index = read_u32(payload, "have", 0)?;
|
|
Ok(PeerWireMessageKind::Have(piece_index))
|
|
}
|
|
|
|
/// Parses a `request` or `cancel` payload into block coordinates.
|
|
fn parse_block_request_payload(
|
|
name: &str,
|
|
payload: &[u8],
|
|
) -> Result<PeerWireBlockRequestModel, String> {
|
|
if payload.len() != 12 {
|
|
return Err(format!(
|
|
"peer-wire {name} payload must be 12 bytes, got {}",
|
|
payload.len()
|
|
));
|
|
}
|
|
Ok(PeerWireBlockRequestModel {
|
|
piece_index: read_u32(payload, name, 0)?,
|
|
block_offset: read_u32(payload, name, 4)?,
|
|
block_length: read_u32(payload, name, 8)?,
|
|
})
|
|
}
|
|
|
|
/// Parses a `piece` payload into block coordinates plus data.
|
|
fn parse_piece_payload(payload: &[u8]) -> Result<PeerWirePieceBlockModel, String> {
|
|
if payload.len() < 8 {
|
|
return Err(format!(
|
|
"peer-wire piece payload must be at least 8 bytes, got {}",
|
|
payload.len()
|
|
));
|
|
}
|
|
Ok(PeerWirePieceBlockModel {
|
|
piece_index: read_u32(payload, "piece", 0)?,
|
|
block_offset: read_u32(payload, "piece", 4)?,
|
|
block: payload[8..].to_vec(),
|
|
})
|
|
}
|
|
|
|
/// Parses a `port` payload into the corresponding peer-wire message variant.
|
|
fn parse_port_payload(payload: &[u8]) -> Result<PeerWireMessageKind, String> {
|
|
if payload.len() != 2 {
|
|
return Err(format!(
|
|
"peer-wire port payload must be 2 bytes, got {}",
|
|
payload.len()
|
|
));
|
|
}
|
|
Ok(PeerWireMessageKind::Port(u16::from_be_bytes([
|
|
payload[0], payload[1],
|
|
])))
|
|
}
|
|
|
|
/// Parses an extension-protocol payload into the typed extension message variant.
|
|
fn parse_extension_payload(payload: &[u8]) -> Result<PeerWireMessageKind, String> {
|
|
let Some((&extension_message_id, rest)) = payload.split_first() else {
|
|
return Err("peer-wire extension payload must include extension message id".to_owned());
|
|
};
|
|
Ok(PeerWireMessageKind::Extension(
|
|
PeerWireExtensionMessageModel {
|
|
extension_message_id,
|
|
payload: rest.to_vec(),
|
|
},
|
|
))
|
|
}
|
|
|
|
/// Parses a synthetic `unknown:<id>` message-type label into a raw peer-wire id.
|
|
fn parse_unknown_message_id(message_type: &str) -> Option<u8> {
|
|
message_type
|
|
.strip_prefix("unknown:")
|
|
.and_then(|value| value.parse::<u8>().ok())
|
|
}
|
|
|
|
/// Encodes request or cancel block coordinates into peer-wire payload bytes.
|
|
fn encode_block_request_payload(request: &PeerWireBlockRequestModel) -> Vec<u8> {
|
|
let mut payload = Vec::with_capacity(12);
|
|
payload.extend_from_slice(&request.piece_index.to_be_bytes());
|
|
payload.extend_from_slice(&request.block_offset.to_be_bytes());
|
|
payload.extend_from_slice(&request.block_length.to_be_bytes());
|
|
payload
|
|
}
|
|
|
|
/// Encodes a piece block into peer-wire payload bytes.
|
|
fn encode_piece_payload(piece: &PeerWirePieceBlockModel) -> Vec<u8> {
|
|
let mut payload = Vec::with_capacity(8 + piece.block.len());
|
|
payload.extend_from_slice(&piece.piece_index.to_be_bytes());
|
|
payload.extend_from_slice(&piece.block_offset.to_be_bytes());
|
|
payload.extend_from_slice(&piece.block);
|
|
payload
|
|
}
|
|
|
|
/// Reads one big-endian `u32` from a peer-wire payload.
|
|
fn read_u32(payload: &[u8], name: &str, start: usize) -> Result<u32, String> {
|
|
let end = start + 4;
|
|
let bytes = payload
|
|
.get(start..end)
|
|
.ok_or_else(|| format!("peer-wire {name} payload truncated at byte offset {start}"))?;
|
|
Ok(u32::from_be_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]))
|
|
}
|