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

use futures::future;
use futures::StreamExt;
use hyper::{
    server::conn::AddrStream,
    service::{make_service_fn, service_fn},
    Body, Request, Response, Server,
};
use lox::bridge_table::{BridgeLine, MAX_BRIDGES_PER_BUCKET};
use lox::{BridgeAuth, BridgeDb};

use rdsys_backend::{proto::ResourceDiff, start_stream};
use serde::Deserialize;
use std::{
    convert::Infallible,
    env,
    fs::File,
    io::BufReader,
    net::SocketAddr,
    sync::{Arc, Mutex},
    time::Duration,
};

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(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(
    context_rx: mpsc::Receiver<Command>,
    mut kill: broadcast::Receiver<()>,
) {
    tokio::select! {
        create_context = context_manager(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(mut context_rx: mpsc::Receiver<Command>) {
    let bridgedb = BridgeDb::new();
    let lox_auth = BridgeAuth::new(bridgedb.pubkey);

    let context = lox_context::LoxServerContext {
        db: Arc::new(Mutex::new(bridgedb)),
        ba: Arc::new(Mutex::new(lox_auth)),
        extra_bridges: Arc::new(Mutex::new(Vec::new())),
        unreplaced_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);
                        for resource in pt.1 {
                            let bridgeline = parse_resource(resource);
                            println!("Now it's a bridgeline: {:?}", bridgeline);
                            if context.unreplaced_bridges.lock().unwrap().len() > 0 {
                                println!("BridgeLine to be replaced: {:?}", bridgeline);
                                let res = context.replace_with_new(bridgeline);
                                if res {
                                    println!("BridgeLine successfully replaced: {:?}", bridgeline);
                                } else {
                                    // Add the bridge to the list of unreplaced 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_unreplaced_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);
                        for resource in pt.1 {
                            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);
                        for resource in pt.1 {
                            let bridgeline = parse_resource(resource);
                            println!("BridgeLine to be replaced: {:?}", bridgeline);
                            let res = context.replace_with_new(bridgeline);
                            if res {
                                println!("BridgeLine successfully replaced: {:?}", bridgeline);
                            } else {
                                // Add the bridge to the list of unreplaced 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_unreplaced_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();
                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: Vec<String> = env::args().collect();
    let file = File::open(&args[1]).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(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;
}