use troll_patrol::{ extra_info::{self, ExtraInfo}, //negative_report::SerializableNegativeReport, //positive_report::SerializablePositiveReport, request_handler::handle, *, }; use clap::Parser; use futures::future; use hyper::{ server::conn::AddrStream, service::{make_service_fn, service_fn}, Body, Request, Response, Server, }; use serde::Deserialize; use sled::Db; use std::{ collections::HashSet, convert::Infallible, fs::File, io::BufReader, net::SocketAddr, path::PathBuf, time::Duration, }; 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 Troll Patrol 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, } #[derive(Debug, Deserialize)] pub struct Config { pub db: DbConfig, //require_bridge_token: bool, port: u16, } #[derive(Debug, Deserialize)] pub struct DbConfig { // The path for the server database, default is "server_db" pub db_path: String, } impl Default for DbConfig { fn default() -> DbConfig { DbConfig { db_path: "server_db".to_owned(), } } } async fn update_extra_infos(db: &Db) { // Track which files have been processed. This is slightly redundant // because we're only downloading files we don't already have, but it // might be a good idea to check in case we downloaded a file but didn't // process it for some reason. let mut processed_extra_infos_files = match db.get(b"extra_infos_files").unwrap() { Some(v) => bincode::deserialize(&v).unwrap(), None => HashSet::::new(), }; let new_files = extra_info::download_extra_infos().await.unwrap(); let mut new_extra_infos = HashSet::::new(); // Make set of new extra-infos for extra_info_file in &new_files { extra_info::add_extra_infos(&extra_info_file, &mut new_extra_infos); processed_extra_infos_files.insert(extra_info_file.to_string()); } // Add new extra-infos data to database for extra_info in new_extra_infos { add_extra_info_to_db(&db, extra_info); } db.insert( b"extra_infos_files", bincode::serialize(&processed_extra_infos_files).unwrap(), ) .unwrap(); } async fn create_context_manager( db_config: DbConfig, context_rx: mpsc::Receiver, mut kill: broadcast::Receiver<()>, ) { tokio::select! { create_context = context_manager(db_config, context_rx) => create_context, _ = kill.recv() => {println!("Shut down manager");}, } } async fn context_manager(db_config: DbConfig, mut context_rx: mpsc::Receiver) { let db: Db = sled::open(&db_config.db_path).unwrap(); while let Some(cmd) = context_rx.recv().await { use Command::*; match cmd { Request { req, sender } => { let response = handle(&db, 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 can be handled #[derive(Debug)] enum Command { Request { req: Request, sender: oneshot::Sender, Infallible>>, }, Shutdown { shutdown_sig: broadcast::Sender<()>, }, } #[tokio::main] async fn main() { // TODO: Currently, we're processing extra-infos here, but we want to: // 2. Periodically (daily): // a) download new extra-infos // b) determine whether we think each bridge is blocked or not // c) report these results to the LA // 3. Store all our data let args: Args = Args::parse(); let config: Config = serde_json::from_reader(BufReader::new( File::open(&args.config).expect("Could not read config file"), )) .expect("Reading config file from JSON failed"); let (request_tx, request_rx) = mpsc::channel(32); let shutdown_cmd_tx = request_tx.clone(); // create the shutdown broadcast channel and clone for every thread let (shutdown_tx, mut shutdown_rx) = broadcast::channel(16); let kill = 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(config.db, request_rx, kill).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(([0, 0, 0, 0], config.port)); 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([context_manager, shutdown_handler]).await; }