mod download;
mod extensions;
mod metadata;
mod tcp;
mod upload;
mod utp;
#[cfg(test)]
mod tests;
#[cfg(test)]
mod testutils;
pub use tcp::run_pwp_listener;
pub use utp::{UtpHandle, launch_utp};
use super::connections::{PeerConnector, PeerReporter};
use super::{ctrl, ctx};
use local_async_utils::prelude::*;
use mtorrent_core::{data, pwp};
use std::io;
use std::net::SocketAddr;
use std::rc::Rc;
use std::time::Duration;
use tokio::sync::broadcast;
use tokio::time::{self, Instant};
use tokio::{runtime, try_join};
use utp::{InboundConnectArgs, OutboundConnectArgs};
type MainHandle = ctx::Handle<ctx::MainCtx>;
type PreliminaryHandle = ctx::Handle<ctx::PreliminaryCtx>;
const EXTENSION_PROTOCOL_ENABLED: bool = true;
const ALL_SUPPORTED_EXTENSIONS: &[pwp::Extension] =
&[pwp::Extension::Metadata, pwp::Extension::PeerExchange];
const CLIENT_NAME: &str = concat!(env!("CARGO_PKG_NAME"), " ", env!("CARGO_PKG_VERSION"));
const LOCAL_REQQ: usize = 1024 * 2;
const PEX_INTERVAL: Duration = sec!(30);
async fn run_download(
mut peer: download::Peer,
remote_ip: SocketAddr,
mut ctx_handle: MainHandle,
) -> io::Result<()> {
define_with_ctx!(ctx_handle);
loop {
with_ctx!(|ctx| ctrl::validate_peer_utility(&remote_ip, ctx))?;
match peer {
download::Peer::Idle(idling_peer) => {
match with_ctx!(|ctx| ctrl::idle_download_next_action(&remote_ip, ctx)) {
ctrl::IdleDownloadAction::ActivateDownload => {
peer = download::activate(idling_peer).await?;
}
ctrl::IdleDownloadAction::WaitForUpdates(timeout) => {
peer =
download::linger(idling_peer.into(), Instant::now() + timeout).await?;
}
}
}
download::Peer::Seeder(seeding_peer) => {
match with_ctx!(|ctx| ctrl::active_download_next_action(&remote_ip, ctx)) {
ctrl::SeederDownloadAction::RequestPieces => {
peer = download::get_pieces(seeding_peer).await?;
}
ctrl::SeederDownloadAction::WaitForUpdates(timeout) => {
peer =
download::linger(seeding_peer.into(), Instant::now() + timeout).await?;
}
ctrl::SeederDownloadAction::DeactivateDownload => {
peer = download::deactivate(seeding_peer).await?.into();
}
}
}
}
}
}
async fn run_upload(
mut peer: upload::Peer,
remote_ip: SocketAddr,
mut ctx_handle: MainHandle,
) -> io::Result<()> {
define_with_ctx!(ctx_handle);
loop {
with_ctx!(|ctx| ctrl::validate_peer_utility(&remote_ip, ctx))?;
match peer {
upload::Peer::Idle(idling_peer) => {
match with_ctx!(|ctx| ctrl::idle_upload_next_action(&remote_ip, &ctx.peer_states)) {
ctrl::IdleUploadAction::ActivateUploadAndServe(duration) => {
let leeching_peer = upload::activate(idling_peer).await?;
peer = upload::serve_pieces(leeching_peer, duration).await?;
}
ctrl::IdleUploadAction::Linger(timeout) => {
peer = upload::linger(idling_peer, timeout).await?;
}
}
}
upload::Peer::Leech(leeching_peer) => {
match with_ctx!(|ctx| ctrl::active_upload_next_action(&remote_ip, &ctx.peer_states))
{
ctrl::LeechUploadAction::DeactivateUploadAndLinger(timeout) => {
let idling_peer = upload::deactivate(leeching_peer).await?;
peer = upload::linger(idling_peer, timeout).await?;
}
ctrl::LeechUploadAction::Serve(duration) => {
peer = upload::serve_pieces(leeching_peer, duration).await?;
}
}
}
}
}
}
async fn run_optional_extensions(peer: Option<extensions::Peer>) -> io::Result<()> {
match peer {
Some(peer) => extensions::run(peer, PEX_INTERVAL).await,
None => Ok(()),
}
}
async fn run_peer_connection(
origin: pwp::PeerOrigin,
transport: pwp::TransportProto,
download_chans: pwp::DownloadChannels,
upload_chans: pwp::UploadChannels,
extended_chans: Option<pwp::ExtendedChannels>,
data: &MainConnectionData,
) -> io::Result<()> {
let remote_ip = *download_chans.0.remote_ip();
let encryption = download_chans.0.is_encrypted();
let pwp::UploadChannels(tx, rx) = upload_chans;
let upload_fut = upload::new_peer(
data.ctx_handle.clone(),
rx,
tx,
data.content_storage.clone(),
data.piece_downloaded_channel.subscribe(),
);
let pwp::DownloadChannels(tx, rx) = download_chans;
let download_fut = download::new_peer(
data.ctx_handle.clone(),
rx,
tx,
data.content_storage.clone(),
data.piece_downloaded_channel.clone(),
);
let handle_copy = data.ctx_handle.clone();
let extensions_fut = async move {
if let Some(pwp::ExtendedChannels(tx, rx)) = extended_chans {
Ok(Some(
extensions::new_peer(
handle_copy,
rx,
tx,
data.metainfo_storage.clone(),
ALL_SUPPORTED_EXTENSIONS,
data.peer_reporter.clone(),
)
.await?,
))
} else {
Ok(None)
}
};
let (extensions, download, (upload, reporter)) =
try_join!(extensions_fut, download_fut, upload_fut)?;
data.ctx_handle
.clone()
.with(|ctx| ctx.peer_states.set_info(&remote_ip, origin, transport, encryption));
try_join!(
run_download(download.into(), remote_ip, data.ctx_handle.clone()),
run_upload(upload.into(), remote_ip, data.ctx_handle.clone()),
run_optional_extensions(extensions),
reporter.run(),
)?;
Ok(())
}
pub struct MainConnectionData {
pub content_storage: data::StorageClient,
pub metainfo_storage: data::StorageClient,
pub ctx_handle: MainHandle,
pub pwp_worker_handle: runtime::Handle,
pub peer_reporter: PeerReporter,
pub piece_downloaded_channel: Rc<broadcast::Sender<usize>>,
pub utp_handle: UtpHandle,
}
impl PeerConnector for MainConnectionData {
type PeerConnection =
(pwp::DownloadChannels, pwp::UploadChannels, Option<pwp::ExtendedChannels>);
fn max_connections(&self) -> usize {
200
}
fn connect_retry_interval(&self) -> Duration {
sec!(15)
}
fn max_connect_retries(&self) -> usize {
2
}
async fn outbound_connect_and_handshake(
&self,
peer_addr: SocketAddr,
use_pe: bool,
deadline: Instant,
) -> io::Result<Self::PeerConnection> {
let mut handle = self.ctx_handle.clone();
define_with_ctx!(handle);
if with_ctx!(|ctx| !ctx.const_data.pwp_outbound_tcp_allowed()) {
return Err(io::Error::from(io::ErrorKind::Unsupported));
}
let (info_hash, data) =
with_ctx!(|ctx| (*ctx.metainfo.info_hash(), ctx.const_data.clone()));
time::timeout_at(
deadline,
tcp::new_outbound_connection(
&data,
&info_hash,
EXTENSION_PROTOCOL_ENABLED,
use_pe,
peer_addr,
&self.pwp_worker_handle,
),
)
.await?
}
async fn outbound_utp_connect_and_handshake(
&self,
peer_addr: SocketAddr,
use_pe: bool,
deadline: Instant,
) -> io::Result<Self::PeerConnection> {
let mut handle = self.ctx_handle.clone();
define_with_ctx!(handle);
if with_ctx!(|ctx| !ctx.const_data.pwp_outbound_utp_allowed()) {
return Err(io::Error::from(io::ErrorKind::Unsupported));
}
let (info_hash, local_peer_id) =
with_ctx!(|ctx| (*ctx.metainfo.info_hash(), *ctx.const_data.local_peer_id()));
self.utp_handle
.outbound_connect(OutboundConnectArgs {
local_peer_id,
info_hash,
extension_protocol_enabled: EXTENSION_PROTOCOL_ENABLED,
protocol_encryption_enabled: use_pe,
peer_addr,
deadline,
})
.await
}
async fn inbound_connect_and_handshake(
&self,
peer_addr: SocketAddr,
deadline: Instant,
stream: tokio::net::TcpStream,
) -> io::Result<Self::PeerConnection> {
let mut handle = self.ctx_handle.clone();
define_with_ctx!(handle);
let (info_hash, local_peer_id) =
with_ctx!(|ctx| (*ctx.metainfo.info_hash(), *ctx.const_data.local_peer_id()));
time::timeout_at(
deadline,
tcp::new_inbound_connection(
&local_peer_id,
&info_hash,
EXTENSION_PROTOCOL_ENABLED,
peer_addr,
stream,
&self.pwp_worker_handle,
),
)
.await?
}
async fn inbound_utp_connect_and_handshake(
&self,
peer_addr: SocketAddr,
deadline: Instant,
data: mtorrent_core::utp::InboundConnectData,
) -> io::Result<Self::PeerConnection> {
let mut handle = self.ctx_handle.clone();
define_with_ctx!(handle);
let (info_hash, local_peer_id) =
with_ctx!(|ctx| (*ctx.metainfo.info_hash(), *ctx.const_data.local_peer_id()));
self.utp_handle
.inbound_connect(InboundConnectArgs {
local_peer_id,
info_hash,
extension_protocol_enabled: EXTENSION_PROTOCOL_ENABLED,
peer_addr,
data,
deadline,
})
.await
}
async fn run_connection(
&self,
origin: pwp::PeerOrigin,
transport: pwp::TransportProto,
connection: Self::PeerConnection,
) -> io::Result<()> {
let (download_chans, upload_chans, extended_chans) = connection;
run_peer_connection(origin, transport, download_chans, upload_chans, extended_chans, self)
.await
}
}
async fn run_metadata_download(
origin: pwp::PeerOrigin,
transport: pwp::TransportProto,
download_chans: pwp::DownloadChannels,
upload_chans: pwp::UploadChannels,
extended_chans: pwp::ExtendedChannels,
mut ctx_handle: PreliminaryHandle,
peer_reporter: PeerReporter,
) -> io::Result<()> {
ctx_handle.with(|ctx| ctx.discovered_peers.insert(*download_chans.0.remote_ip()));
async fn handle_download(download_chans: pwp::DownloadChannels) -> io::Result<()> {
let pwp::DownloadChannels(tx, mut rx) = download_chans;
loop {
let msg = rx.receive_message().await?;
log::trace!("Received {} from {}", msg, tx.remote_ip());
}
}
async fn handle_upload(upload_chans: pwp::UploadChannels) -> io::Result<()> {
let pwp::UploadChannels(tx, mut rx) = upload_chans;
loop {
let msg = rx.receive_message().await?;
log::trace!("Received {} from {}", msg, tx.remote_ip());
}
}
async fn handle_metadata(
origin: pwp::PeerOrigin,
transport: pwp::TransportProto,
extended_chans: pwp::ExtendedChannels,
mut ctx_handle: PreliminaryHandle,
peer_reporter: PeerReporter,
) -> io::Result<()> {
define_with_ctx!(ctx_handle);
let peer_addr = *extended_chans.0.remote_ip();
let encryption = extended_chans.0.is_encrypted();
let mut peer =
metadata::new_peer(ctx_handle.clone(), extended_chans, peer_reporter).await?;
with_ctx!(|ctx| {
ctx.peer_states.set_info(&peer_addr, origin, transport, encryption);
});
while with_ctx!(|ctx| !ctrl::verify_metadata(ctx)) {
match peer {
metadata::Peer::Disabled(disabled) => {
peer = metadata::wait_until_enabled(disabled).await?.into();
}
metadata::Peer::Rejector(rejector) => {
let until = Instant::now() + sec!(10);
peer = metadata::cool_off_rejecting_peer(rejector, until).await?;
}
metadata::Peer::Uploader(uploader) => {
peer = metadata::download_metadata(uploader).await?;
}
}
}
Ok(())
}
try_join!(
handle_download(download_chans),
handle_upload(upload_chans),
handle_metadata(origin, transport, extended_chans, ctx_handle, peer_reporter),
)?;
Ok(())
}
pub struct PreliminaryConnectionData {
pub ctx_handle: PreliminaryHandle,
pub pwp_worker_handle: runtime::Handle,
pub peer_reporter: PeerReporter,
pub utp_handle: UtpHandle,
}
impl PeerConnector for PreliminaryConnectionData {
type PeerConnection = (pwp::DownloadChannels, pwp::UploadChannels, pwp::ExtendedChannels);
fn max_connections(&self) -> usize {
50
}
fn connect_retry_interval(&self) -> Duration {
sec!(5)
}
fn max_connect_retries(&self) -> usize {
1
}
async fn outbound_connect_and_handshake(
&self,
peer_addr: SocketAddr,
use_pe: bool,
deadline: Instant,
) -> io::Result<Self::PeerConnection> {
let mut handle = self.ctx_handle.clone();
define_with_ctx!(handle);
if with_ctx!(|ctx| !ctx.const_data.pwp_outbound_tcp_allowed()) {
return Err(io::Error::from(io::ErrorKind::Unsupported));
}
let (info_hash, data) = with_ctx!(|ctx| (*ctx.magnet.info_hash(), ctx.const_data.clone()));
let (dl, ul, ext) = time::timeout_at(
deadline,
tcp::new_outbound_connection(
&data,
&info_hash,
true, use_pe,
peer_addr,
&self.pwp_worker_handle,
),
)
.await??;
let ext = ext.ok_or_else(|| {
io::Error::new(io::ErrorKind::Unsupported, "peer doesn't support extension protocol")
})?;
Ok((dl, ul, ext))
}
async fn outbound_utp_connect_and_handshake(
&self,
peer_addr: SocketAddr,
use_pe: bool,
deadline: Instant,
) -> io::Result<Self::PeerConnection> {
let mut handle = self.ctx_handle.clone();
define_with_ctx!(handle);
if with_ctx!(|ctx| !ctx.const_data.pwp_outbound_utp_allowed()) {
return Err(io::Error::from(io::ErrorKind::Unsupported));
}
let (info_hash, local_peer_id) =
with_ctx!(|ctx| (*ctx.magnet.info_hash(), *ctx.const_data.local_peer_id()));
let (dl, ul, ext) = self
.utp_handle
.outbound_connect(OutboundConnectArgs {
local_peer_id,
info_hash,
extension_protocol_enabled: true,
protocol_encryption_enabled: use_pe,
peer_addr,
deadline,
})
.await?;
let ext = ext.ok_or_else(|| {
io::Error::new(io::ErrorKind::Unsupported, "peer doesn't support extension protocol")
})?;
Ok((dl, ul, ext))
}
async fn inbound_connect_and_handshake(
&self,
peer_addr: SocketAddr,
deadline: Instant,
stream: tokio::net::TcpStream,
) -> io::Result<Self::PeerConnection> {
let mut handle = self.ctx_handle.clone();
define_with_ctx!(handle);
let (info_hash, local_peer_id) =
with_ctx!(|ctx| (*ctx.magnet.info_hash(), *ctx.const_data.local_peer_id()));
let (dl, ul, ext) = time::timeout_at(
deadline,
tcp::new_inbound_connection(
&local_peer_id,
&info_hash,
true, peer_addr,
stream,
&self.pwp_worker_handle,
),
)
.await??;
let ext = ext.ok_or_else(|| {
io::Error::new(io::ErrorKind::Unsupported, "peer doesn't support extension protocol")
})?;
Ok((dl, ul, ext))
}
async fn inbound_utp_connect_and_handshake(
&self,
peer_addr: SocketAddr,
deadline: Instant,
data: mtorrent_core::utp::InboundConnectData,
) -> io::Result<Self::PeerConnection> {
let mut handle = self.ctx_handle.clone();
define_with_ctx!(handle);
let (info_hash, local_peer_id) =
with_ctx!(|ctx| (*ctx.magnet.info_hash(), *ctx.const_data.local_peer_id()));
let (dl, ul, ext) = self
.utp_handle
.inbound_connect(InboundConnectArgs {
local_peer_id,
info_hash,
extension_protocol_enabled: true,
peer_addr,
data,
deadline,
})
.await?;
let ext = ext.ok_or_else(|| {
io::Error::new(io::ErrorKind::Unsupported, "peer doesn't support extension protocol")
})?;
Ok((dl, ul, ext))
}
async fn run_connection(
&self,
origin: pwp::PeerOrigin,
transport: pwp::TransportProto,
connection: Self::PeerConnection,
) -> io::Result<()> {
let (download_chans, upload_chans, extended_chans) = connection;
run_metadata_download(
origin,
transport,
download_chans,
upload_chans,
extended_chans,
self.ctx_handle.clone(),
self.peer_reporter.clone(),
)
.await
}
}