361 lines
15 KiB
Rust
361 lines
15 KiB
Rust
use super::{
|
|
BTreeMap, BTreeSet, BtPeerInfo, Digest, PeerWireBlockRequestModel,
|
|
PeerWireExtensionHandshakeModel, PeerWireHandshakeModel, PeerWireMessageKind,
|
|
PeerWireMetadataMessageModel, PeerWirePieceBlockModel, PeerWireTransportRequest,
|
|
PeerWireTransportResponse, PieceState, RequestGroup, RpcError, TorrentMessageModel,
|
|
TransportEndpoint, TransportScheme, bt_metadata_piece_count, bt_peer_metadata_key,
|
|
bt_piece_span_bytes, bt_runtime_total_length, hex_string, resolve_bt_info_hash,
|
|
rpc_bt_local_node_id,
|
|
};
|
|
|
|
/// Identifies the peer selected for a peer-wire compatibility exchange.
|
|
#[derive(Clone, Debug)]
|
|
pub(in crate::dispatcher) struct PeerWirePeerTarget {
|
|
/// Original peer index inside the runtime peer list.
|
|
pub(in crate::dispatcher) index: usize,
|
|
/// Peer runtime snapshot used to build the outbound request.
|
|
pub(in crate::dispatcher) peer: BtPeerInfo,
|
|
}
|
|
|
|
/// Holds the outbound peer-wire request and bookkeeping for a compatibility probe.
|
|
#[derive(Clone, Debug)]
|
|
pub(in crate::dispatcher) struct PeerWireExchangePlan {
|
|
/// Peer list index that should receive the parsed response data.
|
|
pub(in crate::dispatcher) peer_index: usize,
|
|
/// Expected info hash validated against the peer handshake.
|
|
pub(in crate::dispatcher) info_hash: [u8; 20],
|
|
/// Transport payload sent to the peer.
|
|
pub(in crate::dispatcher) request: PeerWireTransportRequest,
|
|
/// Optional block request emitted after the handshake.
|
|
pub(in crate::dispatcher) block_request: Option<PeerWireBlockRequestModel>,
|
|
/// Optional known `ut_metadata` extension id for the selected peer.
|
|
pub(in crate::dispatcher) metadata_extension_id: Option<u8>,
|
|
}
|
|
|
|
/// Captures peer-wire handshake and frame state recovered from a peer response.
|
|
#[derive(Clone, Debug, Default)]
|
|
pub(in crate::dispatcher) struct PeerWireExchangeResponseModel {
|
|
/// Remote peer ID emitted by the handshake when present.
|
|
pub(in crate::dispatcher) peer_id: Option<String>,
|
|
/// Parsed remote extended handshake, when observed.
|
|
pub(in crate::dispatcher) extension_handshake: Option<PeerWireExtensionHandshakeModel>,
|
|
/// Latest observed choke state.
|
|
pub(in crate::dispatcher) peer_choked: Option<bool>,
|
|
/// Latest observed interest state.
|
|
pub(in crate::dispatcher) peer_interested: Option<bool>,
|
|
/// Piece set explicitly advertised by a bitfield frame.
|
|
pub(in crate::dispatcher) bitfield_pieces: Option<BTreeSet<u32>>,
|
|
/// Aggregate set of pieces implied by bitfield, have, and piece frames.
|
|
pub(in crate::dispatcher) available_pieces: BTreeSet<u32>,
|
|
/// Piece payload frames recovered from the response.
|
|
pub(in crate::dispatcher) pieces: Vec<PeerWirePieceBlockModel>,
|
|
/// Metadata payloads recovered from BEP 9 messages.
|
|
pub(in crate::dispatcher) metadata_messages: Vec<PeerWireMetadataMessageModel>,
|
|
}
|
|
|
|
/// Builds the outbound peer-wire request for the current BitTorrent runtime state.
|
|
pub(in crate::dispatcher) fn build_peer_wire_exchange_plan(
|
|
group: &RequestGroup,
|
|
) -> Result<PeerWireExchangePlan, RpcError> {
|
|
let bt = group
|
|
.bt()
|
|
.ok_or_else(|| RpcError::unsupported("peer-wire exchange requires bt runtime state"))?;
|
|
let target = pick_bt_peer_target(&bt.peers)?;
|
|
let info_hash_vec = resolve_bt_info_hash(group, bt)?;
|
|
let info_hash: [u8; 20] = info_hash_vec
|
|
.as_slice()
|
|
.try_into()
|
|
.map_err(|_| RpcError::unsupported("peer-wire exchange info hash must be 20 bytes"))?;
|
|
let peer_id_vec = rpc_bt_local_node_id(group.gid());
|
|
let peer_id: [u8; 20] = peer_id_vec
|
|
.clone()
|
|
.try_into()
|
|
.map_err(|_| RpcError::unsupported("peer-wire exchange peer id must be 20 bytes"))?;
|
|
|
|
let peer_key = bt_peer_metadata_key(&target.peer);
|
|
let metadata_extension_id = bt.metadata_extension_ids.get(&peer_key).copied();
|
|
let metadata_request_piece = if bt.metadata_only {
|
|
match (metadata_extension_id, bt.metadata_size) {
|
|
(Some(_), Some(metadata_size)) if metadata_size > 0 => {
|
|
let piece_count = bt_metadata_piece_count(metadata_size);
|
|
(0..piece_count).find(|piece| !bt.metadata_piece_payloads.contains_key(piece))
|
|
}
|
|
_ => None,
|
|
}
|
|
} else {
|
|
None
|
|
};
|
|
|
|
let mut handshake = PeerWireHandshakeModel::new(info_hash, peer_id);
|
|
handshake.reserved[5] |= 0x10;
|
|
let mut payload = handshake.serialize();
|
|
|
|
let mut block_request = None;
|
|
if bt.metadata_only {
|
|
let extension_handshake = PeerWireExtensionHandshakeModel {
|
|
extensions: BTreeMap::from([("ut_metadata".to_owned(), 1_u8)]),
|
|
client_name: Some("aria2-rust-pro".to_owned()),
|
|
metadata_size: None,
|
|
request_queue: Some(16),
|
|
};
|
|
payload.extend_from_slice(
|
|
&TorrentMessageModel::from_peer_wire_kind(PeerWireMessageKind::Extension(
|
|
extension_handshake.to_peer_wire_message(),
|
|
))
|
|
.serialize_peer_wire_frame()
|
|
.map_err(|error| {
|
|
RpcError::unsupported(&format!(
|
|
"peer-wire extension handshake serialization failed: {error}"
|
|
))
|
|
})?,
|
|
);
|
|
if let (Some(extension_message_id), Some(piece)) =
|
|
(metadata_extension_id, metadata_request_piece)
|
|
{
|
|
payload.extend_from_slice(
|
|
&TorrentMessageModel::from_peer_wire_kind(PeerWireMessageKind::Extension(
|
|
PeerWireMetadataMessageModel::request(piece)
|
|
.to_peer_wire_message(extension_message_id),
|
|
))
|
|
.serialize_peer_wire_frame()
|
|
.map_err(|error| {
|
|
RpcError::unsupported(&format!(
|
|
"peer-wire metadata request serialization failed: {error}"
|
|
))
|
|
})?,
|
|
);
|
|
}
|
|
} else {
|
|
let piece_length = group.piece_length().max(1);
|
|
let total_length = bt_runtime_total_length(group);
|
|
let requestable = group.bt_requestable_piece_ids(false, 8);
|
|
let availability = group.piece_availability();
|
|
let selected_piece = requestable
|
|
.iter()
|
|
.copied()
|
|
.find(|piece| availability.get(piece).copied().unwrap_or(0) > 0)
|
|
.or_else(|| requestable.first().copied());
|
|
block_request = selected_piece.map(|piece| PeerWireBlockRequestModel {
|
|
piece_index: piece.0,
|
|
block_offset: 0,
|
|
block_length: bt_piece_span_bytes(piece, piece_length, total_length)
|
|
.min(16_u64 * 1024)
|
|
.max(1) as u32,
|
|
});
|
|
if block_request.is_some() {
|
|
payload.extend_from_slice(
|
|
&TorrentMessageModel::from_peer_wire_kind(PeerWireMessageKind::Interested)
|
|
.serialize_peer_wire_frame()
|
|
.map_err(|error| {
|
|
RpcError::unsupported(&format!(
|
|
"peer-wire interested frame serialization failed: {error}"
|
|
))
|
|
})?,
|
|
);
|
|
}
|
|
if let Some(block_request) = &block_request {
|
|
payload.extend_from_slice(
|
|
&TorrentMessageModel::from_peer_wire_kind(PeerWireMessageKind::Request(
|
|
block_request.clone(),
|
|
))
|
|
.serialize_peer_wire_frame()
|
|
.map_err(|error| {
|
|
RpcError::unsupported(&format!(
|
|
"peer-wire request frame serialization failed: {error}"
|
|
))
|
|
})?,
|
|
);
|
|
}
|
|
}
|
|
|
|
Ok(PeerWireExchangePlan {
|
|
peer_index: target.index,
|
|
info_hash,
|
|
request: PeerWireTransportRequest {
|
|
endpoint: TransportEndpoint {
|
|
scheme: TransportScheme::BitTorrent,
|
|
address: format!("{}:{}", target.peer.ip, target.peer.port),
|
|
},
|
|
info_hash: info_hash_vec,
|
|
peer_id: peer_id_vec,
|
|
payload,
|
|
},
|
|
block_request,
|
|
metadata_extension_id,
|
|
})
|
|
}
|
|
|
|
/// Chooses the best available peer target for a peer-wire exchange.
|
|
pub(super) fn pick_bt_peer_target(peers: &[BtPeerInfo]) -> Result<PeerWirePeerTarget, RpcError> {
|
|
if peers.is_empty() {
|
|
return Err(RpcError::unsupported(
|
|
"peer-wire exchange requires at least one bt peer",
|
|
));
|
|
}
|
|
let mut last_error = None;
|
|
let mut fallback = None;
|
|
for (index, peer) in peers.iter().enumerate() {
|
|
if peer.ip.trim().is_empty() {
|
|
last_error = Some("peer ip must not be empty".to_owned());
|
|
continue;
|
|
}
|
|
if peer.port == 0 {
|
|
last_error = Some("peer port must be non-zero".to_owned());
|
|
continue;
|
|
}
|
|
let target = PeerWirePeerTarget {
|
|
index,
|
|
peer: peer.clone(),
|
|
};
|
|
if !peer.choked {
|
|
return Ok(target);
|
|
}
|
|
if fallback.is_none() {
|
|
fallback = Some(target);
|
|
}
|
|
}
|
|
if let Some(target) = fallback {
|
|
return Ok(target);
|
|
}
|
|
Err(RpcError::unsupported(&format!(
|
|
"peer-wire exchange found no valid bt peers in runtime state{}",
|
|
last_error
|
|
.map(|message| format!(": {message}"))
|
|
.unwrap_or_default()
|
|
)))
|
|
}
|
|
|
|
/// Parses a peer-wire transport payload into normalized runtime update data.
|
|
pub(in crate::dispatcher) fn parse_peer_wire_exchange_response(
|
|
response: &PeerWireTransportResponse,
|
|
expected_info_hash: &[u8; 20],
|
|
known_metadata_extension_id: Option<u8>,
|
|
) -> Result<PeerWireExchangeResponseModel, RpcError> {
|
|
let mut parsed = PeerWireExchangeResponseModel::default();
|
|
let mut cursor = 0;
|
|
let mut negotiated_metadata_extension_id = known_metadata_extension_id;
|
|
if response.payload.first().copied() == Some(19) {
|
|
let (handshake, consumed) = PeerWireHandshakeModel::parse_prefix(&response.payload)
|
|
.map_err(|error| {
|
|
RpcError::unsupported(&format!("invalid peer-wire handshake: {error}"))
|
|
})?;
|
|
if handshake.info_hash != *expected_info_hash {
|
|
return Err(RpcError::unsupported(
|
|
"peer-wire handshake info hash did not match download runtime state",
|
|
));
|
|
}
|
|
parsed.peer_id = Some(hex_string(&handshake.peer_id).to_ascii_lowercase());
|
|
cursor = consumed;
|
|
}
|
|
|
|
while cursor < response.payload.len() {
|
|
let (frame, consumed) = TorrentMessageModel::parse_peer_wire_frame(
|
|
&response.payload[cursor..],
|
|
)
|
|
.map_err(|error| RpcError::unsupported(&format!("invalid peer-wire frame: {error}")))?;
|
|
cursor = cursor.saturating_add(consumed);
|
|
match frame.peer_wire_kind().map_err(|error| {
|
|
RpcError::unsupported(&format!("invalid peer-wire message: {error}"))
|
|
})? {
|
|
PeerWireMessageKind::Choke => parsed.peer_choked = Some(true),
|
|
PeerWireMessageKind::Unchoke => parsed.peer_choked = Some(false),
|
|
PeerWireMessageKind::Interested => parsed.peer_interested = Some(true),
|
|
PeerWireMessageKind::NotInterested => parsed.peer_interested = Some(false),
|
|
PeerWireMessageKind::Have(piece) => {
|
|
parsed.available_pieces.insert(piece);
|
|
}
|
|
PeerWireMessageKind::Bitfield(bitfield) => {
|
|
let available = bitfield
|
|
.to_piece_flags(bitfield.piece_capacity())
|
|
.into_iter()
|
|
.enumerate()
|
|
.filter_map(|(piece, has_piece)| has_piece.then_some(piece as u32))
|
|
.collect::<BTreeSet<_>>();
|
|
parsed.available_pieces.extend(available.iter().copied());
|
|
parsed.bitfield_pieces = Some(available);
|
|
}
|
|
PeerWireMessageKind::Piece(piece) => {
|
|
parsed.available_pieces.insert(piece.piece_index);
|
|
parsed.pieces.push(piece);
|
|
}
|
|
PeerWireMessageKind::Extension(message) => {
|
|
if message.extension_message_id == 0 {
|
|
let handshake = PeerWireExtensionHandshakeModel::from_peer_wire_message(
|
|
&message,
|
|
)
|
|
.map_err(|error| {
|
|
RpcError::unsupported(&format!(
|
|
"invalid peer-wire extended handshake: {error}"
|
|
))
|
|
})?;
|
|
negotiated_metadata_extension_id = handshake
|
|
.ut_metadata_id()
|
|
.or(negotiated_metadata_extension_id);
|
|
parsed.extension_handshake = Some(handshake);
|
|
} else if negotiated_metadata_extension_id
|
|
.is_some_and(|extension_id| extension_id == message.extension_message_id)
|
|
{
|
|
let metadata_message = PeerWireMetadataMessageModel::from_peer_wire_message(
|
|
&message,
|
|
negotiated_metadata_extension_id.expect("checked is_some above"),
|
|
)
|
|
.map_err(|error| {
|
|
RpcError::unsupported(&format!(
|
|
"invalid peer-wire ut_metadata payload: {error}"
|
|
))
|
|
})?;
|
|
parsed.metadata_messages.push(metadata_message);
|
|
}
|
|
}
|
|
PeerWireMessageKind::KeepAlive
|
|
| PeerWireMessageKind::Request(_)
|
|
| PeerWireMessageKind::Cancel(_)
|
|
| PeerWireMessageKind::Port(_)
|
|
| PeerWireMessageKind::Unknown(_) => {}
|
|
}
|
|
}
|
|
|
|
Ok(parsed)
|
|
}
|
|
|
|
/// Computes how many verified bytes overlap a requested byte range.
|
|
pub(in crate::dispatcher) fn verified_length_for_range(
|
|
group: &RequestGroup,
|
|
range_start: u64,
|
|
range_length: u64,
|
|
piece_length: u64,
|
|
total_length: u64,
|
|
) -> u64 {
|
|
if range_length == 0 || piece_length == 0 || total_length == 0 {
|
|
return 0;
|
|
}
|
|
let range_end = range_start.saturating_add(range_length).min(total_length);
|
|
if range_end <= range_start {
|
|
return 0;
|
|
}
|
|
group
|
|
.piece_map()
|
|
.iter()
|
|
.filter(|(_, state)| **state == PieceState::Verified)
|
|
.map(|(piece, _)| {
|
|
let piece_start = u64::from(piece.0).saturating_mul(piece_length);
|
|
let piece_end = piece_start
|
|
.saturating_add(bt_piece_span_bytes(*piece, piece_length, total_length))
|
|
.min(total_length);
|
|
let overlap_start = piece_start.max(range_start);
|
|
let overlap_end = piece_end.min(range_end);
|
|
overlap_end.saturating_sub(overlap_start)
|
|
})
|
|
.sum()
|
|
}
|
|
|
|
/// Returns whether a peer-wire bitfield covers every expected piece index.
|
|
pub(in crate::dispatcher) fn peer_wire_bitfield_is_complete(
|
|
pieces: &BTreeSet<u32>,
|
|
expected_piece_count: usize,
|
|
) -> bool {
|
|
expected_piece_count > 0
|
|
&& pieces.len() >= expected_piece_count
|
|
&& (0..expected_piece_count as u32).all(|piece| pieces.contains(&piece))
|
|
}
|