ptth/crates/ptth_relay/src/relay_state.rs

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 <http_serde::WrappedRequest>),
	ParkedServer (oneshot::Sender <Result <http_serde::WrappedRequest, ShuttingDownError>>),
}

type ResponseRendezvous = oneshot::Sender <Result <(http_serde::ResponseParts, Body), ShuttingDownError>>;

#[derive (Clone)]
pub struct ServerStatus {
	pub last_seen: DateTime <Utc>,
}

impl Default for ServerStatus {
	fn default () -> Self {
		Self {
			last_seen: Utc::now (),
		}
	}
}

pub struct Relay {
	pub (crate) config: RwLock <Config>,
	pub (crate) me_config: RwLock <machine_editable::Config>,
	
	/// The parked clients or parked server, keyed by server
	
	pub (crate) request_rendezvous: Mutex <HashMap <String, RequestRendezvous>>,
	pub (crate) server_status: Mutex <HashMap <String, ServerStatus>>,
	
	/// The parked requests, keyed by request ID
	
	pub (crate) response_rendezvous: RwLock <DashMap <String, ResponseRendezvous>>,
	
	pub (crate) shutdown_watch_tx: watch::Sender <bool>,
	pub (crate) shutdown_watch_rx: watch::Receiver <bool>,
	
	/// List of recently rejected server names (Used to approve servers)
	pub (crate) unregistered_servers: BoundedVec <RejectedServer>,
	
	/// Memory backend for audit logging
	// TODO: Add file / database / network server logging backend
	pub (crate) audit_log: BoundedVec <AuditEvent>,
}

#[derive (Clone)]
pub struct RejectedServer {
	pub name: String,
	pub tripcode: blake3::Hash,
	pub seen: DateTime <Utc>,
}

#[derive (Clone, Debug)]
pub struct AuditEvent {
	time_monotonic: Instant,
	pub time_utc: DateTime <Utc>,
	pub data: AuditData,
}

#[derive (Clone, Debug)]
pub enum AuditData {
	RegisterServer {
		user: Option <String>,
		server: crate::config::file::Server,
	},
	RelayStart,
	ScraperGet {
		key_name: String,
		server_name: String,
		uri: String,
	},
	WebClientGet {
		user: Option <String>,
		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 <T: Clone> {
	bound: usize,
	v: RwLock <Vec <T>>,
}

impl <T: Clone> BoundedVec <T> {
	pub fn new (bound: usize) -> Self {
		Self {
			bound,
			v: Default::default (),
		}
	}
	
	pub async fn to_vec (&self) -> Vec <T> {
		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 <Config> for Relay {
	type Error = RelayError;
	
	fn try_from (config: Config) -> Result <Self, Self::Error> {
		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 <String> {
		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, RelayError> {
		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
	}
}