mirror of
https://github.com/YGGverse/aquatic.git
synced 2026-03-31 17:55:36 +00:00
Move all crates to new crates dir
This commit is contained in:
parent
3835da22ac
commit
9b032f7e24
128 changed files with 27 additions and 26 deletions
56
crates/ws/Cargo.toml
Normal file
56
crates/ws/Cargo.toml
Normal file
|
|
@ -0,0 +1,56 @@
|
|||
[package]
|
||||
name = "aquatic_ws"
|
||||
description = "High-performance open WebTorrent tracker"
|
||||
keywords = ["webtorrent", "websocket", "peer-to-peer", "torrent", "bittorrent"]
|
||||
version.workspace = true
|
||||
authors.workspace = true
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
readme.workspace = true
|
||||
rust-version = "1.70"
|
||||
|
||||
[lib]
|
||||
name = "aquatic_ws"
|
||||
|
||||
[[bin]]
|
||||
name = "aquatic_ws"
|
||||
|
||||
[features]
|
||||
default = ["prometheus"]
|
||||
prometheus = ["metrics", "metrics-exporter-prometheus"]
|
||||
metrics = ["dep:metrics", "metrics-util"]
|
||||
|
||||
[dependencies]
|
||||
aquatic_common = { workspace = true, features = ["rustls", "glommio"] }
|
||||
aquatic_peer_id.workspace = true
|
||||
aquatic_toml_config.workspace = true
|
||||
aquatic_ws_protocol.workspace = true
|
||||
|
||||
anyhow = "1"
|
||||
async-tungstenite = "0.23"
|
||||
cfg-if = "1"
|
||||
futures = "0.3"
|
||||
futures-lite = "1"
|
||||
futures-rustls = "0.24"
|
||||
glommio = "0.8"
|
||||
hashbrown = { version = "0.14", features = ["serde"] }
|
||||
httparse = "1"
|
||||
log = "0.4"
|
||||
metrics = { version = "0.21", optional = true }
|
||||
metrics-util = { version = "0.15", optional = true }
|
||||
metrics-exporter-prometheus = { version = "0.12", optional = true, default-features = false, features = ["http-listener"] }
|
||||
mimalloc = { version = "0.1", default-features = false }
|
||||
privdrop = "0.5"
|
||||
rand = { version = "0.8", features = ["small_rng"] }
|
||||
rustls = "0.21"
|
||||
rustls-pemfile = "1"
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
signal-hook = { version = "0.3" }
|
||||
slab = "0.4"
|
||||
socket2 = { version = "0.5", features = ["all"] }
|
||||
tungstenite = "0.20"
|
||||
|
||||
[dev-dependencies]
|
||||
quickcheck = "1"
|
||||
quickcheck_macros = "1"
|
||||
76
crates/ws/src/common.rs
Normal file
76
crates/ws/src/common.rs
Normal file
|
|
@ -0,0 +1,76 @@
|
|||
use std::{net::IpAddr, sync::Arc};
|
||||
|
||||
use aquatic_common::access_list::AccessListArcSwap;
|
||||
|
||||
pub use aquatic_common::ValidUntil;
|
||||
use aquatic_ws_protocol::{InfoHash, PeerId};
|
||||
|
||||
#[derive(Copy, Clone, Debug)]
|
||||
pub enum IpVersion {
|
||||
V4,
|
||||
V6,
|
||||
}
|
||||
|
||||
impl IpVersion {
|
||||
pub fn canonical_from_ip(ip: IpAddr) -> IpVersion {
|
||||
match ip {
|
||||
IpAddr::V4(_) => Self::V4,
|
||||
IpAddr::V6(addr) => match addr.octets() {
|
||||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0xff, 0xff, _, _, _, _] => Self::V4,
|
||||
_ => Self::V6,
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default, Clone)]
|
||||
pub struct State {
|
||||
pub access_list: Arc<AccessListArcSwap>,
|
||||
}
|
||||
|
||||
#[derive(Copy, Clone, Debug)]
|
||||
pub struct PendingScrapeId(pub u8);
|
||||
|
||||
#[derive(Copy, Clone, Debug)]
|
||||
pub struct ConsumerId(pub u8);
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq)]
|
||||
pub struct ConnectionId(pub usize);
|
||||
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
pub struct InMessageMeta {
|
||||
/// Index of socket worker responsible for this connection. Required for
|
||||
/// sending back response through correct channel to correct worker.
|
||||
pub out_message_consumer_id: ConsumerId,
|
||||
pub connection_id: ConnectionId,
|
||||
pub ip_version: IpVersion,
|
||||
pub pending_scrape_id: Option<PendingScrapeId>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
pub struct OutMessageMeta {
|
||||
/// Index of socket worker responsible for this connection. Required for
|
||||
/// sending back response through correct channel to correct worker.
|
||||
pub out_message_consumer_id: ConsumerId,
|
||||
pub connection_id: ConnectionId,
|
||||
pub pending_scrape_id: Option<PendingScrapeId>,
|
||||
}
|
||||
|
||||
impl Into<OutMessageMeta> for InMessageMeta {
|
||||
fn into(self) -> OutMessageMeta {
|
||||
OutMessageMeta {
|
||||
out_message_consumer_id: self.out_message_consumer_id,
|
||||
connection_id: self.connection_id,
|
||||
pending_scrape_id: self.pending_scrape_id,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
pub enum SwarmControlMessage {
|
||||
ConnectionClosed {
|
||||
info_hash: InfoHash,
|
||||
peer_id: PeerId,
|
||||
ip_version: IpVersion,
|
||||
},
|
||||
}
|
||||
191
crates/ws/src/config.rs
Normal file
191
crates/ws/src/config.rs
Normal file
|
|
@ -0,0 +1,191 @@
|
|||
use std::net::SocketAddr;
|
||||
use std::path::PathBuf;
|
||||
|
||||
use aquatic_common::cpu_pinning::asc::CpuPinningConfigAsc;
|
||||
use aquatic_common::{access_list::AccessListConfig, privileges::PrivilegeConfig};
|
||||
use serde::Deserialize;
|
||||
|
||||
use aquatic_common::cli::LogLevel;
|
||||
use aquatic_toml_config::TomlConfig;
|
||||
|
||||
/// aquatic_ws configuration
|
||||
///
|
||||
/// Running behind a reverse proxy is supported, but IPv4 peer requests have
|
||||
/// to be proxied to IPv4 requests, and IPv6 requests to IPv6 requests.
|
||||
#[derive(Clone, Debug, PartialEq, TomlConfig, Deserialize)]
|
||||
#[serde(default, deny_unknown_fields)]
|
||||
pub struct Config {
|
||||
/// Socket workers receive requests from the socket, parse them and send
|
||||
/// them on to the swarm workers. They then receive responses from the
|
||||
/// swarm workers, encode them and send them back over the socket.
|
||||
pub socket_workers: usize,
|
||||
/// Swarm workers receive a number of requests from socket workers,
|
||||
/// generate responses and send them back to the socket workers.
|
||||
pub swarm_workers: usize,
|
||||
pub log_level: LogLevel,
|
||||
pub network: NetworkConfig,
|
||||
pub protocol: ProtocolConfig,
|
||||
pub cleaning: CleaningConfig,
|
||||
pub privileges: PrivilegeConfig,
|
||||
pub access_list: AccessListConfig,
|
||||
#[cfg(feature = "metrics")]
|
||||
pub metrics: MetricsConfig,
|
||||
pub cpu_pinning: CpuPinningConfigAsc,
|
||||
}
|
||||
|
||||
impl Default for Config {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
socket_workers: 1,
|
||||
swarm_workers: 1,
|
||||
log_level: LogLevel::default(),
|
||||
network: NetworkConfig::default(),
|
||||
protocol: ProtocolConfig::default(),
|
||||
cleaning: CleaningConfig::default(),
|
||||
privileges: PrivilegeConfig::default(),
|
||||
access_list: AccessListConfig::default(),
|
||||
#[cfg(feature = "metrics")]
|
||||
metrics: Default::default(),
|
||||
cpu_pinning: Default::default(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl aquatic_common::cli::Config for Config {
|
||||
fn get_log_level(&self) -> Option<LogLevel> {
|
||||
Some(self.log_level)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, TomlConfig, Deserialize)]
|
||||
#[serde(default, deny_unknown_fields)]
|
||||
pub struct NetworkConfig {
|
||||
/// Bind to this address
|
||||
pub address: SocketAddr,
|
||||
/// Only allow access over IPv6
|
||||
pub only_ipv6: bool,
|
||||
/// Maximum number of pending TCP connections
|
||||
pub tcp_backlog: i32,
|
||||
|
||||
/// Enable TLS
|
||||
pub enable_tls: bool,
|
||||
/// Path to TLS certificate (DER-encoded X.509)
|
||||
pub tls_certificate_path: PathBuf,
|
||||
/// Path to TLS private key (DER-encoded ASN.1 in PKCS#8 or PKCS#1 format)
|
||||
pub tls_private_key_path: PathBuf,
|
||||
|
||||
pub websocket_max_message_size: usize,
|
||||
pub websocket_max_frame_size: usize,
|
||||
pub websocket_write_buffer_size: usize,
|
||||
|
||||
/// Return a HTTP 200 Ok response when receiving GET /health. Can not be
|
||||
/// combined with enable_tls.
|
||||
pub enable_http_health_checks: bool,
|
||||
}
|
||||
|
||||
impl Default for NetworkConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
address: SocketAddr::from(([0, 0, 0, 0], 3000)),
|
||||
only_ipv6: false,
|
||||
tcp_backlog: 1024,
|
||||
|
||||
enable_tls: false,
|
||||
tls_certificate_path: "".into(),
|
||||
tls_private_key_path: "".into(),
|
||||
|
||||
websocket_max_message_size: 64 * 1024,
|
||||
websocket_max_frame_size: 16 * 1024,
|
||||
websocket_write_buffer_size: 8 * 1024,
|
||||
|
||||
enable_http_health_checks: false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, TomlConfig, Deserialize)]
|
||||
#[serde(default, deny_unknown_fields)]
|
||||
pub struct ProtocolConfig {
|
||||
/// Maximum number of torrents to accept in scrape request
|
||||
pub max_scrape_torrents: usize,
|
||||
/// Maximum number of offers to accept in announce request
|
||||
pub max_offers: usize,
|
||||
/// Ask peers to announce this often (seconds)
|
||||
pub peer_announce_interval: usize,
|
||||
}
|
||||
|
||||
impl Default for ProtocolConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
max_scrape_torrents: 255,
|
||||
max_offers: 10,
|
||||
peer_announce_interval: 120,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, TomlConfig, Deserialize)]
|
||||
#[serde(default, deny_unknown_fields)]
|
||||
pub struct CleaningConfig {
|
||||
/// Clean peers this often (seconds)
|
||||
pub torrent_cleaning_interval: u64,
|
||||
/// Remove peers that have not announced for this long (seconds)
|
||||
pub max_peer_age: u32,
|
||||
// Clean connections this often (seconds)
|
||||
pub connection_cleaning_interval: u64,
|
||||
/// Close connections if no responses have been sent to them for this long (seconds)
|
||||
pub max_connection_idle: u32,
|
||||
}
|
||||
|
||||
impl Default for CleaningConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
torrent_cleaning_interval: 30,
|
||||
max_peer_age: 1800,
|
||||
max_connection_idle: 60 * 5,
|
||||
connection_cleaning_interval: 30,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "metrics")]
|
||||
#[derive(Clone, Debug, PartialEq, TomlConfig, Deserialize)]
|
||||
#[serde(default, deny_unknown_fields)]
|
||||
pub struct MetricsConfig {
|
||||
/// Run a prometheus endpoint
|
||||
pub run_prometheus_endpoint: bool,
|
||||
/// Address to run prometheus endpoint on
|
||||
pub prometheus_endpoint_address: SocketAddr,
|
||||
/// Update metrics for torrent count this often (seconds)
|
||||
pub torrent_count_update_interval: u64,
|
||||
/// Serve information on peer clients
|
||||
///
|
||||
/// Expect a certain CPU hit
|
||||
pub peer_clients: bool,
|
||||
/// Serve information on all peer id prefixes
|
||||
///
|
||||
/// Requires `peer_clients` to be activated.
|
||||
///
|
||||
/// Expect a certain CPU hit
|
||||
pub peer_id_prefixes: bool,
|
||||
}
|
||||
|
||||
#[cfg(feature = "metrics")]
|
||||
impl Default for MetricsConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
run_prometheus_endpoint: false,
|
||||
prometheus_endpoint_address: SocketAddr::from(([0, 0, 0, 0], 9000)),
|
||||
torrent_count_update_interval: 10,
|
||||
peer_clients: false,
|
||||
peer_id_prefixes: false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::Config;
|
||||
|
||||
::aquatic_toml_config::gen_serialize_deserialize_test!(Config);
|
||||
}
|
||||
197
crates/ws/src/lib.rs
Normal file
197
crates/ws/src/lib.rs
Normal file
|
|
@ -0,0 +1,197 @@
|
|||
pub mod common;
|
||||
pub mod config;
|
||||
pub mod workers;
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use anyhow::Context;
|
||||
use aquatic_common::cpu_pinning::glommio::{get_worker_placement, set_affinity_for_util_worker};
|
||||
use aquatic_common::cpu_pinning::WorkerIndex;
|
||||
use aquatic_common::rustls_config::create_rustls_config;
|
||||
use aquatic_common::{PanicSentinelWatcher, ServerStartInstant};
|
||||
use glommio::{channels::channel_mesh::MeshBuilder, prelude::*};
|
||||
use signal_hook::{
|
||||
consts::{SIGTERM, SIGUSR1},
|
||||
iterator::Signals,
|
||||
};
|
||||
|
||||
use aquatic_common::access_list::update_access_list;
|
||||
use aquatic_common::privileges::PrivilegeDropper;
|
||||
|
||||
use common::*;
|
||||
use config::Config;
|
||||
|
||||
pub const APP_NAME: &str = "aquatic_ws: WebTorrent tracker";
|
||||
pub const APP_VERSION: &str = env!("CARGO_PKG_VERSION");
|
||||
|
||||
pub const SHARED_IN_CHANNEL_SIZE: usize = 1024;
|
||||
|
||||
pub fn run(config: Config) -> ::anyhow::Result<()> {
|
||||
if config.network.enable_tls && config.network.enable_http_health_checks {
|
||||
return Err(anyhow::anyhow!(
|
||||
"configuration: network.enable_tls and network.enable_http_health_check can't both be set to true"
|
||||
));
|
||||
}
|
||||
|
||||
let mut signals = Signals::new([SIGUSR1, SIGTERM])?;
|
||||
|
||||
#[cfg(feature = "prometheus")]
|
||||
if config.metrics.run_prometheus_endpoint {
|
||||
use metrics_exporter_prometheus::PrometheusBuilder;
|
||||
|
||||
let idle_timeout = config
|
||||
.cleaning
|
||||
.connection_cleaning_interval
|
||||
.max(config.cleaning.torrent_cleaning_interval)
|
||||
.max(config.metrics.torrent_count_update_interval)
|
||||
* 2;
|
||||
|
||||
PrometheusBuilder::new()
|
||||
.idle_timeout(
|
||||
metrics_util::MetricKindMask::GAUGE,
|
||||
Some(Duration::from_secs(idle_timeout)),
|
||||
)
|
||||
.with_http_listener(config.metrics.prometheus_endpoint_address)
|
||||
.install()
|
||||
.with_context(|| {
|
||||
format!(
|
||||
"Install prometheus endpoint on {}",
|
||||
config.metrics.prometheus_endpoint_address
|
||||
)
|
||||
})?;
|
||||
}
|
||||
|
||||
let state = State::default();
|
||||
|
||||
update_access_list(&config.access_list, &state.access_list)?;
|
||||
|
||||
let num_peers = config.socket_workers + config.swarm_workers;
|
||||
|
||||
let request_mesh_builder = MeshBuilder::partial(num_peers, SHARED_IN_CHANNEL_SIZE);
|
||||
let response_mesh_builder = MeshBuilder::partial(num_peers, SHARED_IN_CHANNEL_SIZE * 16);
|
||||
let control_mesh_builder = MeshBuilder::partial(num_peers, SHARED_IN_CHANNEL_SIZE);
|
||||
|
||||
let (sentinel_watcher, sentinel) = PanicSentinelWatcher::create_with_sentinel();
|
||||
let priv_dropper = PrivilegeDropper::new(config.privileges.clone(), config.socket_workers);
|
||||
|
||||
let opt_tls_config = if config.network.enable_tls {
|
||||
Some(Arc::new(
|
||||
create_rustls_config(
|
||||
&config.network.tls_certificate_path,
|
||||
&config.network.tls_private_key_path,
|
||||
)
|
||||
.with_context(|| "create rustls config")?,
|
||||
))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let server_start_instant = ServerStartInstant::new();
|
||||
|
||||
let mut executors = Vec::new();
|
||||
|
||||
for i in 0..(config.socket_workers) {
|
||||
let sentinel = sentinel.clone();
|
||||
let config = config.clone();
|
||||
let state = state.clone();
|
||||
let opt_tls_config = opt_tls_config.clone();
|
||||
let control_mesh_builder = control_mesh_builder.clone();
|
||||
let request_mesh_builder = request_mesh_builder.clone();
|
||||
let response_mesh_builder = response_mesh_builder.clone();
|
||||
let priv_dropper = priv_dropper.clone();
|
||||
|
||||
let placement = get_worker_placement(
|
||||
&config.cpu_pinning,
|
||||
config.socket_workers,
|
||||
config.swarm_workers,
|
||||
WorkerIndex::SocketWorker(i),
|
||||
)?;
|
||||
let builder = LocalExecutorBuilder::new(placement).name(&format!("socket-{:02}", i + 1));
|
||||
|
||||
let executor = builder
|
||||
.spawn(move || async move {
|
||||
workers::socket::run_socket_worker(
|
||||
sentinel,
|
||||
config,
|
||||
state,
|
||||
opt_tls_config,
|
||||
control_mesh_builder,
|
||||
request_mesh_builder,
|
||||
response_mesh_builder,
|
||||
priv_dropper,
|
||||
server_start_instant,
|
||||
i,
|
||||
)
|
||||
.await
|
||||
})
|
||||
.map_err(|err| anyhow::anyhow!("Spawning executor failed: {:#}", err))?;
|
||||
|
||||
executors.push(executor);
|
||||
}
|
||||
|
||||
::log::info!("spawned socket workers");
|
||||
|
||||
for i in 0..(config.swarm_workers) {
|
||||
let sentinel = sentinel.clone();
|
||||
let config = config.clone();
|
||||
let state = state.clone();
|
||||
let control_mesh_builder = control_mesh_builder.clone();
|
||||
let request_mesh_builder = request_mesh_builder.clone();
|
||||
let response_mesh_builder = response_mesh_builder.clone();
|
||||
|
||||
let placement = get_worker_placement(
|
||||
&config.cpu_pinning,
|
||||
config.socket_workers,
|
||||
config.swarm_workers,
|
||||
WorkerIndex::SwarmWorker(i),
|
||||
)?;
|
||||
let builder = LocalExecutorBuilder::new(placement).name(&format!("swarm-{:02}", i + 1));
|
||||
|
||||
let executor = builder
|
||||
.spawn(move || async move {
|
||||
workers::swarm::run_swarm_worker(
|
||||
sentinel,
|
||||
config,
|
||||
state,
|
||||
control_mesh_builder,
|
||||
request_mesh_builder,
|
||||
response_mesh_builder,
|
||||
server_start_instant,
|
||||
i,
|
||||
)
|
||||
.await
|
||||
})
|
||||
.map_err(|err| anyhow::anyhow!("Spawning executor failed: {:#}", err))?;
|
||||
|
||||
executors.push(executor);
|
||||
}
|
||||
|
||||
::log::info!("spawned swarm workers");
|
||||
|
||||
if config.cpu_pinning.active {
|
||||
set_affinity_for_util_worker(
|
||||
&config.cpu_pinning,
|
||||
config.socket_workers,
|
||||
config.swarm_workers,
|
||||
)?;
|
||||
}
|
||||
|
||||
for signal in &mut signals {
|
||||
match signal {
|
||||
SIGUSR1 => {
|
||||
let _ = update_access_list(&config.access_list, &state.access_list);
|
||||
}
|
||||
SIGTERM => {
|
||||
if sentinel_watcher.panic_was_triggered() {
|
||||
return Err(anyhow::anyhow!("worker thread panicked"));
|
||||
} else {
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
_ => unreachable!(),
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
14
crates/ws/src/main.rs
Normal file
14
crates/ws/src/main.rs
Normal file
|
|
@ -0,0 +1,14 @@
|
|||
use aquatic_common::cli::run_app_with_cli_and_config;
|
||||
use aquatic_ws::config::Config;
|
||||
|
||||
#[global_allocator]
|
||||
static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
|
||||
|
||||
fn main() {
|
||||
run_app_with_cli_and_config::<Config>(
|
||||
aquatic_ws::APP_NAME,
|
||||
aquatic_ws::APP_VERSION,
|
||||
aquatic_ws::run,
|
||||
None,
|
||||
)
|
||||
}
|
||||
2
crates/ws/src/workers/mod.rs
Normal file
2
crates/ws/src/workers/mod.rs
Normal file
|
|
@ -0,0 +1,2 @@
|
|||
pub mod socket;
|
||||
pub mod swarm;
|
||||
983
crates/ws/src/workers/socket.rs
Normal file
983
crates/ws/src/workers/socket.rs
Normal file
|
|
@ -0,0 +1,983 @@
|
|||
use std::borrow::Cow;
|
||||
use std::cell::RefCell;
|
||||
use std::collections::BTreeMap;
|
||||
use std::os::unix::prelude::{FromRawFd, IntoRawFd};
|
||||
use std::rc::Rc;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use anyhow::Context;
|
||||
use aquatic_common::access_list::{create_access_list_cache, AccessListArcSwap, AccessListCache};
|
||||
use aquatic_common::privileges::PrivilegeDropper;
|
||||
use aquatic_common::rustls_config::RustlsConfig;
|
||||
use aquatic_common::{PanicSentinel, ServerStartInstant};
|
||||
use aquatic_peer_id::PeerClient;
|
||||
use aquatic_ws_protocol::*;
|
||||
use async_tungstenite::WebSocketStream;
|
||||
use futures::stream::{SplitSink, SplitStream};
|
||||
use futures::{AsyncWriteExt, StreamExt};
|
||||
use futures_lite::future::race;
|
||||
use futures_rustls::TlsAcceptor;
|
||||
use glommio::channels::channel_mesh::{MeshBuilder, Partial, Role, Senders};
|
||||
use glommio::channels::local_channel::{new_bounded, LocalReceiver, LocalSender};
|
||||
use glommio::channels::shared_channel::ConnectedReceiver;
|
||||
use glommio::net::{TcpListener, TcpStream};
|
||||
use glommio::task::JoinHandle;
|
||||
use glommio::timer::{sleep, timeout, TimerActionRepeat};
|
||||
use glommio::{enclose, prelude::*};
|
||||
use hashbrown::hash_map::Entry;
|
||||
use hashbrown::HashMap;
|
||||
use slab::Slab;
|
||||
|
||||
use crate::config::Config;
|
||||
|
||||
use crate::common::*;
|
||||
|
||||
const LOCAL_CHANNEL_SIZE: usize = 16;
|
||||
|
||||
#[cfg(feature = "metrics")]
|
||||
thread_local! { static WORKER_INDEX: ::std::cell::Cell<usize> = Default::default() }
|
||||
|
||||
struct PendingScrapeResponse {
|
||||
pending_worker_out_messages: usize,
|
||||
stats: HashMap<InfoHash, ScrapeStatistics>,
|
||||
}
|
||||
|
||||
struct ConnectionReference {
|
||||
task_handle: Option<JoinHandle<()>>,
|
||||
/// Sender part of channel used to pass on outgoing messages from request
|
||||
/// worker
|
||||
out_message_sender: Rc<LocalSender<(OutMessageMeta, OutMessage)>>,
|
||||
/// Updated after sending message to peer
|
||||
valid_until: ValidUntil,
|
||||
announced_info_hashes: HashMap<InfoHash, PeerId>,
|
||||
ip_version: IpVersion,
|
||||
opt_peer_client: Option<(PeerClient, String)>,
|
||||
}
|
||||
|
||||
pub async fn run_socket_worker(
|
||||
_sentinel: PanicSentinel,
|
||||
config: Config,
|
||||
state: State,
|
||||
opt_tls_config: Option<Arc<RustlsConfig>>,
|
||||
control_message_mesh_builder: MeshBuilder<SwarmControlMessage, Partial>,
|
||||
in_message_mesh_builder: MeshBuilder<(InMessageMeta, InMessage), Partial>,
|
||||
out_message_mesh_builder: MeshBuilder<(OutMessageMeta, OutMessage), Partial>,
|
||||
priv_dropper: PrivilegeDropper,
|
||||
server_start_instant: ServerStartInstant,
|
||||
worker_index: usize,
|
||||
) {
|
||||
#[cfg(feature = "metrics")]
|
||||
WORKER_INDEX.with(|index| index.set(worker_index));
|
||||
|
||||
let config = Rc::new(config);
|
||||
let access_list = state.access_list;
|
||||
|
||||
let listener = create_tcp_listener(&config, priv_dropper).expect("create tcp listener");
|
||||
|
||||
::log::info!("created tcp listener");
|
||||
|
||||
let (control_message_senders, _) = control_message_mesh_builder
|
||||
.join(Role::Producer)
|
||||
.await
|
||||
.unwrap();
|
||||
let control_message_senders = Rc::new(control_message_senders);
|
||||
|
||||
let (in_message_senders, _) = in_message_mesh_builder.join(Role::Producer).await.unwrap();
|
||||
let in_message_senders = Rc::new(in_message_senders);
|
||||
|
||||
let tq_prioritized = executor().create_task_queue(
|
||||
Shares::Static(100),
|
||||
Latency::Matters(Duration::from_millis(1)),
|
||||
"prioritized",
|
||||
);
|
||||
let tq_regular =
|
||||
executor().create_task_queue(Shares::Static(1), Latency::NotImportant, "regular");
|
||||
|
||||
let (_, mut out_message_receivers) =
|
||||
out_message_mesh_builder.join(Role::Consumer).await.unwrap();
|
||||
let out_message_consumer_id = ConsumerId(
|
||||
out_message_receivers
|
||||
.consumer_id()
|
||||
.unwrap()
|
||||
.try_into()
|
||||
.unwrap(),
|
||||
);
|
||||
|
||||
::log::info!("joined channels");
|
||||
|
||||
let connection_slab = Rc::new(RefCell::new(Slab::new()));
|
||||
|
||||
// Periodically clean connections
|
||||
TimerActionRepeat::repeat_into(
|
||||
enclose!((config, connection_slab) move || {
|
||||
clean_connections(
|
||||
config.clone(),
|
||||
connection_slab.clone(),
|
||||
server_start_instant,
|
||||
)
|
||||
}),
|
||||
tq_prioritized,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
for (_, out_message_receiver) in out_message_receivers.streams() {
|
||||
spawn_local_into(
|
||||
receive_out_messages(out_message_receiver, connection_slab.clone()),
|
||||
tq_regular,
|
||||
)
|
||||
.unwrap()
|
||||
.detach();
|
||||
}
|
||||
|
||||
let mut incoming = listener.incoming();
|
||||
|
||||
while let Some(stream) = incoming.next().await {
|
||||
match stream {
|
||||
Ok(stream) => {
|
||||
let ip_version = match stream.peer_addr() {
|
||||
Ok(addr) => IpVersion::canonical_from_ip(addr.ip()),
|
||||
Err(err) => {
|
||||
::log::info!("could not extract ip version (v4 or v6): {:#}", err);
|
||||
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
let (out_message_sender, out_message_receiver) = new_bounded(LOCAL_CHANNEL_SIZE);
|
||||
let out_message_sender = Rc::new(out_message_sender);
|
||||
|
||||
let key = RefCell::borrow_mut(&connection_slab).insert(ConnectionReference {
|
||||
task_handle: None,
|
||||
out_message_sender: out_message_sender.clone(),
|
||||
valid_until: ValidUntil::new(
|
||||
server_start_instant,
|
||||
config.cleaning.max_connection_idle,
|
||||
),
|
||||
announced_info_hashes: Default::default(),
|
||||
ip_version,
|
||||
opt_peer_client: None,
|
||||
});
|
||||
|
||||
::log::trace!("accepting stream, assigning id {}", key);
|
||||
|
||||
let task_handle = spawn_local_into(enclose!((config, access_list, control_message_senders, in_message_senders, connection_slab, opt_tls_config) async move {
|
||||
#[cfg(feature = "metrics")]
|
||||
::metrics::increment_gauge!(
|
||||
"aquatic_active_connections",
|
||||
1.0,
|
||||
"ip_version" => ip_version_to_metrics_str(ip_version),
|
||||
"worker_index" => worker_index.to_string(),
|
||||
);
|
||||
|
||||
if let Err(err) = run_connection(
|
||||
config.clone(),
|
||||
access_list,
|
||||
in_message_senders,
|
||||
tq_prioritized,
|
||||
tq_regular,
|
||||
connection_slab.clone(),
|
||||
out_message_sender,
|
||||
out_message_receiver,
|
||||
server_start_instant,
|
||||
out_message_consumer_id,
|
||||
ConnectionId(key),
|
||||
opt_tls_config,
|
||||
ip_version,
|
||||
stream,
|
||||
).await {
|
||||
::log::debug!("connection error: {:#}", err);
|
||||
}
|
||||
|
||||
// Clean up after closed connection
|
||||
|
||||
#[cfg(feature = "metrics")]
|
||||
::metrics::decrement_gauge!(
|
||||
"aquatic_active_connections",
|
||||
1.0,
|
||||
"ip_version" => ip_version_to_metrics_str(ip_version),
|
||||
"worker_index" => worker_index.to_string(),
|
||||
);
|
||||
|
||||
// Remove reference in separate statement to avoid
|
||||
// multiple RefCell borrows
|
||||
let opt_reference = connection_slab.borrow_mut().try_remove(key);
|
||||
|
||||
// Tell swarm workers to remove peer
|
||||
if let Some(reference) = opt_reference {
|
||||
for (info_hash, peer_id) in reference.announced_info_hashes {
|
||||
let message = SwarmControlMessage::ConnectionClosed {
|
||||
info_hash,
|
||||
peer_id,
|
||||
ip_version: reference.ip_version,
|
||||
};
|
||||
|
||||
let consumer_index =
|
||||
calculate_in_message_consumer_index(&config, info_hash);
|
||||
|
||||
// Only fails when receiver is closed
|
||||
control_message_senders
|
||||
.send_to(
|
||||
consumer_index,
|
||||
message
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[cfg(feature = "prometheus")]
|
||||
if let Some((peer_client, prefix)) = reference.opt_peer_client {
|
||||
::metrics::decrement_gauge!(
|
||||
"aquatic_peer_clients",
|
||||
1.0,
|
||||
"client" => peer_client.to_string(),
|
||||
);
|
||||
|
||||
if config.metrics.peer_id_prefixes {
|
||||
::metrics::decrement_gauge!(
|
||||
"aquatic_peer_id_prefixes",
|
||||
1.0,
|
||||
"prefix_hex" => prefix.to_string(),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}), tq_regular)
|
||||
.unwrap()
|
||||
.detach();
|
||||
|
||||
if let Some(reference) = connection_slab.borrow_mut().get_mut(key) {
|
||||
reference.task_handle = Some(task_handle);
|
||||
}
|
||||
}
|
||||
Err(err) => {
|
||||
::log::error!("accept connection: {:#}", err);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn clean_connections(
|
||||
config: Rc<Config>,
|
||||
connection_slab: Rc<RefCell<Slab<ConnectionReference>>>,
|
||||
server_start_instant: ServerStartInstant,
|
||||
) -> Option<Duration> {
|
||||
let now = server_start_instant.seconds_elapsed();
|
||||
|
||||
connection_slab.borrow_mut().retain(|_, reference| {
|
||||
if reference.valid_until.valid(now) {
|
||||
#[cfg(feature = "prometheus")]
|
||||
if let Some((peer_client, prefix)) = &reference.opt_peer_client {
|
||||
// As long as connection is still alive, increment peer client
|
||||
// gauges by zero to prevent them from being removed due to
|
||||
// idleness
|
||||
|
||||
::metrics::increment_gauge!(
|
||||
"aquatic_peer_clients",
|
||||
0.0,
|
||||
"client" => peer_client.to_string(),
|
||||
);
|
||||
|
||||
if config.metrics.peer_id_prefixes {
|
||||
::metrics::increment_gauge!(
|
||||
"aquatic_peer_id_prefixes",
|
||||
0.0,
|
||||
"prefix_hex" => prefix.to_string(),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
true
|
||||
} else {
|
||||
if let Some(ref handle) = reference.task_handle {
|
||||
handle.cancel();
|
||||
}
|
||||
|
||||
false
|
||||
}
|
||||
});
|
||||
|
||||
connection_slab.borrow_mut().shrink_to_fit();
|
||||
|
||||
#[cfg(feature = "metrics")]
|
||||
{
|
||||
// Increment gauges by zero to prevent them from being removed due to
|
||||
// idleness
|
||||
|
||||
let worker_index = WORKER_INDEX.with(|index| index.get()).to_string();
|
||||
|
||||
if config.network.address.is_ipv4() || !config.network.only_ipv6 {
|
||||
::metrics::increment_gauge!(
|
||||
"aquatic_active_connections",
|
||||
0.0,
|
||||
"ip_version" => "4",
|
||||
"worker_index" => worker_index.clone(),
|
||||
);
|
||||
}
|
||||
if config.network.address.is_ipv6() {
|
||||
::metrics::increment_gauge!(
|
||||
"aquatic_active_connections",
|
||||
0.0,
|
||||
"ip_version" => "6",
|
||||
"worker_index" => worker_index,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
Some(Duration::from_secs(
|
||||
config.cleaning.connection_cleaning_interval,
|
||||
))
|
||||
}
|
||||
|
||||
async fn receive_out_messages(
|
||||
mut out_message_receiver: ConnectedReceiver<(OutMessageMeta, OutMessage)>,
|
||||
connection_references: Rc<RefCell<Slab<ConnectionReference>>>,
|
||||
) {
|
||||
let connection_references = &connection_references;
|
||||
|
||||
while let Some((meta, out_message)) = out_message_receiver.next().await {
|
||||
if let Some(reference) = connection_references.borrow().get(meta.connection_id.0) {
|
||||
::log::trace!(
|
||||
"local channel {} len: {}",
|
||||
meta.connection_id.0,
|
||||
reference.out_message_sender.len()
|
||||
);
|
||||
|
||||
match reference.out_message_sender.try_send((meta, out_message)) {
|
||||
Ok(()) => {}
|
||||
Err(GlommioError::Closed(_)) => {}
|
||||
Err(GlommioError::WouldBlock(_)) => {}
|
||||
Err(err) => {
|
||||
::log::debug!(
|
||||
"Couldn't send out_message from shared channel to local receiver: {:?}",
|
||||
err
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn run_connection(
|
||||
config: Rc<Config>,
|
||||
access_list: Arc<AccessListArcSwap>,
|
||||
in_message_senders: Rc<Senders<(InMessageMeta, InMessage)>>,
|
||||
tq_prioritized: TaskQueueHandle,
|
||||
tq_regular: TaskQueueHandle,
|
||||
connection_slab: Rc<RefCell<Slab<ConnectionReference>>>,
|
||||
out_message_sender: Rc<LocalSender<(OutMessageMeta, OutMessage)>>,
|
||||
out_message_receiver: LocalReceiver<(OutMessageMeta, OutMessage)>,
|
||||
server_start_instant: ServerStartInstant,
|
||||
out_message_consumer_id: ConsumerId,
|
||||
connection_id: ConnectionId,
|
||||
opt_tls_config: Option<Arc<RustlsConfig>>,
|
||||
ip_version: IpVersion,
|
||||
mut stream: TcpStream,
|
||||
) -> anyhow::Result<()> {
|
||||
if let Some(tls_config) = opt_tls_config {
|
||||
let tls_acceptor: TlsAcceptor = tls_config.into();
|
||||
|
||||
let stream = tls_acceptor.accept(stream).await?;
|
||||
|
||||
run_stream_agnostic_connection(
|
||||
config.clone(),
|
||||
access_list,
|
||||
in_message_senders,
|
||||
tq_prioritized,
|
||||
tq_regular,
|
||||
connection_slab.clone(),
|
||||
out_message_sender,
|
||||
out_message_receiver,
|
||||
server_start_instant,
|
||||
out_message_consumer_id,
|
||||
connection_id,
|
||||
stream,
|
||||
ip_version,
|
||||
)
|
||||
.await
|
||||
} else {
|
||||
// Implementing this over TLS is too cumbersome, since the crate used
|
||||
// for TLS streams doesn't support peek and tungstenite doesn't
|
||||
// properly support sending a HTTP error response in accept_hdr
|
||||
// callback.
|
||||
if config.network.enable_http_health_checks {
|
||||
let mut peek_buf = [0u8; 11];
|
||||
|
||||
stream
|
||||
.peek(&mut peek_buf)
|
||||
.await
|
||||
.map_err(|err| anyhow::anyhow!("error peeking: {:#}", err))?;
|
||||
|
||||
if &peek_buf == b"GET /health" {
|
||||
stream
|
||||
.write_all(b"HTTP/1.1 200 Ok\r\nContent-Length: 2\r\n\r\nOk")
|
||||
.await
|
||||
.map_err(|err| {
|
||||
anyhow::anyhow!("error sending health check response: {:#}", err)
|
||||
})?;
|
||||
stream.flush().await.map_err(|err| {
|
||||
anyhow::anyhow!("error flushing health check response: {:#}", err)
|
||||
})?;
|
||||
|
||||
return Err(anyhow::anyhow!(
|
||||
"client requested health check, skipping websocket negotiation"
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
run_stream_agnostic_connection(
|
||||
config.clone(),
|
||||
access_list,
|
||||
in_message_senders,
|
||||
tq_prioritized,
|
||||
tq_regular,
|
||||
connection_slab.clone(),
|
||||
out_message_sender,
|
||||
out_message_receiver,
|
||||
server_start_instant,
|
||||
out_message_consumer_id,
|
||||
connection_id,
|
||||
stream,
|
||||
ip_version,
|
||||
)
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
async fn run_stream_agnostic_connection<
|
||||
S: futures::AsyncRead + futures::AsyncWrite + Unpin + 'static,
|
||||
>(
|
||||
config: Rc<Config>,
|
||||
access_list: Arc<AccessListArcSwap>,
|
||||
in_message_senders: Rc<Senders<(InMessageMeta, InMessage)>>,
|
||||
tq_prioritized: TaskQueueHandle,
|
||||
tq_regular: TaskQueueHandle,
|
||||
connection_slab: Rc<RefCell<Slab<ConnectionReference>>>,
|
||||
out_message_sender: Rc<LocalSender<(OutMessageMeta, OutMessage)>>,
|
||||
out_message_receiver: LocalReceiver<(OutMessageMeta, OutMessage)>,
|
||||
server_start_instant: ServerStartInstant,
|
||||
out_message_consumer_id: ConsumerId,
|
||||
connection_id: ConnectionId,
|
||||
stream: S,
|
||||
ip_version: IpVersion,
|
||||
) -> anyhow::Result<()> {
|
||||
let ws_config = tungstenite::protocol::WebSocketConfig {
|
||||
max_frame_size: Some(config.network.websocket_max_frame_size),
|
||||
max_message_size: Some(config.network.websocket_max_message_size),
|
||||
write_buffer_size: config.network.websocket_write_buffer_size,
|
||||
max_write_buffer_size: config.network.websocket_write_buffer_size * 3,
|
||||
..Default::default()
|
||||
};
|
||||
let stream = async_tungstenite::accept_async_with_config(stream, Some(ws_config)).await?;
|
||||
|
||||
let (ws_out, ws_in) = futures::StreamExt::split(stream);
|
||||
|
||||
let pending_scrape_slab = Rc::new(RefCell::new(Slab::new()));
|
||||
let access_list_cache = create_access_list_cache(&access_list);
|
||||
|
||||
let reader_handle = spawn_local_into(
|
||||
enclose!((config, connection_slab, pending_scrape_slab) async move {
|
||||
let mut reader = ConnectionReader {
|
||||
config,
|
||||
access_list_cache,
|
||||
connection_slab,
|
||||
in_message_senders,
|
||||
out_message_sender,
|
||||
pending_scrape_slab,
|
||||
out_message_consumer_id,
|
||||
ws_in,
|
||||
ip_version,
|
||||
connection_id,
|
||||
};
|
||||
|
||||
let result = reader.run_in_message_loop().await;
|
||||
|
||||
result
|
||||
}),
|
||||
tq_regular,
|
||||
)
|
||||
.unwrap()
|
||||
.detach();
|
||||
|
||||
let writer_handle = spawn_local_into(
|
||||
async move {
|
||||
let mut writer = ConnectionWriter {
|
||||
config,
|
||||
out_message_receiver,
|
||||
connection_slab,
|
||||
ws_out,
|
||||
pending_scrape_slab,
|
||||
connection_id,
|
||||
server_start_instant,
|
||||
ip_version,
|
||||
};
|
||||
|
||||
let result = writer.run_out_message_loop().await;
|
||||
|
||||
result
|
||||
},
|
||||
tq_prioritized,
|
||||
)
|
||||
.unwrap()
|
||||
.detach();
|
||||
|
||||
race(reader_handle, writer_handle)
|
||||
.await
|
||||
.ok_or_else(|| anyhow::anyhow!("reader/writer task closed"))?
|
||||
}
|
||||
|
||||
struct ConnectionReader<S> {
|
||||
config: Rc<Config>,
|
||||
access_list_cache: AccessListCache,
|
||||
connection_slab: Rc<RefCell<Slab<ConnectionReference>>>,
|
||||
in_message_senders: Rc<Senders<(InMessageMeta, InMessage)>>,
|
||||
out_message_sender: Rc<LocalSender<(OutMessageMeta, OutMessage)>>,
|
||||
pending_scrape_slab: Rc<RefCell<Slab<PendingScrapeResponse>>>,
|
||||
out_message_consumer_id: ConsumerId,
|
||||
ws_in: SplitStream<WebSocketStream<S>>,
|
||||
ip_version: IpVersion,
|
||||
connection_id: ConnectionId,
|
||||
}
|
||||
|
||||
impl<S: futures::AsyncRead + futures::AsyncWrite + Unpin> ConnectionReader<S> {
|
||||
async fn run_in_message_loop(&mut self) -> anyhow::Result<()> {
|
||||
loop {
|
||||
while self.out_message_sender.is_full() {
|
||||
sleep(Duration::from_millis(100)).await;
|
||||
|
||||
yield_if_needed().await;
|
||||
}
|
||||
|
||||
let message = self
|
||||
.ws_in
|
||||
.next()
|
||||
.await
|
||||
.ok_or_else(|| anyhow::anyhow!("Stream ended"))??;
|
||||
|
||||
match &message {
|
||||
tungstenite::Message::Text(_) | tungstenite::Message::Binary(_) => {
|
||||
match InMessage::from_ws_message(message) {
|
||||
Ok(in_message) => {
|
||||
self.handle_in_message(in_message).await?;
|
||||
}
|
||||
Err(err) => {
|
||||
::log::debug!("Couldn't parse in_message: {:?}", err);
|
||||
|
||||
self.send_error_response("Invalid request".into(), None, None)
|
||||
.await?;
|
||||
}
|
||||
}
|
||||
}
|
||||
tungstenite::Message::Ping(_) => {
|
||||
::log::trace!("Received ping message");
|
||||
// tungstenite sends a pong response by itself
|
||||
}
|
||||
tungstenite::Message::Pong(_) => {
|
||||
::log::trace!("Received pong message");
|
||||
}
|
||||
tungstenite::Message::Close(_) => {
|
||||
::log::debug!("Client sent close frame");
|
||||
|
||||
break Ok(());
|
||||
}
|
||||
tungstenite::Message::Frame(_) => {
|
||||
::log::warn!("Read raw websocket frame, this should not happen");
|
||||
}
|
||||
}
|
||||
|
||||
yield_if_needed().await;
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_in_message(&mut self, in_message: InMessage) -> anyhow::Result<()> {
|
||||
match in_message {
|
||||
InMessage::AnnounceRequest(announce_request) => {
|
||||
#[cfg(feature = "metrics")]
|
||||
::metrics::increment_counter!(
|
||||
"aquatic_requests_total",
|
||||
"type" => "announce",
|
||||
"ip_version" => ip_version_to_metrics_str(self.ip_version),
|
||||
"worker_index" => WORKER_INDEX.with(|index| index.get()).to_string(),
|
||||
);
|
||||
|
||||
let info_hash = announce_request.info_hash;
|
||||
|
||||
if self
|
||||
.access_list_cache
|
||||
.load()
|
||||
.allows(self.config.access_list.mode, &info_hash.0)
|
||||
{
|
||||
{
|
||||
let mut connection_slab = self.connection_slab.borrow_mut();
|
||||
|
||||
let connection_reference = connection_slab
|
||||
.get_mut(self.connection_id.0)
|
||||
.ok_or_else(|| {
|
||||
anyhow::anyhow!(
|
||||
"connection reference {} not found in slab",
|
||||
self.connection_id.0
|
||||
)
|
||||
})?;
|
||||
|
||||
// Store peer id / check if stored peer id matches
|
||||
match connection_reference
|
||||
.announced_info_hashes
|
||||
.entry(announce_request.info_hash)
|
||||
{
|
||||
Entry::Occupied(entry) => {
|
||||
if *entry.get() != announce_request.peer_id {
|
||||
// Drop Rc borrow before awaiting
|
||||
drop(connection_slab);
|
||||
|
||||
self.send_error_response(
|
||||
"Only one peer id can be used per torrent".into(),
|
||||
Some(ErrorResponseAction::Announce),
|
||||
Some(info_hash),
|
||||
)
|
||||
.await?;
|
||||
|
||||
return Err(anyhow::anyhow!(
|
||||
"Peer used more than one PeerId for a single torrent"
|
||||
));
|
||||
}
|
||||
}
|
||||
Entry::Vacant(entry) => {
|
||||
#[cfg(feature = "prometheus")]
|
||||
if self.config.metrics.run_prometheus_endpoint
|
||||
&& self.config.metrics.peer_clients
|
||||
&& connection_reference.opt_peer_client.is_none()
|
||||
{
|
||||
let peer_id =
|
||||
aquatic_peer_id::PeerId(announce_request.peer_id.0);
|
||||
let client = peer_id.client();
|
||||
let prefix = peer_id.first_8_bytes_hex().to_string();
|
||||
|
||||
::metrics::increment_gauge!(
|
||||
"aquatic_peer_clients",
|
||||
1.0,
|
||||
"client" => client.to_string(),
|
||||
);
|
||||
|
||||
if self.config.metrics.peer_id_prefixes {
|
||||
::metrics::increment_gauge!(
|
||||
"aquatic_peer_id_prefixes",
|
||||
1.0,
|
||||
"prefix_hex" => prefix.to_string(),
|
||||
);
|
||||
}
|
||||
|
||||
connection_reference.opt_peer_client = Some((client, prefix));
|
||||
};
|
||||
|
||||
entry.insert(announce_request.peer_id);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let in_message = InMessage::AnnounceRequest(announce_request);
|
||||
|
||||
let consumer_index =
|
||||
calculate_in_message_consumer_index(&self.config, info_hash);
|
||||
|
||||
// Only fails when receiver is closed
|
||||
self.in_message_senders
|
||||
.send_to(
|
||||
consumer_index,
|
||||
(self.make_connection_meta(None), in_message),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
} else {
|
||||
self.send_error_response(
|
||||
"Info hash not allowed".into(),
|
||||
Some(ErrorResponseAction::Announce),
|
||||
Some(info_hash),
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
}
|
||||
InMessage::ScrapeRequest(ScrapeRequest { info_hashes, .. }) => {
|
||||
#[cfg(feature = "metrics")]
|
||||
::metrics::increment_counter!(
|
||||
"aquatic_requests_total",
|
||||
"type" => "scrape",
|
||||
"ip_version" => ip_version_to_metrics_str(self.ip_version),
|
||||
"worker_index" => WORKER_INDEX.with(|index| index.get()).to_string(),
|
||||
);
|
||||
|
||||
let info_hashes = if let Some(info_hashes) = info_hashes {
|
||||
info_hashes
|
||||
} else {
|
||||
// If request.info_hashes is empty, don't return scrape for all
|
||||
// torrents, even though reference server does it. It is too expensive.
|
||||
self.send_error_response(
|
||||
"Full scrapes are not allowed".into(),
|
||||
Some(ErrorResponseAction::Scrape),
|
||||
None,
|
||||
)
|
||||
.await?;
|
||||
|
||||
return Ok(());
|
||||
};
|
||||
|
||||
let mut info_hashes_by_worker: BTreeMap<usize, Vec<InfoHash>> = BTreeMap::new();
|
||||
|
||||
for info_hash in info_hashes.as_vec() {
|
||||
let info_hashes = info_hashes_by_worker
|
||||
.entry(calculate_in_message_consumer_index(&self.config, info_hash))
|
||||
.or_default();
|
||||
|
||||
info_hashes.push(info_hash);
|
||||
}
|
||||
|
||||
let pending_worker_out_messages = info_hashes_by_worker.len();
|
||||
|
||||
let pending_scrape_response = PendingScrapeResponse {
|
||||
pending_worker_out_messages,
|
||||
stats: Default::default(),
|
||||
};
|
||||
|
||||
let pending_scrape_id: u8 = self
|
||||
.pending_scrape_slab
|
||||
.borrow_mut()
|
||||
.insert(pending_scrape_response)
|
||||
.try_into()
|
||||
.with_context(|| "Reached 256 pending scrape responses")?;
|
||||
|
||||
let meta = self.make_connection_meta(Some(PendingScrapeId(pending_scrape_id)));
|
||||
|
||||
for (consumer_index, info_hashes) in info_hashes_by_worker {
|
||||
let in_message = InMessage::ScrapeRequest(ScrapeRequest {
|
||||
action: ScrapeAction,
|
||||
info_hashes: Some(ScrapeRequestInfoHashes::Multiple(info_hashes)),
|
||||
});
|
||||
|
||||
// Only fails when receiver is closed
|
||||
self.in_message_senders
|
||||
.send_to(consumer_index, (meta, in_message))
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn send_error_response(
|
||||
&self,
|
||||
failure_reason: Cow<'static, str>,
|
||||
action: Option<ErrorResponseAction>,
|
||||
info_hash: Option<InfoHash>,
|
||||
) -> anyhow::Result<()> {
|
||||
let out_message = OutMessage::ErrorResponse(ErrorResponse {
|
||||
action,
|
||||
failure_reason,
|
||||
info_hash,
|
||||
});
|
||||
|
||||
let result = self
|
||||
.out_message_sender
|
||||
.send((self.make_connection_meta(None).into(), out_message))
|
||||
.await
|
||||
.map_err(|err| {
|
||||
anyhow::anyhow!("ConnectionReader::send_error_response failed: {}", err)
|
||||
});
|
||||
|
||||
#[cfg(feature = "metrics")]
|
||||
::metrics::increment_counter!(
|
||||
"aquatic_responses_total",
|
||||
"type" => "error",
|
||||
"ip_version" => ip_version_to_metrics_str(self.ip_version),
|
||||
"worker_index" => WORKER_INDEX.with(|index| index.get()).to_string(),
|
||||
);
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
fn make_connection_meta(&self, pending_scrape_id: Option<PendingScrapeId>) -> InMessageMeta {
|
||||
InMessageMeta {
|
||||
connection_id: self.connection_id,
|
||||
out_message_consumer_id: self.out_message_consumer_id,
|
||||
ip_version: self.ip_version,
|
||||
pending_scrape_id,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct ConnectionWriter<S> {
|
||||
config: Rc<Config>,
|
||||
out_message_receiver: LocalReceiver<(OutMessageMeta, OutMessage)>,
|
||||
connection_slab: Rc<RefCell<Slab<ConnectionReference>>>,
|
||||
ws_out: SplitSink<WebSocketStream<S>, tungstenite::Message>,
|
||||
pending_scrape_slab: Rc<RefCell<Slab<PendingScrapeResponse>>>,
|
||||
server_start_instant: ServerStartInstant,
|
||||
connection_id: ConnectionId,
|
||||
ip_version: IpVersion,
|
||||
}
|
||||
|
||||
impl<S: futures::AsyncRead + futures::AsyncWrite + Unpin> ConnectionWriter<S> {
|
||||
async fn run_out_message_loop(&mut self) -> anyhow::Result<()> {
|
||||
loop {
|
||||
let (meta, out_message) = self.out_message_receiver.recv().await.ok_or_else(|| {
|
||||
anyhow::anyhow!("ConnectionWriter couldn't receive message, sender is closed")
|
||||
})?;
|
||||
|
||||
match out_message {
|
||||
OutMessage::ScrapeResponse(out_message) => {
|
||||
let pending_scrape_id = meta
|
||||
.pending_scrape_id
|
||||
.expect("meta.pending_scrape_id not set");
|
||||
|
||||
let finished = if let Some(pending) = Slab::get_mut(
|
||||
&mut RefCell::borrow_mut(&self.pending_scrape_slab),
|
||||
pending_scrape_id.0 as usize,
|
||||
) {
|
||||
pending.stats.extend(out_message.files);
|
||||
pending.pending_worker_out_messages -= 1;
|
||||
|
||||
pending.pending_worker_out_messages == 0
|
||||
} else {
|
||||
return Err(anyhow::anyhow!("pending scrape not found in slab"));
|
||||
};
|
||||
|
||||
if finished {
|
||||
let out_message = {
|
||||
let mut slab = RefCell::borrow_mut(&self.pending_scrape_slab);
|
||||
|
||||
let pending = slab.remove(pending_scrape_id.0 as usize);
|
||||
|
||||
slab.shrink_to_fit();
|
||||
|
||||
OutMessage::ScrapeResponse(ScrapeResponse {
|
||||
action: ScrapeAction,
|
||||
files: pending.stats,
|
||||
})
|
||||
};
|
||||
|
||||
self.send_out_message(&out_message).await?;
|
||||
}
|
||||
}
|
||||
out_message => {
|
||||
self.send_out_message(&out_message).await?;
|
||||
}
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
async fn send_out_message(&mut self, out_message: &OutMessage) -> anyhow::Result<()> {
|
||||
let result = timeout(Duration::from_secs(10), async {
|
||||
let result =
|
||||
futures::SinkExt::send(&mut self.ws_out, out_message.to_ws_message()).await;
|
||||
|
||||
Ok(result)
|
||||
})
|
||||
.await;
|
||||
|
||||
match result {
|
||||
Ok(Ok(())) => {
|
||||
#[cfg(feature = "metrics")]
|
||||
{
|
||||
let out_message_type = match &out_message {
|
||||
OutMessage::Offer(_) => "offer",
|
||||
OutMessage::Answer(_) => "offer_answer",
|
||||
OutMessage::AnnounceResponse(_) => "announce",
|
||||
OutMessage::ScrapeResponse(_) => "scrape",
|
||||
OutMessage::ErrorResponse(_) => "error",
|
||||
};
|
||||
|
||||
::metrics::increment_counter!(
|
||||
"aquatic_responses_total",
|
||||
"type" => out_message_type,
|
||||
"ip_version" => ip_version_to_metrics_str(self.ip_version),
|
||||
"worker_index" => WORKER_INDEX.with(|index| index.get()).to_string(),
|
||||
);
|
||||
}
|
||||
|
||||
self.connection_slab
|
||||
.borrow_mut()
|
||||
.get_mut(self.connection_id.0)
|
||||
.ok_or_else(|| {
|
||||
anyhow::anyhow!(
|
||||
"connection reference {} not found in slab",
|
||||
self.connection_id.0
|
||||
)
|
||||
})?
|
||||
.valid_until = ValidUntil::new(
|
||||
self.server_start_instant,
|
||||
self.config.cleaning.max_connection_idle,
|
||||
);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
Ok(Err(err)) => Err(err.into()),
|
||||
Err(err) => {
|
||||
::log::debug!("send_out_message: sending to peer took too long: {}", err);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn calculate_in_message_consumer_index(config: &Config, info_hash: InfoHash) -> usize {
|
||||
(info_hash.0[0] as usize) % config.swarm_workers
|
||||
}
|
||||
|
||||
fn create_tcp_listener(
|
||||
config: &Config,
|
||||
priv_dropper: PrivilegeDropper,
|
||||
) -> anyhow::Result<TcpListener> {
|
||||
let domain = if config.network.address.is_ipv4() {
|
||||
socket2::Domain::IPV4
|
||||
} else {
|
||||
socket2::Domain::IPV6
|
||||
};
|
||||
|
||||
::log::info!("creating socket..");
|
||||
|
||||
let socket = socket2::Socket::new(domain, socket2::Type::STREAM, Some(socket2::Protocol::TCP))
|
||||
.with_context(|| "create socket")?;
|
||||
|
||||
if config.network.only_ipv6 {
|
||||
::log::info!("setting socket to ipv6 only..");
|
||||
|
||||
socket
|
||||
.set_only_v6(true)
|
||||
.with_context(|| "socket: set only ipv6")?;
|
||||
}
|
||||
|
||||
::log::info!("setting SO_REUSEPORT on socket..");
|
||||
|
||||
socket
|
||||
.set_reuse_port(true)
|
||||
.with_context(|| "socket: set reuse port")?;
|
||||
|
||||
::log::info!("binding socket..");
|
||||
|
||||
socket
|
||||
.bind(&config.network.address.into())
|
||||
.with_context(|| format!("socket: bind to {}", config.network.address))?;
|
||||
|
||||
::log::info!("listening on socket..");
|
||||
|
||||
socket
|
||||
.listen(config.network.tcp_backlog)
|
||||
.with_context(|| format!("socket: listen {}", config.network.address))?;
|
||||
|
||||
::log::info!("running PrivilegeDropper::after_socket_creation..");
|
||||
|
||||
priv_dropper.after_socket_creation()?;
|
||||
|
||||
::log::info!("casting socket to glommio TcpListener..");
|
||||
|
||||
Ok(unsafe { TcpListener::from_raw_fd(socket.into_raw_fd()) })
|
||||
}
|
||||
|
||||
#[cfg(feature = "metrics")]
|
||||
fn ip_version_to_metrics_str(ip_version: IpVersion) -> &'static str {
|
||||
match ip_version {
|
||||
IpVersion::V4 => "4",
|
||||
IpVersion::V6 => "6",
|
||||
}
|
||||
}
|
||||
539
crates/ws/src/workers/swarm.rs
Normal file
539
crates/ws/src/workers/swarm.rs
Normal file
|
|
@ -0,0 +1,539 @@
|
|||
use std::cell::RefCell;
|
||||
use std::rc::Rc;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use aquatic_common::access_list::{create_access_list_cache, AccessListArcSwap, AccessListCache};
|
||||
use futures::StreamExt;
|
||||
use glommio::channels::channel_mesh::{MeshBuilder, Partial, Role, Senders};
|
||||
use glommio::enclose;
|
||||
use glommio::prelude::*;
|
||||
use glommio::timer::TimerActionRepeat;
|
||||
use hashbrown::HashMap;
|
||||
use rand::{rngs::SmallRng, SeedableRng};
|
||||
|
||||
use aquatic_common::{
|
||||
extract_response_peers, IndexMap, PanicSentinel, SecondsSinceServerStart, ServerStartInstant,
|
||||
};
|
||||
use aquatic_ws_protocol::*;
|
||||
|
||||
use crate::common::*;
|
||||
use crate::config::Config;
|
||||
use crate::SHARED_IN_CHANNEL_SIZE;
|
||||
|
||||
#[cfg(feature = "metrics")]
|
||||
thread_local! { static WORKER_INDEX: ::std::cell::Cell<usize> = Default::default() }
|
||||
|
||||
#[derive(PartialEq, Eq, Clone, Copy, Debug)]
|
||||
enum PeerStatus {
|
||||
Seeding,
|
||||
Leeching,
|
||||
Stopped,
|
||||
}
|
||||
|
||||
impl PeerStatus {
|
||||
/// Determine peer status from announce event and number of bytes left.
|
||||
///
|
||||
/// Likely, the last branch will be taken most of the time.
|
||||
#[inline]
|
||||
fn from_event_and_bytes_left(event: AnnounceEvent, opt_bytes_left: Option<usize>) -> Self {
|
||||
if let AnnounceEvent::Stopped = event {
|
||||
Self::Stopped
|
||||
} else if let Some(0) = opt_bytes_left {
|
||||
Self::Seeding
|
||||
} else {
|
||||
Self::Leeching
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
struct Peer {
|
||||
pub consumer_id: ConsumerId,
|
||||
pub connection_id: ConnectionId,
|
||||
pub seeder: bool,
|
||||
pub valid_until: ValidUntil,
|
||||
}
|
||||
|
||||
type PeerMap = IndexMap<PeerId, Peer>;
|
||||
|
||||
struct TorrentData {
|
||||
pub peers: PeerMap,
|
||||
pub num_seeders: usize,
|
||||
}
|
||||
|
||||
impl Default for TorrentData {
|
||||
#[inline]
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
peers: Default::default(),
|
||||
num_seeders: 0,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl TorrentData {
|
||||
pub fn remove_peer(&mut self, peer_id: PeerId) {
|
||||
if let Some(peer) = self.peers.remove(&peer_id) {
|
||||
if peer.seeder {
|
||||
self.num_seeders -= 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn num_leechers(&self) -> usize {
|
||||
self.peers.len() - self.num_seeders
|
||||
}
|
||||
}
|
||||
|
||||
type TorrentMap = IndexMap<InfoHash, TorrentData>;
|
||||
|
||||
#[derive(Default)]
|
||||
struct TorrentMaps {
|
||||
pub ipv4: TorrentMap,
|
||||
pub ipv6: TorrentMap,
|
||||
}
|
||||
|
||||
impl TorrentMaps {
|
||||
fn clean(
|
||||
&mut self,
|
||||
config: &Config,
|
||||
access_list: &Arc<AccessListArcSwap>,
|
||||
server_start_instant: ServerStartInstant,
|
||||
) {
|
||||
let mut access_list_cache = create_access_list_cache(access_list);
|
||||
let now = server_start_instant.seconds_elapsed();
|
||||
|
||||
Self::clean_torrent_map(config, &mut access_list_cache, &mut self.ipv4, now, "4");
|
||||
Self::clean_torrent_map(config, &mut access_list_cache, &mut self.ipv6, now, "6");
|
||||
}
|
||||
|
||||
fn clean_torrent_map(
|
||||
config: &Config,
|
||||
access_list_cache: &mut AccessListCache,
|
||||
torrent_map: &mut TorrentMap,
|
||||
now: SecondsSinceServerStart,
|
||||
ip_version: &'static str,
|
||||
) {
|
||||
let mut total_num_peers = 0u64;
|
||||
|
||||
torrent_map.retain(|info_hash, torrent_data| {
|
||||
if !access_list_cache
|
||||
.load()
|
||||
.allows(config.access_list.mode, &info_hash.0)
|
||||
{
|
||||
return false;
|
||||
}
|
||||
|
||||
let num_seeders = &mut torrent_data.num_seeders;
|
||||
|
||||
torrent_data.peers.retain(|_, peer| {
|
||||
let keep = peer.valid_until.valid(now);
|
||||
|
||||
if (!keep) & peer.seeder {
|
||||
*num_seeders -= 1;
|
||||
}
|
||||
|
||||
keep
|
||||
});
|
||||
|
||||
total_num_peers += torrent_data.peers.len() as u64;
|
||||
|
||||
!torrent_data.peers.is_empty()
|
||||
});
|
||||
|
||||
torrent_map.shrink_to_fit();
|
||||
|
||||
let total_num_peers = total_num_peers as f64;
|
||||
|
||||
#[cfg(feature = "metrics")]
|
||||
::metrics::gauge!(
|
||||
"aquatic_peers",
|
||||
total_num_peers,
|
||||
"ip_version" => ip_version,
|
||||
"worker_index" => WORKER_INDEX.with(|index| index.get()).to_string(),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn run_swarm_worker(
|
||||
_sentinel: PanicSentinel,
|
||||
config: Config,
|
||||
state: State,
|
||||
control_message_mesh_builder: MeshBuilder<SwarmControlMessage, Partial>,
|
||||
in_message_mesh_builder: MeshBuilder<(InMessageMeta, InMessage), Partial>,
|
||||
out_message_mesh_builder: MeshBuilder<(OutMessageMeta, OutMessage), Partial>,
|
||||
server_start_instant: ServerStartInstant,
|
||||
worker_index: usize,
|
||||
) {
|
||||
#[cfg(feature = "metrics")]
|
||||
WORKER_INDEX.with(|index| index.set(worker_index));
|
||||
|
||||
let (_, mut control_message_receivers) = control_message_mesh_builder
|
||||
.join(Role::Consumer)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let (_, mut in_message_receivers) = in_message_mesh_builder.join(Role::Consumer).await.unwrap();
|
||||
let (out_message_senders, _) = out_message_mesh_builder.join(Role::Producer).await.unwrap();
|
||||
|
||||
let out_message_senders = Rc::new(out_message_senders);
|
||||
|
||||
let torrents = Rc::new(RefCell::new(TorrentMaps::default()));
|
||||
let access_list = state.access_list;
|
||||
|
||||
// Periodically clean torrents
|
||||
TimerActionRepeat::repeat(enclose!((config, torrents, access_list) move || {
|
||||
enclose!((config, torrents, access_list) move || async move {
|
||||
torrents.borrow_mut().clean(&config, &access_list, server_start_instant);
|
||||
|
||||
Some(Duration::from_secs(config.cleaning.torrent_cleaning_interval))
|
||||
})()
|
||||
}));
|
||||
|
||||
// Periodically update torrent count metrics
|
||||
#[cfg(feature = "metrics")]
|
||||
TimerActionRepeat::repeat(enclose!((config, torrents) move || {
|
||||
enclose!((config, torrents) move || async move {
|
||||
let torrents = torrents.borrow_mut();
|
||||
|
||||
::metrics::gauge!(
|
||||
"aquatic_torrents",
|
||||
torrents.ipv4.len() as f64,
|
||||
"ip_version" => "4",
|
||||
"worker_index" => worker_index.to_string(),
|
||||
);
|
||||
::metrics::gauge!(
|
||||
"aquatic_torrents",
|
||||
torrents.ipv6.len() as f64,
|
||||
"ip_version" => "6",
|
||||
"worker_index" => worker_index.to_string(),
|
||||
);
|
||||
|
||||
Some(Duration::from_secs(config.metrics.torrent_count_update_interval))
|
||||
})()
|
||||
}));
|
||||
|
||||
let mut handles = Vec::new();
|
||||
|
||||
for (_, receiver) in control_message_receivers.streams() {
|
||||
let handle =
|
||||
spawn_local(handle_control_message_stream(torrents.clone(), receiver)).detach();
|
||||
|
||||
handles.push(handle);
|
||||
}
|
||||
|
||||
for (_, receiver) in in_message_receivers.streams() {
|
||||
let handle = spawn_local(handle_request_stream(
|
||||
config.clone(),
|
||||
torrents.clone(),
|
||||
server_start_instant,
|
||||
out_message_senders.clone(),
|
||||
receiver,
|
||||
))
|
||||
.detach();
|
||||
|
||||
handles.push(handle);
|
||||
}
|
||||
|
||||
for handle in handles {
|
||||
handle.await;
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_control_message_stream<S>(torrents: Rc<RefCell<TorrentMaps>>, mut stream: S)
|
||||
where
|
||||
S: futures_lite::Stream<Item = SwarmControlMessage> + ::std::marker::Unpin,
|
||||
{
|
||||
while let Some(message) = stream.next().await {
|
||||
match message {
|
||||
SwarmControlMessage::ConnectionClosed {
|
||||
info_hash,
|
||||
peer_id,
|
||||
ip_version,
|
||||
} => {
|
||||
::log::debug!("Removing peer from torrents because connection was closed");
|
||||
|
||||
if let IpVersion::V4 = ip_version {
|
||||
if let Some(torrent_data) = torrents.borrow_mut().ipv4.get_mut(&info_hash) {
|
||||
torrent_data.remove_peer(peer_id);
|
||||
}
|
||||
} else {
|
||||
if let Some(torrent_data) = torrents.borrow_mut().ipv6.get_mut(&info_hash) {
|
||||
torrent_data.remove_peer(peer_id);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_request_stream<S>(
|
||||
config: Config,
|
||||
torrents: Rc<RefCell<TorrentMaps>>,
|
||||
server_start_instant: ServerStartInstant,
|
||||
out_message_senders: Rc<Senders<(OutMessageMeta, OutMessage)>>,
|
||||
stream: S,
|
||||
) where
|
||||
S: futures_lite::Stream<Item = (InMessageMeta, InMessage)> + ::std::marker::Unpin,
|
||||
{
|
||||
let rng = Rc::new(RefCell::new(SmallRng::from_entropy()));
|
||||
|
||||
let max_peer_age = config.cleaning.max_peer_age;
|
||||
let peer_valid_until = Rc::new(RefCell::new(ValidUntil::new(
|
||||
server_start_instant,
|
||||
max_peer_age,
|
||||
)));
|
||||
|
||||
TimerActionRepeat::repeat(enclose!((peer_valid_until) move || {
|
||||
enclose!((peer_valid_until) move || async move {
|
||||
*peer_valid_until.borrow_mut() = ValidUntil::new(server_start_instant, max_peer_age);
|
||||
|
||||
Some(Duration::from_secs(1))
|
||||
})()
|
||||
}));
|
||||
|
||||
let config = &config;
|
||||
let torrents = &torrents;
|
||||
let peer_valid_until = &peer_valid_until;
|
||||
let rng = &rng;
|
||||
let out_message_senders = &out_message_senders;
|
||||
|
||||
stream
|
||||
.for_each_concurrent(
|
||||
SHARED_IN_CHANNEL_SIZE,
|
||||
move |(meta, in_message)| async move {
|
||||
let mut out_messages = Vec::new();
|
||||
|
||||
match in_message {
|
||||
InMessage::AnnounceRequest(request) => handle_announce_request(
|
||||
&config,
|
||||
&mut rng.borrow_mut(),
|
||||
&mut torrents.borrow_mut(),
|
||||
&mut out_messages,
|
||||
peer_valid_until.borrow().to_owned(),
|
||||
meta,
|
||||
request,
|
||||
),
|
||||
InMessage::ScrapeRequest(request) => handle_scrape_request(
|
||||
&config,
|
||||
&mut torrents.borrow_mut(),
|
||||
&mut out_messages,
|
||||
meta,
|
||||
request,
|
||||
),
|
||||
};
|
||||
|
||||
for (meta, out_message) in out_messages.drain(..) {
|
||||
out_message_senders
|
||||
.send_to(meta.out_message_consumer_id.0 as usize, (meta, out_message))
|
||||
.await
|
||||
.expect("failed sending out_message to socket worker");
|
||||
|
||||
::log::debug!("swarm worker sent OutMessage to socket worker");
|
||||
}
|
||||
},
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
fn handle_announce_request(
|
||||
config: &Config,
|
||||
rng: &mut SmallRng,
|
||||
torrent_maps: &mut TorrentMaps,
|
||||
out_messages: &mut Vec<(OutMessageMeta, OutMessage)>,
|
||||
valid_until: ValidUntil,
|
||||
request_sender_meta: InMessageMeta,
|
||||
request: AnnounceRequest,
|
||||
) {
|
||||
let (torrent_data, ip_version): (&mut TorrentData, &'static str) =
|
||||
if let IpVersion::V4 = request_sender_meta.ip_version {
|
||||
(torrent_maps.ipv4.entry(request.info_hash).or_default(), "4")
|
||||
} else {
|
||||
(torrent_maps.ipv6.entry(request.info_hash).or_default(), "6")
|
||||
};
|
||||
|
||||
// If there is already a peer with this peer_id, check that connection id
|
||||
// is same as that of request sender. Otherwise, ignore request. Since
|
||||
// peers have access to each others peer_id's, they could send requests
|
||||
// using them, causing all sorts of issues.
|
||||
if let Some(previous_peer) = torrent_data.peers.get(&request.peer_id) {
|
||||
if request_sender_meta.connection_id != previous_peer.connection_id {
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
::log::trace!("received request from {:?}", request_sender_meta);
|
||||
|
||||
// Insert/update/remove peer who sent this request
|
||||
{
|
||||
let peer_status = PeerStatus::from_event_and_bytes_left(
|
||||
request.event.unwrap_or_default(),
|
||||
request.bytes_left,
|
||||
);
|
||||
|
||||
let opt_removed_peer = match peer_status {
|
||||
PeerStatus::Leeching => {
|
||||
let peer = Peer {
|
||||
connection_id: request_sender_meta.connection_id,
|
||||
consumer_id: request_sender_meta.out_message_consumer_id,
|
||||
seeder: false,
|
||||
valid_until,
|
||||
};
|
||||
|
||||
torrent_data.peers.insert(request.peer_id, peer)
|
||||
}
|
||||
PeerStatus::Seeding => {
|
||||
torrent_data.num_seeders += 1;
|
||||
|
||||
let peer = Peer {
|
||||
connection_id: request_sender_meta.connection_id,
|
||||
consumer_id: request_sender_meta.out_message_consumer_id,
|
||||
seeder: true,
|
||||
valid_until,
|
||||
};
|
||||
|
||||
torrent_data.peers.insert(request.peer_id, peer)
|
||||
}
|
||||
PeerStatus::Stopped => torrent_data.peers.remove(&request.peer_id),
|
||||
};
|
||||
|
||||
if let Some(&Peer { seeder: true, .. }) = opt_removed_peer.as_ref() {
|
||||
torrent_data.num_seeders -= 1;
|
||||
}
|
||||
|
||||
#[cfg(feature = "metrics")]
|
||||
match peer_status {
|
||||
PeerStatus::Stopped if opt_removed_peer.is_some() => {
|
||||
::metrics::decrement_gauge!(
|
||||
"aquatic_peers",
|
||||
1.0,
|
||||
"ip_version" => ip_version,
|
||||
"worker_index" => WORKER_INDEX.with(|index| index.get()).to_string(),
|
||||
);
|
||||
}
|
||||
PeerStatus::Leeching | PeerStatus::Seeding if opt_removed_peer.is_none() => {
|
||||
::metrics::increment_gauge!(
|
||||
"aquatic_peers",
|
||||
1.0,
|
||||
"ip_version" => ip_version,
|
||||
"worker_index" => WORKER_INDEX.with(|index| index.get()).to_string(),
|
||||
);
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
// If peer sent offers, send them on to random peers
|
||||
if let Some(offers) = request.offers {
|
||||
// FIXME: config: also maybe check this when parsing request
|
||||
let max_num_peers_to_take = offers.len().min(config.protocol.max_offers);
|
||||
|
||||
#[inline]
|
||||
fn f(peer: &Peer) -> Peer {
|
||||
*peer
|
||||
}
|
||||
|
||||
let offer_receivers: Vec<Peer> = extract_response_peers(
|
||||
rng,
|
||||
&torrent_data.peers,
|
||||
max_num_peers_to_take,
|
||||
request.peer_id,
|
||||
f,
|
||||
);
|
||||
|
||||
for (offer, offer_receiver) in offers.into_iter().zip(offer_receivers) {
|
||||
let middleman_offer = MiddlemanOfferToPeer {
|
||||
action: AnnounceAction,
|
||||
info_hash: request.info_hash,
|
||||
peer_id: request.peer_id,
|
||||
offer: offer.offer,
|
||||
offer_id: offer.offer_id,
|
||||
};
|
||||
|
||||
let meta = OutMessageMeta {
|
||||
out_message_consumer_id: offer_receiver.consumer_id,
|
||||
connection_id: offer_receiver.connection_id,
|
||||
pending_scrape_id: None,
|
||||
};
|
||||
|
||||
out_messages.push((meta, OutMessage::Offer(middleman_offer)));
|
||||
::log::trace!("sending middleman offer to {:?}", meta);
|
||||
}
|
||||
}
|
||||
|
||||
// If peer sent answer, send it on to relevant peer
|
||||
if let (Some(answer), Some(answer_receiver_id), Some(offer_id)) =
|
||||
(request.answer, request.to_peer_id, request.offer_id)
|
||||
{
|
||||
if let Some(answer_receiver) = torrent_data.peers.get(&answer_receiver_id) {
|
||||
let middleman_answer = MiddlemanAnswerToPeer {
|
||||
action: AnnounceAction,
|
||||
peer_id: request.peer_id,
|
||||
info_hash: request.info_hash,
|
||||
answer,
|
||||
offer_id,
|
||||
};
|
||||
|
||||
let meta = OutMessageMeta {
|
||||
out_message_consumer_id: answer_receiver.consumer_id,
|
||||
connection_id: answer_receiver.connection_id,
|
||||
pending_scrape_id: None,
|
||||
};
|
||||
|
||||
out_messages.push((meta, OutMessage::Answer(middleman_answer)));
|
||||
::log::trace!("sending middleman answer to {:?}", meta);
|
||||
}
|
||||
}
|
||||
|
||||
let out_message = OutMessage::AnnounceResponse(AnnounceResponse {
|
||||
action: AnnounceAction,
|
||||
info_hash: request.info_hash,
|
||||
complete: torrent_data.num_seeders,
|
||||
incomplete: torrent_data.num_leechers(),
|
||||
announce_interval: config.protocol.peer_announce_interval,
|
||||
});
|
||||
|
||||
out_messages.push((request_sender_meta.into(), out_message));
|
||||
}
|
||||
|
||||
fn handle_scrape_request(
|
||||
config: &Config,
|
||||
torrent_maps: &mut TorrentMaps,
|
||||
out_messages: &mut Vec<(OutMessageMeta, OutMessage)>,
|
||||
meta: InMessageMeta,
|
||||
request: ScrapeRequest,
|
||||
) {
|
||||
let info_hashes = if let Some(info_hashes) = request.info_hashes {
|
||||
info_hashes.as_vec()
|
||||
} else {
|
||||
return;
|
||||
};
|
||||
|
||||
let num_to_take = info_hashes.len().min(config.protocol.max_scrape_torrents);
|
||||
|
||||
let mut out_message = ScrapeResponse {
|
||||
action: ScrapeAction,
|
||||
files: HashMap::with_capacity(num_to_take),
|
||||
};
|
||||
|
||||
let torrent_map: &mut TorrentMap = if let IpVersion::V4 = meta.ip_version {
|
||||
&mut torrent_maps.ipv4
|
||||
} else {
|
||||
&mut torrent_maps.ipv6
|
||||
};
|
||||
|
||||
for info_hash in info_hashes.into_iter().take(num_to_take) {
|
||||
if let Some(torrent_data) = torrent_map.get(&info_hash) {
|
||||
let stats = ScrapeStatistics {
|
||||
complete: torrent_data.num_seeders,
|
||||
downloaded: 0, // No implementation planned
|
||||
incomplete: torrent_data.num_leechers(),
|
||||
};
|
||||
|
||||
out_message.files.insert(info_hash, stats);
|
||||
}
|
||||
}
|
||||
|
||||
out_messages.push((meta.into(), OutMessage::ScrapeResponse(out_message)));
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue