use amberfork_layout::{Document, PayloadResponse, SlotAddress, ViewModel};
use axum::Router;
use axum::body::Bytes;
use axum::extract::{Json, Request, State};
use axum::http::{HeaderMap, HeaderValue, StatusCode, Uri, header};
use axum::middleware::{self, Next};
use axum::response::{IntoResponse, Response};
use axum::routing::{get, post};
use rust_embed::{EmbeddedFile, RustEmbed};
use sha2::{Digest, Sha256};
use std::borrow::Cow;
use std::fmt;
use std::io;
use std::net::{Ipv4Addr, SocketAddr};
use std::sync::Arc;
use tokio::net::TcpListener;
pub const DOCUMENT_ROUTE: &str = "/api/document";
pub const PAYLOAD_ROUTE: &str = "/api/payload";
const INDEX_HTML: &str = "index.html";
#[derive(RustEmbed)]
#[folder = "ui-dist"]
struct UiAssets;
struct Content {
body: Bytes,
etag: HeaderValue,
}
struct ServerState {
content: Content,
full_view: ViewModel,
}
#[derive(Debug)]
pub struct Server {
listener: TcpListener,
router: Router,
local_addr: SocketAddr,
}
#[derive(Debug)]
pub enum ServeError {
Bind { addr: SocketAddr, source: io::Error },
Serve(io::Error),
BundleMissing,
}
impl fmt::Display for ServeError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Bind { addr, source } => write!(f, "cannot bind {addr}: {source}"),
Self::Serve(source) => write!(f, "server stopped: {source}"),
Self::BundleMissing => write!(
f,
"web UI bundle missing — this build has no index.html embedded; build the \
web UI into crates/amberfork-server/ui-dist/ first (release artifacts ship \
with it built in)"
),
}
}
}
impl std::error::Error for ServeError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
Self::Bind { source, .. } | Self::Serve(source) => Some(source),
Self::BundleMissing => None,
}
}
}
impl Server {
pub async fn bind(
document: &Document,
full_view: &ViewModel,
port: u16,
) -> Result<Self, ServeError> {
Self::bind_with_assets::<UiAssets>(document, full_view, port).await
}
pub async fn bind_with_assets<A: RustEmbed + 'static>(
document: &Document,
full_view: &ViewModel,
port: u16,
) -> Result<Self, ServeError> {
if A::get(INDEX_HTML).is_none() {
return Err(ServeError::BundleMissing);
}
let addr = SocketAddr::from((Ipv4Addr::LOCALHOST, port));
let listener = TcpListener::bind(addr)
.await
.map_err(|source| ServeError::Bind { addr, source })?;
let local_addr = listener
.local_addr()
.map_err(|source| ServeError::Bind { addr, source })?;
Ok(Self {
listener,
router: router::<A>(document, full_view),
local_addr,
})
}
pub fn local_addr(&self) -> SocketAddr {
self.local_addr
}
pub async fn serve(self) -> Result<(), ServeError> {
axum::serve(self.listener, self.router)
.await
.map_err(ServeError::Serve)
}
}
fn router<A: RustEmbed + 'static>(document: &Document, full_view: &ViewModel) -> Router {
let json = serde_json::to_string(document)
.expect("Document serialization is infallible (no non-string map keys)");
let etag = HeaderValue::from_str(&format!("\"{:x}\"", Sha256::digest(json.as_bytes())))
.expect("a quoted hex digest is valid header ASCII");
let state = Arc::new(ServerState {
content: Content {
body: Bytes::from(json),
etag,
},
full_view: full_view.clone(),
});
Router::new()
.route(DOCUMENT_ROUTE, get(serve_document))
.route(PAYLOAD_ROUTE, post(serve_payload))
.fallback(get(serve_ui::<A>))
.layer(middleware::from_fn(require_local_host))
.with_state(state)
}
async fn serve_ui<A: RustEmbed + 'static>(uri: Uri) -> Response {
let path = uri.path();
if path.starts_with("/api/") {
return (
StatusCode::NOT_FOUND,
"amberfork: no such endpoint (the content endpoint is /api/document)\n",
)
.into_response();
}
let asset_path = path.trim_start_matches('/');
match A::get(asset_path) {
Some(asset) => asset_response(asset_path, asset),
None => match A::get(INDEX_HTML) {
Some(index) => asset_response(INDEX_HTML, index),
None => (
StatusCode::INTERNAL_SERVER_ERROR,
"amberfork: UI bundle lost index.html\n",
)
.into_response(),
},
}
}
fn asset_response(path: &str, asset: EmbeddedFile) -> Response {
let body = match asset.data {
Cow::Borrowed(bytes) => Bytes::from_static(bytes),
Cow::Owned(bytes) => Bytes::from(bytes),
};
let mime = mime_guess::from_path(path).first_or_octet_stream();
let mime = HeaderValue::from_str(mime.as_ref())
.unwrap_or_else(|_| HeaderValue::from_static("application/octet-stream"));
([(header::CONTENT_TYPE, mime)], body).into_response()
}
async fn serve_document(State(state): State<Arc<ServerState>>, headers: HeaderMap) -> Response {
let content = &state.content;
let revalidate = [
(header::ETAG, content.etag.clone()),
(header::CACHE_CONTROL, HeaderValue::from_static("no-cache")),
];
if headers.get(header::IF_NONE_MATCH) == Some(&content.etag) {
return (StatusCode::NOT_MODIFIED, revalidate).into_response();
}
(
revalidate,
[(
header::CONTENT_TYPE,
HeaderValue::from_static("application/json"),
)],
content.body.clone(),
)
.into_response()
}
async fn serve_payload(
State(state): State<Arc<ServerState>>,
Json(address): Json<SlotAddress>,
) -> Response {
match state.full_view.full_text(&address) {
Some(text) => Json(PayloadResponse {
text: text.to_string(),
})
.into_response(),
None => (
StatusCode::NOT_FOUND,
"amberfork: no payload at that address\n",
)
.into_response(),
}
}
async fn require_local_host(request: Request, next: Next) -> Response {
let allowed = request
.headers()
.get(header::HOST)
.and_then(|value| value.to_str().ok())
.is_some_and(host_is_local);
if !allowed {
return (
StatusCode::FORBIDDEN,
"amberfork: forbidden — non-local Host header\n",
)
.into_response();
}
next.run(request).await
}
fn host_is_local(host: &str) -> bool {
let name = if host.starts_with('[') {
let Some(end) = host.find(']') else {
return false;
};
match &host[end + 1..] {
"" => &host[..=end],
port if port.starts_with(':') => &host[..=end],
_ => return false,
}
} else {
host.split(':').next().unwrap_or(host)
};
let name = name.to_ascii_lowercase();
name == "localhost" || name == "127.0.0.1" || name == "[::1]"
}
#[cfg(test)]
mod tests {
use super::host_is_local;
#[test]
fn localhost_forms_pass() {
for host in [
"localhost",
"LOCALHOST",
"localhost:7777",
"127.0.0.1",
"127.0.0.1:7777",
"[::1]",
"[::1]:7777",
] {
assert!(host_is_local(host), "{host:?} names loopback");
}
}
#[test]
fn everything_else_is_refused() {
for host in [
"",
"evil.example",
"evil.example:7777",
"localhost.evil.example",
"127.0.0.1.evil.example",
"[::2]",
"[::1",
"[::1]evil",
"0.0.0.0",
"192.168.1.10:7777",
] {
assert!(!host_is_local(host), "{host:?} must not pass the guard");
}
}
}