mirror of
https://github.com/YGGverse/aquatic.git
synced 2026-04-02 18:55:32 +00:00
aquatic_ws: use signals for access list updates
This commit is contained in:
parent
9a1993d72e
commit
446fd0b1f4
6 changed files with 93 additions and 94 deletions
1
Cargo.lock
generated
1
Cargo.lock
generated
|
|
@ -259,6 +259,7 @@ dependencies = [
|
||||||
"rand",
|
"rand",
|
||||||
"rustls-pemfile",
|
"rustls-pemfile",
|
||||||
"serde",
|
"serde",
|
||||||
|
"signal-hook",
|
||||||
"slab",
|
"slab",
|
||||||
"tungstenite",
|
"tungstenite",
|
||||||
]
|
]
|
||||||
|
|
|
||||||
|
|
@ -35,6 +35,7 @@ privdrop = "0.5"
|
||||||
rand = { version = "0.8", features = ["small_rng"] }
|
rand = { version = "0.8", features = ["small_rng"] }
|
||||||
rustls-pemfile = "0.2"
|
rustls-pemfile = "0.2"
|
||||||
serde = { version = "1", features = ["derive"] }
|
serde = { version = "1", features = ["derive"] }
|
||||||
|
signal-hook = { version = "0.3" }
|
||||||
slab = "0.4"
|
slab = "0.4"
|
||||||
tungstenite = "0.15"
|
tungstenite = "0.15"
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -1,13 +1,8 @@
|
||||||
use std::borrow::Borrow;
|
|
||||||
use std::cell::RefCell;
|
|
||||||
use std::net::{IpAddr, SocketAddr};
|
use std::net::{IpAddr, SocketAddr};
|
||||||
use std::rc::Rc;
|
use std::sync::Arc;
|
||||||
use std::time::Instant;
|
use std::time::Instant;
|
||||||
|
|
||||||
use aquatic_common::access_list::AccessList;
|
use aquatic_common::access_list::{create_access_list_cache, AccessListArcSwap, AccessListCache};
|
||||||
use futures_lite::AsyncBufReadExt;
|
|
||||||
use glommio::io::{BufferedFile, StreamReaderBuilder};
|
|
||||||
use glommio::yield_if_needed;
|
|
||||||
use hashbrown::HashMap;
|
use hashbrown::HashMap;
|
||||||
use indexmap::IndexMap;
|
use indexmap::IndexMap;
|
||||||
|
|
||||||
|
|
@ -99,16 +94,25 @@ pub struct TorrentMaps {
|
||||||
}
|
}
|
||||||
|
|
||||||
impl TorrentMaps {
|
impl TorrentMaps {
|
||||||
pub fn clean(&mut self, config: &Config, access_list: &AccessList) {
|
pub fn clean(&mut self, config: &Config, access_list: &Arc<AccessListArcSwap>) {
|
||||||
Self::clean_torrent_map(config, access_list, &mut self.ipv4);
|
let mut access_list_cache = create_access_list_cache(access_list);
|
||||||
Self::clean_torrent_map(config, access_list, &mut self.ipv6);
|
|
||||||
|
Self::clean_torrent_map(config, &mut access_list_cache, &mut self.ipv4);
|
||||||
|
Self::clean_torrent_map(config, &mut access_list_cache, &mut self.ipv6);
|
||||||
}
|
}
|
||||||
|
|
||||||
fn clean_torrent_map(config: &Config, access_list: &AccessList, torrent_map: &mut TorrentMap) {
|
fn clean_torrent_map(
|
||||||
|
config: &Config,
|
||||||
|
access_list_cache: &mut AccessListCache,
|
||||||
|
torrent_map: &mut TorrentMap,
|
||||||
|
) {
|
||||||
let now = Instant::now();
|
let now = Instant::now();
|
||||||
|
|
||||||
torrent_map.retain(|info_hash, torrent_data| {
|
torrent_map.retain(|info_hash, torrent_data| {
|
||||||
if !access_list.allows(config.access_list.mode, &info_hash.0) {
|
if !access_list_cache
|
||||||
|
.load()
|
||||||
|
.allows(config.access_list.mode, &info_hash.0)
|
||||||
|
{
|
||||||
return false;
|
return false;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -140,44 +144,7 @@ impl TorrentMaps {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn update_access_list<C: Borrow<Config>>(
|
#[derive(Default, Clone)]
|
||||||
config: C,
|
pub struct State {
|
||||||
access_list: Rc<RefCell<AccessList>>,
|
pub access_list: Arc<AccessListArcSwap>,
|
||||||
) {
|
|
||||||
if config.borrow().access_list.mode.is_on() {
|
|
||||||
match BufferedFile::open(&config.borrow().access_list.path).await {
|
|
||||||
Ok(file) => {
|
|
||||||
let mut reader = StreamReaderBuilder::new(file).build();
|
|
||||||
let mut new_access_list = AccessList::default();
|
|
||||||
|
|
||||||
loop {
|
|
||||||
let mut buf = String::with_capacity(42);
|
|
||||||
|
|
||||||
match reader.read_line(&mut buf).await {
|
|
||||||
Ok(_) => {
|
|
||||||
if let Err(err) = new_access_list.insert_from_line(&buf) {
|
|
||||||
::log::error!(
|
|
||||||
"Couln't parse access list line '{}': {:?}",
|
|
||||||
buf,
|
|
||||||
err
|
|
||||||
);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
Err(err) => {
|
|
||||||
::log::error!("Couln't read access list line {:?}", err);
|
|
||||||
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
yield_if_needed().await;
|
|
||||||
}
|
|
||||||
|
|
||||||
*access_list.borrow_mut() = new_access_list;
|
|
||||||
}
|
|
||||||
Err(err) => {
|
|
||||||
::log::error!("Couldn't open access list file: {:?}", err)
|
|
||||||
}
|
|
||||||
};
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -2,7 +2,6 @@ use std::cell::RefCell;
|
||||||
use std::rc::Rc;
|
use std::rc::Rc;
|
||||||
use std::time::Duration;
|
use std::time::Duration;
|
||||||
|
|
||||||
use aquatic_common::access_list::AccessList;
|
|
||||||
use aquatic_common::extract_response_peers;
|
use aquatic_common::extract_response_peers;
|
||||||
use futures_lite::StreamExt;
|
use futures_lite::StreamExt;
|
||||||
use glommio::channels::channel_mesh::{MeshBuilder, Partial, Role, Senders};
|
use glommio::channels::channel_mesh::{MeshBuilder, Partial, Role, Senders};
|
||||||
|
|
@ -19,9 +18,9 @@ use crate::config::Config;
|
||||||
|
|
||||||
pub async fn run_request_worker(
|
pub async fn run_request_worker(
|
||||||
config: Config,
|
config: Config,
|
||||||
|
state: State,
|
||||||
in_message_mesh_builder: MeshBuilder<(ConnectionMeta, InMessage), Partial>,
|
in_message_mesh_builder: MeshBuilder<(ConnectionMeta, InMessage), Partial>,
|
||||||
out_message_mesh_builder: MeshBuilder<(ConnectionMeta, OutMessage), Partial>,
|
out_message_mesh_builder: MeshBuilder<(ConnectionMeta, OutMessage), Partial>,
|
||||||
access_list: AccessList,
|
|
||||||
) {
|
) {
|
||||||
let (_, mut in_message_receivers) = in_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, _) = out_message_mesh_builder.join(Role::Producer).await.unwrap();
|
||||||
|
|
@ -29,14 +28,12 @@ pub async fn run_request_worker(
|
||||||
let out_message_senders = Rc::new(out_message_senders);
|
let out_message_senders = Rc::new(out_message_senders);
|
||||||
|
|
||||||
let torrents = Rc::new(RefCell::new(TorrentMaps::default()));
|
let torrents = Rc::new(RefCell::new(TorrentMaps::default()));
|
||||||
let access_list = Rc::new(RefCell::new(access_list));
|
let access_list = state.access_list;
|
||||||
|
|
||||||
// Periodically clean torrents and update access list
|
// Periodically clean torrents
|
||||||
TimerActionRepeat::repeat(enclose!((config, torrents, access_list) move || {
|
TimerActionRepeat::repeat(enclose!((config, torrents, access_list) move || {
|
||||||
enclose!((config, torrents, access_list) move || async move {
|
enclose!((config, torrents, access_list) move || async move {
|
||||||
update_access_list(&config, access_list.clone()).await;
|
torrents.borrow_mut().clean(&config, &access_list);
|
||||||
|
|
||||||
torrents.borrow_mut().clean(&config, &*access_list.borrow());
|
|
||||||
|
|
||||||
Some(Duration::from_secs(config.cleaning.interval))
|
Some(Duration::from_secs(config.cleaning.interval))
|
||||||
})()
|
})()
|
||||||
|
|
|
||||||
|
|
@ -4,9 +4,12 @@ use std::{
|
||||||
sync::{atomic::AtomicUsize, Arc},
|
sync::{atomic::AtomicUsize, Arc},
|
||||||
};
|
};
|
||||||
|
|
||||||
use aquatic_common::{access_list::AccessList, privileges::drop_privileges_after_socket_binding};
|
use aquatic_common::{
|
||||||
use common::TlsConfig;
|
access_list::AccessListQuery, privileges::drop_privileges_after_socket_binding,
|
||||||
|
};
|
||||||
|
use common::{State, TlsConfig};
|
||||||
use glommio::{channels::channel_mesh::MeshBuilder, prelude::*};
|
use glommio::{channels::channel_mesh::MeshBuilder, prelude::*};
|
||||||
|
use signal_hook::{consts::SIGUSR1, iterator::Signals};
|
||||||
|
|
||||||
use crate::config::Config;
|
use crate::config::Config;
|
||||||
|
|
||||||
|
|
@ -19,18 +22,44 @@ pub const APP_NAME: &str = "aquatic_ws: WebTorrent tracker";
|
||||||
|
|
||||||
const SHARED_CHANNEL_SIZE: usize = 1024;
|
const SHARED_CHANNEL_SIZE: usize = 1024;
|
||||||
|
|
||||||
pub fn run(config: Config) -> anyhow::Result<()> {
|
pub fn run(config: Config) -> ::anyhow::Result<()> {
|
||||||
if config.cpu_pinning.active {
|
if config.cpu_pinning.active {
|
||||||
core_affinity::set_for_current(core_affinity::CoreId {
|
core_affinity::set_for_current(core_affinity::CoreId {
|
||||||
id: config.cpu_pinning.offset,
|
id: config.cpu_pinning.offset,
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
let access_list = if config.access_list.mode.is_on() {
|
let state = State::default();
|
||||||
AccessList::create_from_path(&config.access_list.path).expect("Load access list")
|
|
||||||
} else {
|
update_access_list(&config, &state)?;
|
||||||
AccessList::default()
|
|
||||||
};
|
let mut signals = Signals::new(::std::iter::once(SIGUSR1))?;
|
||||||
|
|
||||||
|
{
|
||||||
|
let config = config.clone();
|
||||||
|
let state = state.clone();
|
||||||
|
|
||||||
|
::std::thread::spawn(move || run_inner(config, state));
|
||||||
|
}
|
||||||
|
|
||||||
|
for signal in &mut signals {
|
||||||
|
match signal {
|
||||||
|
SIGUSR1 => {
|
||||||
|
let _ = update_access_list(&config, &state);
|
||||||
|
}
|
||||||
|
_ => unreachable!(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn run_inner(config: Config, state: State) -> anyhow::Result<()> {
|
||||||
|
if config.cpu_pinning.active {
|
||||||
|
core_affinity::set_for_current(core_affinity::CoreId {
|
||||||
|
id: config.cpu_pinning.offset,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
let num_peers = config.socket_workers + config.request_workers;
|
let num_peers = config.socket_workers + config.request_workers;
|
||||||
|
|
||||||
|
|
@ -45,11 +74,11 @@ pub fn run(config: Config) -> anyhow::Result<()> {
|
||||||
|
|
||||||
for i in 0..(config.socket_workers) {
|
for i in 0..(config.socket_workers) {
|
||||||
let config = config.clone();
|
let config = config.clone();
|
||||||
|
let state = state.clone();
|
||||||
let tls_config = tls_config.clone();
|
let tls_config = tls_config.clone();
|
||||||
let request_mesh_builder = request_mesh_builder.clone();
|
let request_mesh_builder = request_mesh_builder.clone();
|
||||||
let response_mesh_builder = response_mesh_builder.clone();
|
let response_mesh_builder = response_mesh_builder.clone();
|
||||||
let num_bound_sockets = num_bound_sockets.clone();
|
let num_bound_sockets = num_bound_sockets.clone();
|
||||||
let access_list = access_list.clone();
|
|
||||||
|
|
||||||
let mut builder = LocalExecutorBuilder::default();
|
let mut builder = LocalExecutorBuilder::default();
|
||||||
|
|
||||||
|
|
@ -60,11 +89,11 @@ pub fn run(config: Config) -> anyhow::Result<()> {
|
||||||
let executor = builder.spawn(|| async move {
|
let executor = builder.spawn(|| async move {
|
||||||
network::run_socket_worker(
|
network::run_socket_worker(
|
||||||
config,
|
config,
|
||||||
|
state,
|
||||||
tls_config,
|
tls_config,
|
||||||
request_mesh_builder,
|
request_mesh_builder,
|
||||||
response_mesh_builder,
|
response_mesh_builder,
|
||||||
num_bound_sockets,
|
num_bound_sockets,
|
||||||
access_list,
|
|
||||||
)
|
)
|
||||||
.await
|
.await
|
||||||
});
|
});
|
||||||
|
|
@ -74,9 +103,9 @@ pub fn run(config: Config) -> anyhow::Result<()> {
|
||||||
|
|
||||||
for i in 0..(config.request_workers) {
|
for i in 0..(config.request_workers) {
|
||||||
let config = config.clone();
|
let config = config.clone();
|
||||||
|
let state = state.clone();
|
||||||
let request_mesh_builder = request_mesh_builder.clone();
|
let request_mesh_builder = request_mesh_builder.clone();
|
||||||
let response_mesh_builder = response_mesh_builder.clone();
|
let response_mesh_builder = response_mesh_builder.clone();
|
||||||
let access_list = access_list.clone();
|
|
||||||
|
|
||||||
let mut builder = LocalExecutorBuilder::default();
|
let mut builder = LocalExecutorBuilder::default();
|
||||||
|
|
||||||
|
|
@ -85,12 +114,7 @@ pub fn run(config: Config) -> anyhow::Result<()> {
|
||||||
}
|
}
|
||||||
|
|
||||||
let executor = builder.spawn(|| async move {
|
let executor = builder.spawn(|| async move {
|
||||||
handlers::run_request_worker(
|
handlers::run_request_worker(config, state, request_mesh_builder, response_mesh_builder)
|
||||||
config,
|
|
||||||
request_mesh_builder,
|
|
||||||
response_mesh_builder,
|
|
||||||
access_list,
|
|
||||||
)
|
|
||||||
.await
|
.await
|
||||||
});
|
});
|
||||||
|
|
||||||
|
|
@ -142,3 +166,20 @@ fn create_tls_config(config: &Config) -> anyhow::Result<TlsConfig> {
|
||||||
|
|
||||||
Ok(tls_config)
|
Ok(tls_config)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn update_access_list(config: &Config, state: &State) -> anyhow::Result<()> {
|
||||||
|
if config.access_list.mode.is_on() {
|
||||||
|
match state.access_list.update(&config.access_list) {
|
||||||
|
Ok(()) => {
|
||||||
|
::log::info!("Access list updated")
|
||||||
|
}
|
||||||
|
Err(err) => {
|
||||||
|
::log::error!("Updating access list failed: {:#}", err);
|
||||||
|
|
||||||
|
return Err(err);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
|
||||||
|
|
@ -7,7 +7,7 @@ use std::sync::atomic::{AtomicUsize, Ordering};
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
use std::time::Duration;
|
use std::time::Duration;
|
||||||
|
|
||||||
use aquatic_common::access_list::AccessList;
|
use aquatic_common::access_list::{create_access_list_cache, AccessListArcSwap, AccessListCache};
|
||||||
use aquatic_common::convert_ipv4_mapped_ipv6;
|
use aquatic_common::convert_ipv4_mapped_ipv6;
|
||||||
use aquatic_ws_protocol::*;
|
use aquatic_ws_protocol::*;
|
||||||
use async_tungstenite::WebSocketStream;
|
use async_tungstenite::WebSocketStream;
|
||||||
|
|
@ -40,14 +40,14 @@ struct ConnectionReference {
|
||||||
|
|
||||||
pub async fn run_socket_worker(
|
pub async fn run_socket_worker(
|
||||||
config: Config,
|
config: Config,
|
||||||
|
state: State,
|
||||||
tls_config: Arc<TlsConfig>,
|
tls_config: Arc<TlsConfig>,
|
||||||
in_message_mesh_builder: MeshBuilder<(ConnectionMeta, InMessage), Partial>,
|
in_message_mesh_builder: MeshBuilder<(ConnectionMeta, InMessage), Partial>,
|
||||||
out_message_mesh_builder: MeshBuilder<(ConnectionMeta, OutMessage), Partial>,
|
out_message_mesh_builder: MeshBuilder<(ConnectionMeta, OutMessage), Partial>,
|
||||||
num_bound_sockets: Arc<AtomicUsize>,
|
num_bound_sockets: Arc<AtomicUsize>,
|
||||||
access_list: AccessList,
|
|
||||||
) {
|
) {
|
||||||
let config = Rc::new(config);
|
let config = Rc::new(config);
|
||||||
let access_list = Rc::new(RefCell::new(access_list));
|
let access_list = state.access_list;
|
||||||
|
|
||||||
let listener = TcpListener::bind(config.network.address).expect("bind socket");
|
let listener = TcpListener::bind(config.network.address).expect("bind socket");
|
||||||
num_bound_sockets.fetch_add(1, Ordering::SeqCst);
|
num_bound_sockets.fetch_add(1, Ordering::SeqCst);
|
||||||
|
|
@ -62,15 +62,6 @@ pub async fn run_socket_worker(
|
||||||
let connection_slab = Rc::new(RefCell::new(Slab::new()));
|
let connection_slab = Rc::new(RefCell::new(Slab::new()));
|
||||||
let connections_to_remove = Rc::new(RefCell::new(Vec::new()));
|
let connections_to_remove = Rc::new(RefCell::new(Vec::new()));
|
||||||
|
|
||||||
// Periodically update access list
|
|
||||||
TimerActionRepeat::repeat(enclose!((config, access_list) move || {
|
|
||||||
enclose!((config, access_list) move || async move {
|
|
||||||
update_access_list(config.clone(), access_list.clone()).await;
|
|
||||||
|
|
||||||
Some(Duration::from_secs(config.cleaning.interval))
|
|
||||||
})()
|
|
||||||
}));
|
|
||||||
|
|
||||||
// Periodically remove closed connections
|
// Periodically remove closed connections
|
||||||
TimerActionRepeat::repeat(
|
TimerActionRepeat::repeat(
|
||||||
enclose!((config, connection_slab, connections_to_remove) move || {
|
enclose!((config, connection_slab, connections_to_remove) move || {
|
||||||
|
|
@ -176,7 +167,7 @@ struct Connection;
|
||||||
impl Connection {
|
impl Connection {
|
||||||
async fn run(
|
async fn run(
|
||||||
config: Rc<Config>,
|
config: Rc<Config>,
|
||||||
access_list: Rc<RefCell<AccessList>>,
|
access_list: Arc<AccessListArcSwap>,
|
||||||
in_message_senders: Rc<Senders<(ConnectionMeta, InMessage)>>,
|
in_message_senders: Rc<Senders<(ConnectionMeta, InMessage)>>,
|
||||||
out_message_sender: Rc<LocalSender<(ConnectionMeta, OutMessage)>>,
|
out_message_sender: Rc<LocalSender<(ConnectionMeta, OutMessage)>>,
|
||||||
out_message_receiver: LocalReceiver<(ConnectionMeta, OutMessage)>,
|
out_message_receiver: LocalReceiver<(ConnectionMeta, OutMessage)>,
|
||||||
|
|
@ -201,11 +192,12 @@ impl Connection {
|
||||||
let (ws_out, ws_in) = futures::StreamExt::split(stream);
|
let (ws_out, ws_in) = futures::StreamExt::split(stream);
|
||||||
|
|
||||||
let pending_scrape_slab = Rc::new(RefCell::new(Slab::new()));
|
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(enclose!((pending_scrape_slab) async move {
|
let reader_handle = spawn_local(enclose!((pending_scrape_slab) async move {
|
||||||
let mut reader = ConnectionReader {
|
let mut reader = ConnectionReader {
|
||||||
config,
|
config,
|
||||||
access_list,
|
access_list_cache,
|
||||||
in_message_senders,
|
in_message_senders,
|
||||||
out_message_sender,
|
out_message_sender,
|
||||||
pending_scrape_slab,
|
pending_scrape_slab,
|
||||||
|
|
@ -237,7 +229,7 @@ impl Connection {
|
||||||
|
|
||||||
struct ConnectionReader {
|
struct ConnectionReader {
|
||||||
config: Rc<Config>,
|
config: Rc<Config>,
|
||||||
access_list: Rc<RefCell<AccessList>>,
|
access_list_cache: AccessListCache,
|
||||||
in_message_senders: Rc<Senders<(ConnectionMeta, InMessage)>>,
|
in_message_senders: Rc<Senders<(ConnectionMeta, InMessage)>>,
|
||||||
out_message_sender: Rc<LocalSender<(ConnectionMeta, OutMessage)>>,
|
out_message_sender: Rc<LocalSender<(ConnectionMeta, OutMessage)>>,
|
||||||
pending_scrape_slab: Rc<RefCell<Slab<PendingScrapeResponse>>>,
|
pending_scrape_slab: Rc<RefCell<Slab<PendingScrapeResponse>>>,
|
||||||
|
|
@ -275,8 +267,8 @@ impl ConnectionReader {
|
||||||
let info_hash = announce_request.info_hash;
|
let info_hash = announce_request.info_hash;
|
||||||
|
|
||||||
if self
|
if self
|
||||||
.access_list
|
.access_list_cache
|
||||||
.borrow()
|
.load()
|
||||||
.allows(self.config.access_list.mode, &info_hash.0)
|
.allows(self.config.access_list.mode, &info_hash.0)
|
||||||
{
|
{
|
||||||
let in_message = InMessage::AnnounceRequest(announce_request);
|
let in_message = InMessage::AnnounceRequest(announce_request);
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue