ptth/crates/ptth_relay/src/relay_state.rs

331 lines
7.0 KiB
Rust

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,
path: 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
}
}