use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use axum::Router;
use axum::extract::{DefaultBodyLimit, Request, State};
use axum::http::{StatusCode, header};
use axum::middleware::{Next, from_fn_with_state};
use axum::response::{IntoResponse, Response};
use axum::routing::{MethodRouter, get, post};
use kernel::{Address, AxError, B3Hash, Sealed, Seq};
use serde::Deserialize;
use tokio::sync::broadcast;
use crate::answer::Answer;
use crate::assets::ClientAssets;
use crate::auth::Pairing;
use crate::command::{Command, WireCommand};
use crate::frames::Query;
use crate::reception::{Admission, BindFace, Door, decide_admission, offered_pairing};
use crate::reply::{AcpProgress, Reply};
use super::committed::Committed;
mod enrolment;
use super::bundle::bundle_routes;
use super::socket::upgrade;
use super::uploads::{accept_acp, accept_drop, accept_recording};
use enrolment::accept_enrolment;
pub type Answering = Arc<dyn Fn(Query) -> (Seq, Result<Answer, AxError>) + Send + Sync>;
pub type SecretSink =
Arc<dyn Fn(Command<Sealed<String>>, Reply) -> Result<(), AxError> + Send + Sync>;
const ENROLMENT_PATIENCE: std::time::Duration = std::time::Duration::from_secs(2);
pub struct ServeConfig {
pub client: Arc<ClientAssets>,
pub transcribe_sink: TranscribeSink,
pub drop_sink: DropSink,
pub acp: AcpSink,
pub commands: Arc<dyn Fn(WireCommand, Reply) -> Result<(), AxError> + Send + Sync>,
pub events: broadcast::Sender<Committed>,
pub deltas: broadcast::Sender<crate::frames::Delta>,
pub logs: broadcast::Sender<crate::frames::LogLine>,
pub outputs: broadcast::Sender<crate::frames::LiveOutput>,
pub outputs_so_far: Arc<dyn Fn() -> Vec<crate::frames::LiveOutput> + Send + Sync>,
pub monitor: MonitorFeed,
pub queries: Answering,
pub secrets: SecretSink,
pub city: Option<Address>,
pub head: Arc<LedgerHead>,
pub epoch: Option<B3Hash>,
}
#[derive(Debug, Default)]
pub struct LedgerHead(AtomicU64);
impl LedgerHead {
#[must_use]
pub fn at(head_seq: Option<Seq>) -> Self {
let head = Self::default();
if let Some(seq) = head_seq {
head.advance(seq);
}
head
}
pub fn advance(&self, seq: Seq) {
if let Some(stored) = seq.value().checked_add(1) {
self.0.store(stored, Ordering::Release);
}
}
#[must_use]
pub fn read(&self) -> Option<Seq> {
self.0.load(Ordering::Acquire).checked_sub(1).map(Seq::new)
}
}
pub(crate) struct ShellState {
pub(crate) commands: Arc<dyn Fn(WireCommand, Reply) -> Result<(), AxError> + Send + Sync>,
pub(crate) events: broadcast::Sender<Committed>,
pub(crate) deltas: broadcast::Sender<crate::frames::Delta>,
pub(crate) logs: broadcast::Sender<crate::frames::LogLine>,
pub(crate) outputs: broadcast::Sender<crate::frames::LiveOutput>,
pub(crate) outputs_so_far: Arc<dyn Fn() -> Vec<crate::frames::LiveOutput> + Send + Sync>,
pub(crate) monitor: MonitorFeed,
pub(crate) queries: Answering,
pub(crate) secrets: SecretSink,
pub(crate) acp: AcpSink,
pub(crate) transcribe_sink: TranscribeSink,
pub(crate) drop_sink: DropSink,
pub(crate) face: BindFace,
pub(crate) city: Option<Address>,
pub(crate) head: Arc<LedgerHead>,
pub(crate) epoch: Option<B3Hash>,
}
#[derive(Clone)]
pub struct MonitorFeed {
pub watch: Arc<dyn Fn(crate::frames::Watched) -> Box<dyn Send> + Send + Sync>,
pub samples: broadcast::Sender<crate::frames::Sample>,
}
pub type AcpSink =
Arc<dyn Fn(&serde_json::Value, Pairing) -> Result<AcpProgress, AxError> + Send + Sync>;
pub type TranscribeSink = Arc<dyn Fn(Vec<u8>, String) -> Result<String, AxError> + Send + Sync>;
pub type DropSink = Arc<dyn Fn(&str, &[u8]) -> Result<String, AxError> + Send + Sync>;
pub const DROP_BYTES_MAX: usize = 64 * 1024 * 1024;
#[derive(Debug, Deserialize)]
pub struct EnrollBody {
pub realm: String,
pub name: String,
pub value: String,
}
pub fn router(config: &ServeConfig, face: BindFace) -> Router {
let state = Arc::new(ShellState {
commands: Arc::clone(&config.commands),
events: config.events.clone(),
deltas: config.deltas.clone(),
logs: config.logs.clone(),
outputs: config.outputs.clone(),
outputs_so_far: Arc::clone(&config.outputs_so_far),
monitor: config.monitor.clone(),
queries: Arc::clone(&config.queries),
secrets: Arc::clone(&config.secrets),
acp: Arc::clone(&config.acp),
transcribe_sink: Arc::clone(&config.transcribe_sink),
drop_sink: Arc::clone(&config.drop_sink),
face,
city: config.city.clone(),
head: Arc::clone(&config.head),
epoch: config.epoch,
});
Router::new()
.route("/ws", get(upgrade))
.route(
"/transcribe",
paired(Door::Transcribe, &state, post(accept_recording)),
)
.route(
"/enroll",
paired(Door::Enroll, &state, post(accept_enrolment)),
)
.route(
"/drop",
paired(
Door::Drop,
&state,
post(accept_drop).layer(DefaultBodyLimit::max(DROP_BYTES_MAX)),
),
)
.route("/acp", post(accept_acp))
.merge(bundle_routes(Arc::clone(&config.client)))
.with_state(state)
}
pub(crate) fn refusal_text(err: &AxError) -> String {
format!("{}: {}", err.action(), err.recovery())
}
fn paired(
door: Door,
state: &Arc<ShellState>,
handler: MethodRouter<Arc<ShellState>>,
) -> MethodRouter<Arc<ShellState>> {
handler.route_layer(from_fn_with_state((door, Arc::clone(state)), admit_request))
}
async fn admit_request(
State((door, state)): State<(Door, Arc<ShellState>)>,
request: Request,
next: Next,
) -> Response {
let offered = request
.headers()
.get(header::AUTHORIZATION)
.and_then(|value| value.to_str().ok());
match decide_admission(door, offered_pairing(offered), &state.face) {
Admission::Admit(Pairing::Held | Pairing::Absent) => next.run(request).await,
Admission::Refuse(err) => (StatusCode::FORBIDDEN, refusal_text(&err)).into_response(),
}
}