use std::{ collections::HashMap, convert::TryFrom, path::Path, time::Instant, }; use chrono::{DateTime, Utc}; use dashmap::DashMap; use tokio::sync::{ Mutex, RwLock, oneshot, watch, }; use crate::{ Body, Config, RelayError, ShuttingDownError, config::machine_editable, }; use ptth_core::{ http_serde, prelude::*, }; /* Here's what we need to handle: When a request comes in: - Park the client in response_rendezvous - Look up the server ID in request_rendezvous - If a server is parked, unpark it and send the request - Otherwise, queue the request When a server comes to listen: - Look up the server ID in request_rendezvous - Either return all pending requests, or park the server When a server comes to respond: - Look up the parked client in response_rendezvous - Unpark the client and begin streaming So we need these lookups to be fast: - Server IDs, where (1 server) or (0 or many clients) can be parked - Request IDs, where 1 client is parked */ pub enum RequestRendezvous { ParkedClients (Vec ), ParkedServer (oneshot::Sender >), } pub (crate) struct ResponseRendezvous { pub timeout: Instant, pub tx: oneshot::Sender >, } #[derive (Clone)] pub struct ServerStatus { pub last_seen: DateTime , } impl Default for ServerStatus { fn default () -> Self { Self { last_seen: Utc::now (), } } } pub struct Relay { pub (crate) config: RwLock , pub (crate) me_config: RwLock , /// The parked clients or parked server, keyed by server pub (crate) request_rendezvous: Mutex >, pub (crate) server_status: Mutex >, /// The parked requests, keyed by request ID pub (crate) response_rendezvous: RwLock >, pub (crate) shutdown_watch_tx: watch::Sender , pub (crate) shutdown_watch_rx: watch::Receiver , /// List of recently rejected server names (Used to approve servers) pub (crate) unregistered_servers: BoundedVec , /// Memory backend for audit logging // TODO: Add file / database / network server logging backend pub (crate) audit_log: BoundedVec , } #[derive (Clone)] pub struct RejectedServer { pub name: String, pub tripcode: blake3::Hash, pub seen: DateTime , } #[derive (Clone, Debug)] pub struct AuditEvent { time_monotonic: Instant, pub time_utc: DateTime , pub data: AuditData, } #[derive (Clone, Debug)] pub enum AuditData { RegisterServer { user: Option , server: crate::config::file::Server, }, RelayStart, ScraperGet { key_name: String, path: String, }, WebClientGet { user: Option , server_name: String, uri: String, }, } impl AuditEvent { pub fn new (data: AuditData) -> Self { Self { time_monotonic: Instant::now (), time_utc: Utc::now (), data, } } } pub struct BoundedVec { bound: usize, v: RwLock >, } impl BoundedVec { pub fn new (bound: usize) -> Self { Self { bound, v: Default::default (), } } pub async fn to_vec (&self) -> Vec { let guard = self.v.read ().await; (*guard).clone () } // Not mut because we have a RwLock // One of the only problems with Rust is that // 'mut' doesn't _really_ -mean 'mutable' once you're // multi-threading or using async pub async fn push (&self, x: T) { let mut guard = self.v.write ().await; guard.push (x); while guard.len () > self.bound { guard.remove (0); } } } impl TryFrom for Relay { type Error = RelayError; fn try_from (config: Config) -> Result { let (shutdown_watch_tx, shutdown_watch_rx) = watch::channel (false); let me_config = machine_editable::Config::default (); let me_config = match machine_editable::Config::from_file (Path::new ("data/ptth_relay_me_config.toml")) { Err (e) => { warn! ("Can't load machine-editable config: {:?}", e); me_config }, Ok (x) => x, }; Ok (Self { config: config.into (), me_config: me_config.into (), request_rendezvous: Default::default (), server_status: Default::default (), response_rendezvous: Default::default (), shutdown_watch_tx, shutdown_watch_rx, unregistered_servers: BoundedVec::new (20), audit_log: BoundedVec::new (256), }) } } impl Relay { /// Returns a `Vec` of the names of currently connected servers pub async fn list_servers (&self) -> Vec { self.request_rendezvous.lock ().await.iter () .map (|(k, _)| (*k).clone ()) .collect () } pub async fn server_exists (&self, name: &str) -> bool { { let config = self.config.read ().await; if config.servers.contains_key (name) { return true; } } { let config = self.me_config.read ().await; if config.servers.contains_key (name) { return true; } } false } #[must_use] pub fn build () -> Builder { Builder::default () } pub async fn park_client ( &self, server_name: &str, req: http_serde::RequestParts, req_id: &str, ) { use RequestRendezvous::*; let mut request_rendezvous = self.request_rendezvous.lock ().await; let wrapped = http_serde::WrappedRequest { id: req_id.to_string (), req, }; let new_rendezvous = match request_rendezvous.remove (server_name) { Some (ParkedClients (mut v)) => { debug! ("Parking request {} ({} already queued)", req_id, v.len ()); v.push (wrapped); ParkedClients (v) }, Some (ParkedServer (s)) => { // If sending to the server fails, queue it match s.send (Ok (wrapped)) { Ok (()) => { // TODO: This can actually still fail, if the server // disconnects right as we're sending this. // Then what? trace! ( "Sending request {} directly to server {}", req_id, server_name, ); ParkedClients (vec! []) }, Err (Ok (wrapped)) => { debug! ("Parking request {}", req_id); ParkedClients (vec! [wrapped]) }, Err (_) => unreachable! (), } }, None => { debug! ("Parking request {}", req_id); ParkedClients (vec! [wrapped]) }, }; request_rendezvous.insert (server_name.to_string (), new_rendezvous); } } #[derive (Default)] pub struct Builder { config: Config, } impl Builder { pub fn build (self) -> Result { Relay::try_from (self.config) } pub fn address (mut self, addr: std::net::IpAddr) -> Self { self.config.address = addr; self } pub fn enable_scraper_api (mut self, b: bool) -> Self { self.config.iso.enable_scraper_api = b; self } pub fn port (mut self, port: u16) -> Self { self.config.port = Some (port); self } pub fn scraper_key (mut self, key: crate::key_validity::ScraperKey) -> Self { self.config.scraper_keys.insert (key.hash.encode_base64 (), key); self } pub fn server (mut self, server: crate::config::file::Server) -> Self { self.config.servers.insert (server.name.clone (), server); self } }