lox/crates/lox-distributor/src/main.rs

402 lines
17 KiB
Rust

use chrono::Utc;
use clap::Parser;
use futures::future;
use futures::StreamExt;
use hyper::{
server::conn::AddrStream,
service::{make_service_fn, service_fn},
Body, Request, Response, Server,
};
use lox_library::bridge_table::{BridgeLine, MAX_BRIDGES_PER_BUCKET};
use lox_library::{BridgeAuth, BridgeDb};
use rdsys_backend::{proto::ResourceDiff, start_stream};
use serde::Deserialize;
use std::{
convert::Infallible,
fs::File,
io::BufReader,
net::SocketAddr,
path::PathBuf,
sync::{Arc, Mutex},
time::Duration,
};
mod file_reader;
use file_reader::{check_db_exists, read_context_from_file, write_context_to_file};
mod lox_context;
mod request_handler;
use request_handler::handle;
mod resource_parser;
use resource_parser::parse_resource;
use tokio::{
signal, spawn,
sync::{broadcast, mpsc, oneshot},
time::sleep,
};
async fn shutdown_signal() {
tokio::signal::ctrl_c()
.await
.expect("failed to listen for ctrl+c signal");
println!("Shut down Lox Server");
}
#[derive(Parser, Debug)]
#[command(author, version, about, long_about = None)]
struct Args {
/// Name/path of the configuration file
#[arg(short, long, default_value = "config.json")]
config: PathBuf,
/// Optional name/path of the db file
#[arg(short, long)]
db: Option<PathBuf>,
}
#[derive(Debug, Deserialize)]
struct ResourceInfo {
endpoint: String,
name: String,
token: String,
types: Vec<String>,
}
// Populate Bridgedb from rdsys
// Rdsys sender creates a ResourceStream with the api_endpoint, resource token and type specified
// in the config.json file.
// TODO: ensure this stream gracefully shutdowns on the ctrl_c command.
async fn rdsys_stream(
rtype: ResourceInfo,
tx: mpsc::Sender<ResourceDiff>,
mut kill: broadcast::Receiver<()>,
) {
let mut rstream = start_stream(rtype.endpoint, rtype.name, rtype.token, rtype.types)
.await
.expect("rdsys stream initialization failed. Start rdsys or check config.json");
loop {
tokio::select! {
res = rstream.next() => {
match res {
Some(diff) => tx.send(diff).await.unwrap(),
None => return,
}
},
_ = kill.recv() => {println!("Shut down rdsys stream"); return},
}
}
}
async fn rdsys_bridge_parser(
rdsys_tx: mpsc::Sender<Command>,
rx: mpsc::Receiver<ResourceDiff>,
mut kill: broadcast::Receiver<()>,
) {
tokio::select! {
start_bridge_parser = parse_bridges(rdsys_tx, rx) => start_bridge_parser ,
_ = kill.recv() => {println!("Shut down bridge_parser");},
}
}
// Parse Bridges receives a ResourceDiff from rdsys_sender and sends it to the
// Context Manager to be parsed and added to the BridgeDB
async fn parse_bridges(rdsys_tx: mpsc::Sender<Command>, mut rx: mpsc::Receiver<ResourceDiff>) {
loop {
let resourcediff = rx.recv().await.unwrap();
let cmd = Command::Rdsys { resourcediff };
rdsys_tx.send(cmd).await.unwrap();
sleep(Duration::from_secs(1)).await;
}
}
async fn create_context_manager(
db_path: Option<PathBuf>,
context_rx: mpsc::Receiver<Command>,
mut kill: broadcast::Receiver<()>,
) {
tokio::select! {
create_context = context_manager(db_path, context_rx) => create_context,
_ = kill.recv() => {println!("Shut down context_manager");},
}
}
// Context Manager handles the Lox BridgeDB and Bridge Authority, ensuring
// that the DB can be updated from the rdsys stream and client requests
// can be responded to with an updated BridgeDB state
async fn context_manager(db_path: Option<PathBuf>, mut context_rx: mpsc::Receiver<Command>) {
let context: lox_context::LoxServerContext;
if let Some(existing_db) = db_path.as_deref() {
context = read_context_from_file(existing_db).unwrap();
} else if let Some(last_modified_file) = check_db_exists() {
println!("Reading from file {:?}", last_modified_file);
context = read_context_from_file(&last_modified_file.path()).unwrap();
} else {
let new_db = BridgeDb::new();
let new_ba = BridgeAuth::new(new_db.pubkey);
context = lox_context::LoxServerContext {
db: Arc::new(Mutex::new(new_db)),
ba: Arc::new(Mutex::new(new_ba)),
extra_bridges: Arc::new(Mutex::new(Vec::new())),
to_be_replaced_bridges: Arc::new(Mutex::new(Vec::new())),
}
}
while let Some(cmd) = context_rx.recv().await {
use Command::*;
match cmd {
Rdsys { resourcediff } => {
if let Some(new_resources) = resourcediff.new {
let mut count = 0;
let mut bucket = [BridgeLine::default(); MAX_BRIDGES_PER_BUCKET];
for pt in new_resources {
println!("A NEW RESOURCE: {:?}", pt);
if let Some(resources) = pt.1 {
for resource in resources {
let bridgeline = parse_resource(resource);
println!("Now it's a bridgeline: {:?}", bridgeline);
if context.to_be_replaced_bridges.lock().unwrap().len() > 0 {
println!("BridgeLine to be replaced: {:?}", bridgeline);
let res = context.replace_with_new(bridgeline);
if res == lox_library::ReplaceSuccess::NotFound {
println!(
"BridgeLine not found in bridge_table, already updated {:?}",
bridgeline
);
} else if res == lox_library::ReplaceSuccess::Replaced {
println!(
"BridgeLine successfully replaced: {:?}",
bridgeline
);
} else {
assert!(
res == lox_library::ReplaceSuccess::NotReplaced,
"ReplaceSuccess incorrectly set somehow"
);
// Add the bridge to the list of to_be_replaced bridges in the Lox context and try
// again to replace at the next update (nothing changes in the Lox Authority)
println!("'Gone' BridgeLine NOT replaced, saved for next update! : {:?}", bridgeline);
context.new_to_be_replaced_bridge(bridgeline);
}
} else if count < MAX_BRIDGES_PER_BUCKET {
bucket[count] = bridgeline;
count += 1;
} else {
// TODO: Decide the circumstances under which a bridge is allocated to an open_inv or spare bucket,
// eventually also do some more fancy grouping of new resources, i.e., by type or region
context.add_openinv_bucket(bucket);
count = 0;
bucket = [BridgeLine::default(); MAX_BRIDGES_PER_BUCKET];
}
}
}
}
// Handle the extra buckets that were not allocated already
if count != 0 {
for val in 0..count {
if context.extra_bridges.lock().unwrap().len()
< (MAX_BRIDGES_PER_BUCKET)
{
context.append_extra_bridges(bucket[val]);
} else {
bucket = context.remove_extra_bridges();
context.add_spare_bucket(bucket);
}
}
}
}
if let Some(changed_resources) = resourcediff.changed {
for pt in changed_resources {
println!("A NEW CHANGED RESOURCE: {:?}", pt);
if let Some(resources) = pt.1 {
for resource in resources {
let bridgeline = parse_resource(resource);
println!("BridgeLine to be changed: {:?}", bridgeline);
let res = context.update_bridge(bridgeline);
if res {
println!("BridgeLine successfully updated: {:?}", bridgeline);
} else {
println!("BridgeLine: {:?} not found in Lox's Bridgetable. Save it as a new resource for now!", bridgeline);
if context.extra_bridges.lock().unwrap().len() < 2 {
context.append_extra_bridges(bridgeline);
} else {
let bucket = context.remove_extra_bridges();
context.add_spare_bucket(bucket);
}
}
}
}
}
}
// gone resources are not the same as blocked resources.
// Instead, these are bridges which have either failed to pass tests for some period
// or have expired bridge descriptors. In both cases, the bridge is unusable, but this
// is not likely due to censorship. Therefore, we replace gone resources with new resources
// TODO: create a notion of blocked resources from information collected through various means:
// https://gitlab.torproject.org/tpo/anti-censorship/censorship-analysis/-/issues/40035
if let Some(gone_resources) = resourcediff.gone {
for pt in gone_resources {
println!("A NEW GONE RESOURCE: {:?}", pt);
if let Some(resources) = pt.1 {
for resource in resources {
// If resource last passed tests 3 hours ago, it should be replaced with a working
// resource and be removed from the bridgetable. If it has been gone for more than 7 hours,
// we should stop trying to remove it from the bridge table and assume it has successfully been
// removed already
if resource.last_passed < (Utc::now() - chrono::Duration::hours(3))
|| resource.last_passed
> (Utc::now() - chrono::Duration::hours(7))
{
let bridgeline = parse_resource(resource);
println!("BridgeLine to be replaced: {:?}", bridgeline);
let res = context.replace_with_new(bridgeline);
if res == lox_library::ReplaceSuccess::Replaced {
println!(
"BridgeLine successfully replaced: {:?}",
bridgeline
);
} else if res == lox_library::ReplaceSuccess::NotReplaced {
// Add the bridge to the list of to_be_replaced bridges in the Lox context and try
// again to replace at the next update (nothing changes in the Lox Authority)
println!(
"'Gone' BridgeLine NOT replaced, saved for next update! : {:?}",
bridgeline
);
context.new_to_be_replaced_bridge(bridgeline);
}
}
}
}
}
}
/* Functionality for marking bridges as unreachable/blocked is currently not enabled as there is not
yet a reliable way to determine that a bridge is blocked. This means that migrations to unblocked bridges do not
currently work but can be easily enabled with a list of `blocked resources` from rdsys or another source with something
like the following:
println!("BridgeLine to be removed: {:?}", bridgeline);
let res = context.add_unreachable(bridgeline);
if res {
println!(
"BridgeLine successfully marked unreachable: {:?}",
bridgeline
);
} else {
println!("'Gone' BridgeLine NOT REMOVED!! : {:?}", bridgeline);
//TODO probably do something else here
}
*/
context.allocate_leftover_bridges();
context.encrypt_table();
write_context_to_file(context.clone());
sleep(Duration::from_millis(1)).await;
}
Request { req, sender } => {
let response = handle(context.clone(), req).await;
if let Err(e) = sender.send(response) {
eprintln!("Server Response Error: {:?}", e);
};
sleep(Duration::from_millis(1)).await;
}
Shutdown { shutdown_sig } => {
println!("Sending Shutdown Signal, all threads should shutdown.");
drop(shutdown_sig);
println!("Shutdown Sent.");
}
}
}
}
// Each of the commands that the Context Manager handles
#[derive(Debug)]
enum Command {
Rdsys {
resourcediff: ResourceDiff,
},
Request {
req: Request<Body>,
sender: oneshot::Sender<Result<Response<Body>, Infallible>>,
},
Shutdown {
shutdown_sig: broadcast::Sender<()>,
},
}
#[tokio::main]
async fn main() {
let args: Args = Args::parse();
let file = File::open(&args.config).expect("Should have been able to read config.json file");
let reader = BufReader::new(file);
// Read the JSON contents of the file as a ResourceInfo
let rtype: ResourceInfo =
serde_json::from_reader(reader).expect("Reading ResourceInfo from JSON failed.");
let (rdsys_tx, context_rx) = mpsc::channel(32);
let request_tx = rdsys_tx.clone();
let shutdown_cmd_tx = rdsys_tx.clone();
// create the shutdown broadcast channel and clone for every thread
let (shutdown_tx, mut shutdown_rx) = broadcast::channel(16);
let kill_stream = shutdown_tx.subscribe();
let kill_parser = shutdown_tx.subscribe();
let kill_context = shutdown_tx.subscribe();
// Listen for ctrl_c, send signal to broadcast shutdown to all threads by dropping shutdown_tx
let shutdown_handler = spawn(async move {
tokio::select! {
_ = signal::ctrl_c() => {
let cmd = Command::Shutdown {
shutdown_sig: shutdown_tx,
};
shutdown_cmd_tx.send(cmd).await.unwrap();
sleep(Duration::from_secs(1)).await;
_ = shutdown_rx.recv().await;
}
}
});
let context_manager =
spawn(async move { create_context_manager(args.db, context_rx, kill_context).await });
let (tx, rx) = mpsc::channel(32);
let rdsys_stream_handler = spawn(async { rdsys_stream(rtype, tx, kill_stream).await });
let rdsys_resource_receiver =
spawn(async { rdsys_bridge_parser(rdsys_tx, rx, kill_parser).await });
let make_service = make_service_fn(move |_conn: &AddrStream| {
let request_tx = request_tx.clone();
let service = service_fn(move |req| {
let request_tx = request_tx.clone();
let (response_tx, response_rx) = oneshot::channel();
let cmd = Command::Request {
req,
sender: response_tx,
};
async move {
request_tx.send(cmd).await.unwrap();
response_rx.await.unwrap()
}
});
async move { Ok::<_, Infallible>(service) }
});
let addr = SocketAddr::from(([127, 0, 0, 1], 8001));
let server = Server::bind(&addr).serve(make_service);
let graceful = server.with_graceful_shutdown(shutdown_signal());
println!("Listening on {}", addr);
if let Err(e) = graceful.await {
eprintln!("server error: {}", e);
}
future::join_all([
rdsys_stream_handler,
rdsys_resource_receiver,
context_manager,
shutdown_handler,
])
.await;
}