ptth/crates/ptth_relay/src/relay_state.rs

use std::{
	collections::HashMap,
	convert::TryFrom,
	time::Instant,
};

use chrono::{DateTime, Utc};
use dashmap::DashMap;
use tokio::sync::{
	Mutex,
	RwLock,
	oneshot,
	watch,
};

use crate::{
	Body,
	Config,
	RelayError,
	ShuttingDownError,
};

use ptth_core::http_serde;

/*

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>,
	
	// Key: Server ID
	pub (crate) request_rendezvous: Mutex <HashMap <String, RequestRendezvous>>,
	pub (crate) server_status: Mutex <HashMap <String, ServerStatus>>,
	
	// Key: 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,
	time_utc: DateTime <Utc>,
	data: AuditData,
}

#[derive (Clone, Debug)]
pub enum AuditData {
	RelayStart,
	WebClientGet {
		req_id: String,
		server_name: 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);
		
		Ok (Self {
			config: 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 fn build () -> Builder {
		Builder::default ()
	}
}

#[derive (Default)]
pub struct Builder {
	config: Config,
}

impl Builder {
	pub fn build (self) -> Result <Relay, RelayError> {
		Relay::try_from (self.config)
	}
	
	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 <crate::key_validity::Valid30Days>) 
	-> 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
	}
}