334 lines
7.0 KiB
Rust
334 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>>),
|
|
}
|
|
|
|
pub (crate) struct ResponseRendezvous {
|
|
pub timeout: Instant,
|
|
pub tx: 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
|
|
}
|
|
}
|