use std::cell::RefCell; 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::privileges::PrivilegeDropper; use aquatic_common::rustls_config::RustlsConfig; use aquatic_common::{PanicSentinel, ServerStartInstant}; use aquatic_ws_protocol::common::InfoHash; use aquatic_ws_protocol::incoming::InMessage; use aquatic_ws_protocol::outgoing::OutMessage; use arc_swap::ArcSwap; use futures::StreamExt; use glommio::channels::channel_mesh::{MeshBuilder, Partial, Role}; use glommio::channels::local_channel::{new_bounded, LocalSender}; use glommio::channels::shared_channel::ConnectedReceiver; use glommio::net::TcpListener; use glommio::timer::TimerActionRepeat; use glommio::{enclose, prelude::*}; use slotmap::HopSlotMap; use crate::config::Config; use crate::common::*; use crate::workers::socket::connection::ConnectionRunner; mod connection; type ConnectionHandles = HopSlotMap; const LOCAL_CHANNEL_SIZE: usize = 16; #[cfg(feature = "metrics")] thread_local! { static WORKER_INDEX: ::std::cell::Cell = Default::default() } /// Used to interact with the connection tasks struct ConnectionHandle { close_conn_sender: LocalSender<()>, /// Sender part of channel used to pass on outgoing messages from request /// worker out_message_sender: Rc>, /// Updated after sending message to peer valid_until: Rc>, /// The TLS config used for this connection opt_tls_config: Option>, valid_until_after_tls_update: Option, } #[allow(clippy::too_many_arguments)] pub async fn run_socket_worker( _sentinel: PanicSentinel, config: Config, state: State, opt_tls_config: Option>>, control_message_mesh_builder: MeshBuilder, 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_handles = Rc::new(RefCell::new(ConnectionHandles::default())); // Periodically clean connections TimerActionRepeat::repeat_into( enclose!((config, connection_handles, opt_tls_config) move || { clean_connections( config.clone(), connection_handles.clone(), server_start_instant, opt_tls_config.clone(), ) }), tq_prioritized, ) .unwrap(); for (_, out_message_receiver) in out_message_receivers.streams() { spawn_local_into( receive_out_messages(out_message_receiver, connection_handles.clone()), tq_regular, ) .unwrap() .detach(); } let mut incoming = listener.incoming(); while let Some(stream) = incoming.next().await { match stream { Err(err) => { ::log::error!("accept connection: {:#}", err); } 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 (close_conn_sender, close_conn_receiver) = new_bounded(1); let connection_valid_until = Rc::new(RefCell::new(ValidUntil::new( server_start_instant, config.cleaning.max_connection_idle, ))); let connection_handle = ConnectionHandle { close_conn_sender, out_message_sender: out_message_sender.clone(), valid_until: connection_valid_until.clone(), opt_tls_config: opt_tls_config.as_ref().map(|c| c.load_full()), valid_until_after_tls_update: None, }; let connection_id = connection_handles.borrow_mut().insert(connection_handle); spawn_local_into( enclose!(( config, access_list, in_message_senders, connection_valid_until, opt_tls_config, control_message_senders, connection_handles ) async move { let runner = ConnectionRunner { config, access_list, in_message_senders, connection_valid_until, out_message_sender, out_message_receiver, server_start_instant, out_message_consumer_id, connection_id, opt_tls_config, ip_version }; runner.run(control_message_senders, close_conn_receiver, stream).await; connection_handles.borrow_mut().remove(connection_id); }), tq_regular, ) .unwrap() .detach(); } } } } async fn clean_connections( config: Rc, connection_slab: Rc>, server_start_instant: ServerStartInstant, opt_tls_config: Option>>, ) -> Option { let now = server_start_instant.seconds_elapsed(); let opt_current_tls_config = opt_tls_config.map(|c| c.load_full()); connection_slab.borrow_mut().retain(|_, reference| { let mut keep = true; // Handle case when connection runs on an old TLS certificate if let Some(valid_until) = reference.valid_until_after_tls_update { if !valid_until.valid(now) { keep = false; } } else if let Some(false) = opt_current_tls_config .as_ref() .zip(reference.opt_tls_config.as_ref()) .map(|(a, b)| Arc::ptr_eq(a, b)) { reference.valid_until_after_tls_update = Some(ValidUntil::new( server_start_instant, config.cleaning.close_after_tls_update_grace_period, )); } keep &= reference.valid_until.borrow().valid(now); if keep { true } else { if let Err(err) = reference.close_conn_sender.try_send(()) { ::log::info!("couldn't tell connection to close: {:#}", err); } false } }); #[cfg(feature = "metrics")] { ::log::info!( "cleaned connections in worker {}, {} references remaining", WORKER_INDEX.get(), connection_slab.borrow_mut().len() ); // 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::gauge!( "aquatic_active_connections", "ip_version" => "4", "worker_index" => worker_index.clone(), ) .increment(0.0); } if config.network.address.is_ipv6() { ::metrics::gauge!( "aquatic_active_connections", "ip_version" => "6", "worker_index" => worker_index, ) .increment(0.0); } } Some(Duration::from_secs( config.cleaning.connection_cleaning_interval, )) } async fn receive_out_messages( mut out_message_receiver: ConnectedReceiver<(OutMessageMeta, OutMessage)>, connection_references: Rc>, ) { 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) { match reference.out_message_sender.try_send((meta, out_message)) { Ok(()) => {} Err(GlommioError::Closed(_)) => {} Err(GlommioError::WouldBlock(_)) => { ::log::debug!( "couldn't send OutMessage over local channel to Connection, channel full" ); } Err(err) => { ::log::debug!( "couldn't send OutMessage over local channel to Connection: {:?}", err ); } } } } } fn create_tcp_listener( config: &Config, priv_dropper: PrivilegeDropper, ) -> anyhow::Result { 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", } } fn calculate_in_message_consumer_index(config: &Config, info_hash: InfoHash) -> usize { (info_hash.0[0] as usize) % config.swarm_workers }