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

use std::collections::BTreeMap;

use super::{
    bencode::{
        BencodeValue, dict_bytes, dict_int, encode_bencode_root, parse_bencode_root_exact,
        parse_bencode_root_prefix,
    },
    utils::{bytes_to_string, i64_to_u64},
};

/// BEP 10 extension-protocol handshake and metadata helpers.
mod extension;
/// Peer-wire frame parsing and serialization helpers.
mod framing;
/// `BitTorrent` handshake parsing and serialization helpers.
mod handshake;

/// Generic torrent message wrapper reused by higher-level peer-wire helpers.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct TorrentMessageModel {
    /// Canonical internal message type name.
    pub message_type: String,
    /// Raw payload bytes excluding transport framing.
    pub payload: Vec<u8>,
}

/// `BitTorrent` peer-wire handshake header.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct PeerWireHandshakeModel {
    /// Reserved extension bits.
    pub reserved: [u8; 8],
    /// 20-byte torrent info-hash.
    pub info_hash: [u8; 20],
    /// 20-byte local peer id.
    pub peer_id: [u8; 20],
}

/// One parsed peer-wire message, optionally associated with a peer id.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct PeerWireMessageModel {
    /// Optional peer id attached by higher-level wrappers.
    pub peer_id: Option<[u8; 20]>,
    /// Parsed message payload.
    pub message: TorrentMessageModel,
}

/// Lightweight inspection result for a framed peer-wire message.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct PeerWireFrameHeaderModel {
    /// Optional peer-wire message id. `None` represents keepalive.
    pub message_id: Option<u8>,
    /// Payload length excluding the length prefix and optional message id byte.
    pub payload_len: usize,
}

/// Packed peer-wire bitfield bytes.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct PeerWireBitfieldModel {
    /// Raw bitfield bytes in network order.
    pub bytes: Vec<u8>,
}

/// Piece request or cancel coordinates for the peer-wire protocol.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct PeerWireBlockRequestModel {
    /// Zero-based piece index.
    pub piece_index: u32,
    /// Byte offset within the piece.
    pub block_offset: u32,
    /// Requested block length in bytes.
    pub block_length: u32,
}

/// Piece payload delivered through the peer-wire protocol.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct PeerWirePieceBlockModel {
    /// Zero-based piece index.
    pub piece_index: u32,
    /// Byte offset within the piece.
    pub block_offset: u32,
    /// Raw block bytes.
    pub block: Vec<u8>,
}

/// Extension-protocol message payload.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct PeerWireExtensionMessageModel {
    /// Extension message id.
    pub extension_message_id: u8,
    /// Extension payload bytes after the extension id.
    pub payload: Vec<u8>,
}

/// Parsed extended-handshake payload from BEP 10.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct PeerWireExtensionHandshakeModel {
    /// Named extension ids announced under the `m` dictionary.
    pub extensions: BTreeMap<String, u8>,
    /// Optional peer/client version string from `v`.
    pub client_name: Option<String>,
    /// Optional BEP 9 metadata byte length.
    pub metadata_size: Option<u32>,
    /// Optional request queue depth from `reqq`.
    pub request_queue: Option<u32>,
}

/// BEP 9 `ut_metadata` message type.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum PeerWireMetadataMessageType {
    /// Requests one metadata piece.
    Request,
    /// Carries one metadata piece payload.
    Data,
    /// Rejects one metadata piece request.
    Reject,
}

impl PeerWireMetadataMessageType {
    /// Returns the BEP 9 wire value for the metadata message type.
    #[must_use]
    pub const fn wire_value(self) -> u8 {
        match self {
            Self::Request => 0,
            Self::Data => 1,
            Self::Reject => 2,
        }
    }

    /// Decodes one BEP 9 wire value into the typed metadata message kind.
    fn from_wire_value(value: i64) -> Result<Self, String> {
        match value {
            0 => Ok(Self::Request),
            1 => Ok(Self::Data),
            2 => Ok(Self::Reject),
            _ => Err(format!("unsupported ut_metadata msg_type: {value}")),
        }
    }
}

/// Default BEP 9 metadata piece size in bytes.
pub const PEER_WIRE_METADATA_PIECE_SIZE: u32 = 16 * 1024;

/// Parsed BEP 9 `ut_metadata` message with header and optional payload bytes.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct PeerWireMetadataMessageModel {
    /// Metadata message subtype.
    pub message_type: PeerWireMetadataMessageType,
    /// Metadata piece index addressed by the message.
    pub piece: u32,
    /// Total metadata byte length when included in `data` messages.
    pub total_size: Option<u32>,
    /// Metadata payload bytes for `data` messages.
    pub payload: Vec<u8>,
}

/// Unknown peer-wire message payload retained losslessly.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct PeerWireUnknownMessageModel {
    /// Raw peer-wire message id.
    pub message_id: u8,
    /// Unparsed message payload bytes.
    pub payload: Vec<u8>,
}

/// Supported peer-wire message variants.
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum PeerWireMessageKind {
    /// Zero-length keepalive frame.
    KeepAlive,
    /// Choke control message.
    Choke,
    /// Unchoke control message.
    Unchoke,
    /// Interested control message.
    Interested,
    /// Not-interested control message.
    NotInterested,
    /// `have` message naming one completed piece.
    Have(u32),
    /// `bitfield` message.
    Bitfield(PeerWireBitfieldModel),
    /// `request` message.
    Request(PeerWireBlockRequestModel),
    /// `piece` message.
    Piece(PeerWirePieceBlockModel),
    /// `cancel` message.
    Cancel(PeerWireBlockRequestModel),
    /// `port` DHT advertisement message.
    Port(u16),
    /// Extension-protocol message.
    Extension(PeerWireExtensionMessageModel),
    /// Unknown message preserved losslessly.
    Unknown(PeerWireUnknownMessageModel),
}

/// Canonical protocol string embedded in peer-wire handshakes.
const PEER_WIRE_PROTOCOL_NAME: &str = "BitTorrent protocol";
/// Byte length of [`PEER_WIRE_PROTOCOL_NAME`].
const PEER_WIRE_PROTOCOL_LEN: u8 = 19;
/// Handshake bytes following the protocol-length octet and protocol string.
const PEER_WIRE_HANDSHAKE_PREFIX_LEN: usize = 49;
/// Peer-wire message id for `choke`.
const PEER_WIRE_CHOKE_ID: u8 = 0;
/// Peer-wire message id for `unchoke`.
const PEER_WIRE_UNCHOKE_ID: u8 = 1;
/// Peer-wire message id for `interested`.
const PEER_WIRE_INTERESTED_ID: u8 = 2;
/// Peer-wire message id for `not interested`.
const PEER_WIRE_NOT_INTERESTED_ID: u8 = 3;
/// Peer-wire message id for `have`.
const PEER_WIRE_HAVE_ID: u8 = 4;
/// Peer-wire message id for `bitfield`.
const PEER_WIRE_BITFIELD_ID: u8 = 5;
/// Peer-wire message id for `request`.
const PEER_WIRE_REQUEST_ID: u8 = 6;
/// Peer-wire message id for `piece`.
const PEER_WIRE_PIECE_ID: u8 = 7;
/// Peer-wire message id for `cancel`.
const PEER_WIRE_CANCEL_ID: u8 = 8;
/// Peer-wire message id for `port`.
const PEER_WIRE_PORT_ID: u8 = 9;
/// Peer-wire message id for extension-protocol payloads.
const PEER_WIRE_EXTENSION_ID: u8 = 20;