use async_trait::async_trait;
use axum::body::Body;
use axum::extract::{Request, State};
use axum::http::header::{ALLOW, CACHE_CONTROL, CONTENT_TYPE, HOST};
use axum::http::{HeaderMap, HeaderName, HeaderValue, Method, StatusCode, Uri};
use axum::middleware::{self, Next};
use axum::response::Response;
use axum::routing::{MethodRouter, get};
use axum::{Json, Router};
use kcode_k1_http_replay::{ReplayError, ReplayWindow};
use kcode_k1_http_request::{Envelope, VerifyError, registration_public_key};
use serde::Serialize;
use std::fmt::{self, Display, Formatter};
use std::sync::Arc;
use tower_http::cors::{Any, CorsLayer};
pub use kcode_k1_http_request::CanonicalUsername;
pub struct Config {
pub server_id: String,
pub public_origin: String,
pub max_body_bytes: usize,
}
#[derive(Clone, Copy, Eq, PartialEq)]
pub struct Identity {
user_id: [u8; 12],
public_key: [u8; 32],
}
impl Identity {
pub fn new(user_id: [u8; 12], public_key: [u8; 32]) -> Self {
Self {
user_id,
public_key,
}
}
pub fn user_id(&self) -> &[u8; 12] {
&self.user_id
}
pub fn public_key(&self) -> &[u8; 32] {
&self.public_key
}
}
#[derive(Clone)]
pub struct Principal {
user_id: [u8; 12],
username: CanonicalUsername,
}
impl Principal {
pub fn user_id(&self) -> &[u8; 12] {
&self.user_id
}
pub fn username(&self) -> &CanonicalUsername {
&self.username
}
}
#[derive(Clone)]
pub struct RegistrationPrincipal {
username: CanonicalUsername,
public_key: [u8; 32],
}
impl RegistrationPrincipal {
pub fn username(&self) -> &CanonicalUsername {
&self.username
}
pub fn public_key(&self) -> &[u8; 32] {
&self.public_key
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum IdentityError {
Unavailable,
}
impl Display for IdentityError {
fn fmt(&self, formatter: &mut Formatter<'_>) -> fmt::Result {
formatter.write_str("identity provider unavailable")
}
}
impl std::error::Error for IdentityError {}
#[async_trait]
pub trait IdentityProvider: Send + Sync + 'static {
async fn lookup(&self, username: &CanonicalUsername)
-> Result<Option<Identity>, IdentityError>;
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ConfigError {
EmptyServerId,
InvalidPublicOrigin,
InvalidBodyLimit,
}
impl Display for ConfigError {
fn fmt(&self, formatter: &mut Formatter<'_>) -> fmt::Result {
formatter.write_str(match self {
Self::EmptyServerId => "server ID is empty",
Self::InvalidPublicOrigin => "public origin is invalid",
Self::InvalidBodyLimit => "body limit is invalid",
})
}
}
impl std::error::Error for ConfigError {}
pub struct K1Http {
state: Arc<AppState>,
}
struct AppState {
server_id: String,
public_origin: String,
authority: String,
max_body_bytes: usize,
replay: ReplayWindow,
identities: Arc<dyn IdentityProvider>,
}
#[derive(Serialize)]
struct ConfigResponse {
protocol: &'static str,
server_id: String,
public_origin: String,
}
impl K1Http {
pub fn new(
config: Config,
replay: kcode_k1_http_replay::ReplayWindow,
identities: Arc<dyn IdentityProvider>,
) -> Result<Self, ConfigError> {
if config.server_id.is_empty() {
return Err(ConfigError::EmptyServerId);
}
if config.max_body_bytes == 0 {
return Err(ConfigError::InvalidBodyLimit);
}
let authority =
origin_authority(&config.public_origin).ok_or(ConfigError::InvalidPublicOrigin)?;
Ok(Self {
state: Arc::new(AppState {
server_id: config.server_id,
public_origin: config.public_origin,
authority,
max_body_bytes: config.max_body_bytes,
replay,
identities,
}),
})
}
pub fn router(
&self,
registration: MethodRouter,
terms: MethodRouter,
authenticated: Router,
) -> Router {
let authenticated = authenticated.layer(middleware::from_fn_with_state(
self.state.clone(),
authenticate,
));
let registration = registration.layer(middleware::from_fn_with_state(
self.state.clone(),
authenticate_registration,
));
let configuration = get(configuration).with_state(self.state.clone());
Router::new()
.nest("/api", authenticated)
.route("/api/register", registration)
.route("/api/terms", terms)
.route("/api/config.json", configuration)
.layer(middleware::from_fn(terms_method_guard))
.layer(cors())
.layer(middleware::from_fn_with_state(self.state.clone(), api_gate))
}
}
async fn configuration(State(state): State<Arc<AppState>>) -> Json<ConfigResponse> {
Json(ConfigResponse {
protocol: "K1-HTTP-1",
server_id: state.server_id.clone(),
public_origin: state.public_origin.clone(),
})
}
fn origin_authority(value: &str) -> Option<String> {
let uri: Uri = value.parse().ok()?;
let scheme = uri.scheme_str()?;
if !scheme.eq_ignore_ascii_case("http") && !scheme.eq_ignore_ascii_case("https") {
return None;
}
let authority = uri.authority()?.as_str();
if authority.is_empty() || authority.contains('@') || uri.query().is_some() {
return None;
}
if !matches!(uri.path(), "" | "/") {
return None;
}
Some(authority.to_owned())
}
fn cors() -> CorsLayer {
CorsLayer::new()
.allow_origin(Any)
.allow_methods([
Method::GET,
Method::POST,
Method::PUT,
Method::PATCH,
Method::DELETE,
Method::OPTIONS,
Method::HEAD,
])
.allow_headers([
CONTENT_TYPE,
header("k1-username"),
header("k1-epoch"),
header("k1-nonce"),
header("k1-body-sha256"),
header("k1-signature"),
header("k1-public-key"),
])
.expose_headers([header("k1-epoch")])
}
fn header(value: &'static str) -> HeaderName {
HeaderName::from_static(value)
}
fn set(headers: &mut HeaderMap, name: HeaderName, value: &'static str) {
headers.insert(name, HeaderValue::from_static(value));
}
async fn api_gate(State(state): State<Arc<AppState>>, request: Request, next: Next) -> Response {
let authority_ok = single(request.headers(), HOST.as_str())
.is_ok_and(|value| value.as_bytes() == state.authority.as_bytes());
if !authority_ok {
return finish(
&state,
error(StatusCode::MISDIRECTED_REQUEST, "invalid_request_authority"),
)
.await;
}
if state.replay.current_epoch().await.is_err() {
return decorate(
error(StatusCode::SERVICE_UNAVAILABLE, "epoch_unavailable"),
None,
);
}
finish(&state, next.run(request).await).await
}
async fn terms_method_guard(request: Request, next: Next) -> Response {
if request.uri().path() == "/api/terms"
&& request.method() != Method::GET
&& request.method() != Method::HEAD
{
return method_not_allowed("GET, HEAD");
}
next.run(request).await
}
async fn finish(state: &AppState, response: Response) -> Response {
match state.replay.current_epoch().await {
Ok(epoch) => decorate(response, Some(epoch)),
Err(_) => decorate(
error(StatusCode::SERVICE_UNAVAILABLE, "epoch_unavailable"),
None,
),
}
}
fn decorate(mut response: Response, epoch: Option<u64>) -> Response {
let headers = response.headers_mut();
set(headers, CACHE_CONTROL, "no-store");
set(headers, header("x-content-type-options"), "nosniff");
set(headers, header("access-control-allow-origin"), "*");
set(headers, header("access-control-expose-headers"), "k1-epoch");
if let Some(epoch) = epoch
&& let Ok(value) = HeaderValue::from_str(&epoch.to_string())
{
headers.insert(header("k1-epoch"), value);
}
response
}
fn single<'a>(headers: &'a HeaderMap, name: &str) -> Result<&'a HeaderValue, ()> {
let mut values = headers.get_all(name).iter();
let value = values.next().ok_or(())?;
if values.next().is_some() {
return Err(());
}
Ok(value)
}
async fn authenticate(
State(state): State<Arc<AppState>>,
request: Request,
next: Next,
) -> Response {
proceed(authenticated_request(&state, request).await, next).await
}
async fn authenticate_registration(
State(state): State<Arc<AppState>>,
request: Request,
next: Next,
) -> Response {
if request.method() != Method::POST {
return method_not_allowed("POST");
}
proceed(registration_request(&state, request).await, next).await
}
async fn proceed(result: Result<Request, Response>, next: Next) -> Response {
match result {
Ok(request) => next.run(request).await,
Err(response) => response,
}
}
async fn authenticated_request(state: &AppState, request: Request) -> Result<Request, Response> {
let envelope = Envelope::parse(request.headers())
.map_err(|_| error(StatusCode::BAD_REQUEST, "malformed_authentication_envelope"))?;
validate_epoch(state, envelope.epoch()).await?;
let identity = match state.identities.lookup(envelope.username()).await {
Ok(Some(identity)) => identity,
Ok(None) => return Err(error(StatusCode::UNAUTHORIZED, "authentication_failed")),
Err(IdentityError::Unavailable) => {
return Err(error(
StatusCode::SERVICE_UNAVAILABLE,
"identity_provider_unavailable",
));
}
};
let mut request = verified(state, &envelope, request, identity.public_key()).await?;
let mut replay_id = [0_u8; 32];
replay_id[..12].copy_from_slice(identity.user_id());
state
.replay
.admit(&replay_id, envelope.epoch(), *envelope.nonce())
.await
.map_err(replay_error)?;
request.extensions_mut().insert(Principal {
user_id: *identity.user_id(),
username: envelope.username().clone(),
});
Ok(request)
}
async fn registration_request(state: &AppState, request: Request) -> Result<Request, Response> {
let envelope = Envelope::parse(request.headers())
.map_err(|_| error(StatusCode::BAD_REQUEST, "malformed_authentication_envelope"))?;
if single(request.headers(), "k1-public-key").is_err() {
return Err(error(
StatusCode::BAD_REQUEST,
"malformed_authentication_envelope",
));
}
let candidate = registration_public_key(request.headers())
.map_err(|_| error(StatusCode::UNAUTHORIZED, "authentication_failed"))?;
validate_epoch(state, envelope.epoch()).await?;
let mut request = verified(state, &envelope, request, &candidate).await?;
request.extensions_mut().insert(RegistrationPrincipal {
username: envelope.username().clone(),
public_key: candidate,
});
Ok(request)
}
async fn verified(
state: &AppState,
envelope: &Envelope,
request: Request,
public_key: &[u8; 32],
) -> Result<Request, Response> {
envelope
.verify(
request,
&state.server_id,
&state.public_origin,
public_key,
state.max_body_bytes,
)
.await
.map_err(verify_error)
}
async fn validate_epoch(state: &AppState, epoch: u64) -> Result<(), Response> {
state
.replay
.validate_epoch(epoch)
.await
.map(|_| ())
.map_err(replay_error)
}
fn replay_error(cause: ReplayError) -> Response {
match cause {
ReplayError::EpochOutsideWindow { .. } => error(StatusCode::UNAUTHORIZED, "stale_epoch"),
ReplayError::Replay { .. } => error(StatusCode::CONFLICT, "replay"),
ReplayError::CapacityExceeded { .. } => {
error(StatusCode::TOO_MANY_REQUESTS, "nonce_capacity")
}
_ => error(StatusCode::SERVICE_UNAVAILABLE, "epoch_unavailable"),
}
}
fn verify_error(cause: VerifyError) -> Response {
match cause {
VerifyError::AuthenticationFailed => {
error(StatusCode::UNAUTHORIZED, "authentication_failed")
}
VerifyError::BodyDigestMismatch => error(StatusCode::BAD_REQUEST, "body_digest_mismatch"),
VerifyError::BodyTooLarge => error(StatusCode::PAYLOAD_TOO_LARGE, "body_too_large"),
}
}
fn method_not_allowed(allow: &'static str) -> Response {
let mut response = Response::new(Body::empty());
*response.status_mut() = StatusCode::METHOD_NOT_ALLOWED;
set(response.headers_mut(), ALLOW, allow);
response
}
fn error(status: StatusCode, code: &'static str) -> Response {
let mut response = Response::new(Body::from(format!("{{\"error\":\"{code}\"}}")));
*response.status_mut() = status;
set(response.headers_mut(), CONTENT_TYPE, "application/json");
response
}