use super::super::{PeerReporter, ctrl, ctx};
use super::{CLIENT_NAME, LOCAL_REQQ};
use local_async_utils::prelude::*;
use mtorrent_core::{data, pwp};
use std::collections::{HashMap, HashSet};
use std::net::SocketAddr;
use std::time::Duration;
use std::{cmp, io};
use tokio::{select, time};
type CtxHandle = ctx::Handle<ctx::MainCtx>;
struct Data {
handle: CtxHandle,
rx: pwp::ExtendedRxChannel,
tx: pwp::ExtendedTxChannel,
metadata_storage: data::StorageClient,
remote_extensions: HashMap<pwp::Extension, u8>,
sent_metadata_pieces: pwp::Bitfield,
last_shared_peers: HashSet<SocketAddr>,
peer_reporter: PeerReporter,
}
pub struct Peer(Box<Data>);
pub async fn new_peer(
mut handle: CtxHandle,
rx: pwp::ExtendedRxChannel,
tx: pwp::ExtendedTxChannel,
metadata_storage: data::StorageClient,
enabled_extensions: impl IntoIterator<Item = &pwp::Extension>,
peer_reporter: PeerReporter,
) -> io::Result<Peer> {
let metadata_len = handle.with(|ctx| ctx.metainfo.size());
let mut inner = Box::new(Data {
handle,
rx,
tx,
metadata_storage,
remote_extensions: Default::default(),
sent_metadata_pieces: pwp::Bitfield::repeat(false, metadata_len),
last_shared_peers: Default::default(),
peer_reporter,
});
define_with_ctx!(inner.handle);
let local_handshake = Box::new(pwp::ExtendedHandshake {
extensions: enabled_extensions.into_iter().map(|e| (*e, e.local_id())).collect(),
listen_port: Some(with_ctx!(|ctx| ctx.const_data.pwp_external_port())),
client_type: Some(CLIENT_NAME.to_string()),
yourip: Some(inner.rx.remote_ip().ip()),
metadata_size: Some(with_ctx!(|ctx| ctx.metainfo.size())),
request_limit: Some(LOCAL_REQQ),
..Default::default()
});
log::debug!("Sending extended handshake to {}: {}", inner.rx.remote_ip(), local_handshake);
inner
.tx
.send_message((pwp::ExtendedMessage::Handshake(local_handshake), 0))
.await?;
match inner.rx.receive_message_timed(sec!(1)).await {
Ok(pwp::ExtendedMessage::Handshake(mut hs)) => {
log::debug!("Received extended handshake from {}: {}", inner.rx.remote_ip(), hs);
hs.extensions.retain(|_ext, id| *id != 0); inner.remote_extensions.clone_from(&hs.extensions);
with_ctx!(|ctx| ctx.peer_states.set_extended_handshake(inner.rx.remote_ip(), hs));
}
Ok(msg) => {
return Err(io::Error::other(format!("unexpected first extended message: {msg}")));
}
Err(pwp::ChannelError::Timeout) => (),
Err(e) => return Err(e.into()),
}
Ok(Peer(inner))
}
pub async fn run(peer: Peer, pex_interval: Duration) -> io::Result<()> {
let Peer(mut inner) = peer;
let mut pex_timer = time::interval(pex_interval);
loop {
select! {
biased;
received = inner.rx.receive_message() => {
handle_incoming_message(&mut inner, received?).await?;
}
_ = pex_timer.tick() => {
send_pex(&mut inner).await?;
}
}
}
}
async fn send_pex(inner: &mut Data) -> io::Result<()> {
define_with_ctx!(inner.handle);
if let Some(&id) = inner.remote_extensions.get(&pwp::Extension::PeerExchange) {
let connected_peers: HashSet<SocketAddr> =
with_ctx!(|ctx| ctx.peer_states.iter().map(|(ip, _)| *ip).collect());
let pex = Box::new(pwp::PeerExchangeData {
added: connected_peers.difference(&inner.last_shared_peers).cloned().collect(),
dropped: inner.last_shared_peers.difference(&connected_peers).cloned().collect(),
});
if !pex.added.is_empty() || !pex.dropped.is_empty() {
log::debug!("Sending PEX message to {}: {}", inner.tx.remote_ip(), pex);
inner.tx.send_message((pwp::ExtendedMessage::PeerExchange(pex), id)).await?;
inner.last_shared_peers = connected_peers;
}
}
Ok(())
}
async fn handle_incoming_message(inner: &mut Data, msg: pwp::ExtendedMessage) -> io::Result<()> {
define_with_ctx!(inner.handle);
let remote_ip = *inner.rx.remote_ip();
let serve_metadata = with_ctx!(|ctx| ctrl::can_serve_metadata(&remote_ip, ctx));
match msg {
pwp::ExtendedMessage::Handshake(mut hs) => {
log::debug!("Received extended handshake from {remote_ip}: {hs}");
for (ext, id) in &hs.extensions {
if id == &0 {
inner.remote_extensions.remove(ext);
} else {
inner.remote_extensions.insert(*ext, *id);
}
}
hs.extensions.clone_from(&inner.remote_extensions);
with_ctx!(|ctx| ctx.peer_states.set_extended_handshake(&remote_ip, hs));
}
pwp::ExtendedMessage::PeerExchange(pex) => {
log::debug!("Received PEX message from {remote_ip}: {pex}");
if !pex.added.is_empty() {
let connected_peers: HashSet<SocketAddr> =
with_ctx!(|ctx| ctx.peer_states.iter().map(|(ip, _)| *ip).collect());
for peer_addr in pex.added.difference(&connected_peers) {
inner.peer_reporter.report_discovered(*peer_addr, pwp::PeerOrigin::Pex).await;
}
};
}
pwp::ExtendedMessage::MetadataRequest { piece } => {
if let Some(&id) = inner.remote_extensions.get(&pwp::Extension::Metadata) {
if serve_metadata && inner.sent_metadata_pieces.get(piece).is_some_and(|sent| !sent)
{
log::debug!("Serving metadata request from {remote_ip}: (piece={piece})");
let metadata_len = with_ctx!(|ctx| ctx.metainfo.size());
let global_offset = piece * pwp::MAX_BLOCK_SIZE;
let length = cmp::min(pwp::MAX_BLOCK_SIZE, metadata_len - global_offset);
let data = inner.metadata_storage.read_block(global_offset, length).await?;
inner
.tx
.send_message((
pwp::ExtendedMessage::MetadataBlock {
piece,
total_size: metadata_len,
data,
},
id,
))
.await?;
inner.sent_metadata_pieces.set(piece, true);
} else {
log::debug!("Rejecting metadata request from {remote_ip}: (piece={piece})");
inner
.tx
.send_message((pwp::ExtendedMessage::MetadataReject { piece }, id))
.await?;
}
}
}
msg => {
log::warn!("Unexpected extension message from {remote_ip}: {msg}");
}
}
Ok(())
}