use async_trait::async_trait;
use axum::Router;
use axum::body::{Body, to_bytes};
use axum::extract::{OriginalUri, 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;
use base64::Engine as _;
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
use kcode_k1_http_replay::{ReplayError, ReplayWindow};
use kcode_k1_http_signature::{RequestBinding, verify};
use sha2::{Digest, Sha256};
use std::fmt::{self, Display, Formatter};
use std::sync::Arc;
use tower_http::cors::{Any, CorsLayer};
pub use kcode_k1_http_signature::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>,
}
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, 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,
));
Router::new()
.nest("/api", authenticated)
.route("/api/register", registration)
.layer(cors())
.layer(middleware::from_fn_with_state(self.state.clone(), api_gate))
}
}
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)
}
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,
);
}
let response = next.run(request).await;
finish(&state, response).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();
headers.insert(CACHE_CONTROL, HeaderValue::from_static("no-store"));
headers.insert(
HeaderName::from_static("x-content-type-options"),
HeaderValue::from_static("nosniff"),
);
headers.insert(
HeaderName::from_static("access-control-allow-origin"),
HeaderValue::from_static("*"),
);
headers.insert(
HeaderName::from_static("access-control-expose-headers"),
HeaderValue::from_static("k1-epoch"),
);
if let Some(epoch) = epoch
&& let Ok(value) = HeaderValue::from_str(&epoch.to_string())
{
headers.insert(HeaderName::from_static("k1-epoch"), value);
}
response
}
struct Envelope {
username: CanonicalUsername,
epoch: u64,
nonce: [u8; 16],
body_sha256: [u8; 32],
signature: [u8; 64],
content_type: String,
}
impl Envelope {
fn parse(headers: &HeaderMap) -> Result<Self, ()> {
let username = CanonicalUsername::parse(text(headers, "k1-username")?).map_err(|_| ())?;
let epoch = text(headers, "k1-epoch")?.parse().map_err(|_| ())?;
let nonce = fixed(text(headers, "k1-nonce")?)?;
let body_sha256 = fixed(text(headers, "k1-body-sha256")?)?;
let signature = fixed(text(headers, "k1-signature")?)?;
let content_type = optional_ascii(headers, CONTENT_TYPE.as_str())?;
Ok(Self {
username,
epoch,
nonce,
body_sha256,
signature,
content_type,
})
}
}
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)
}
fn text<'a>(headers: &'a HeaderMap, name: &str) -> Result<&'a str, ()> {
single(headers, name)?.to_str().map_err(|_| ())
}
fn optional_ascii(headers: &HeaderMap, name: &str) -> Result<String, ()> {
let mut values = headers.get_all(name).iter();
let Some(value) = values.next() else {
return Ok(String::new());
};
if values.next().is_some() || !value.as_bytes().is_ascii() {
return Err(());
}
Ok(std::str::from_utf8(value.as_bytes())
.map_err(|_| ())?
.to_owned())
}
fn fixed<const N: usize>(value: &str) -> Result<[u8; N], ()> {
let decoded = URL_SAFE_NO_PAD.decode(value).map_err(|_| ())?;
if decoded.len() != N || URL_SAFE_NO_PAD.encode(&decoded) != value {
return Err(());
}
decoded.try_into().map_err(|_| ())
}
fn request_target(request: &Request) -> String {
let uri = request
.extensions()
.get::<OriginalUri>()
.map(|original| &original.0)
.unwrap_or_else(|| request.uri());
uri.path_and_query()
.map(|value| value.as_str())
.unwrap_or("/")
.to_owned()
}
async fn authenticate(
State(state): State<Arc<AppState>>,
mut request: Request,
next: Next,
) -> Response {
let envelope = match Envelope::parse(request.headers()) {
Ok(value) => value,
Err(()) => return error(StatusCode::BAD_REQUEST, "malformed_authentication_envelope"),
};
if let Err(response) = validate_epoch(&state, envelope.epoch).await {
return response;
}
let identity = match state.identities.lookup(&envelope.username).await {
Ok(Some(identity)) => identity,
Ok(None) => return error(StatusCode::UNAUTHORIZED, "authentication_failed"),
Err(IdentityError::Unavailable) => {
return error(
StatusCode::SERVICE_UNAVAILABLE,
"identity_provider_unavailable",
);
}
};
let method = request.method().as_str().to_owned();
let target = request_target(&request);
if !valid_signature(&state, &envelope, identity.public_key(), &method, &target) {
return error(StatusCode::UNAUTHORIZED, "authentication_failed");
}
if let Err(response) =
check_body(&mut request, envelope.body_sha256, state.max_body_bytes).await
{
return response;
}
let mut replay_id = [0_u8; 32];
replay_id[..12].copy_from_slice(identity.user_id());
match state
.replay
.admit(&replay_id, envelope.epoch, envelope.nonce)
.await
{
Ok(_) => {}
Err(ReplayError::EpochOutsideWindow { .. }) => {
return error(StatusCode::UNAUTHORIZED, "stale_epoch");
}
Err(ReplayError::Replay { .. }) => return error(StatusCode::CONFLICT, "replay"),
Err(ReplayError::CapacityExceeded { .. }) => {
return error(StatusCode::TOO_MANY_REQUESTS, "nonce_capacity");
}
Err(_) => return error(StatusCode::SERVICE_UNAVAILABLE, "epoch_unavailable"),
}
request.extensions_mut().insert(Principal {
user_id: *identity.user_id(),
username: envelope.username,
});
next.run(request).await
}
async fn authenticate_registration(
State(state): State<Arc<AppState>>,
mut request: Request,
next: Next,
) -> Response {
if request.method() != Method::POST {
let mut response = Response::new(Body::empty());
*response.status_mut() = StatusCode::METHOD_NOT_ALLOWED;
response
.headers_mut()
.insert(ALLOW, HeaderValue::from_static("POST"));
return response;
}
let envelope = match Envelope::parse(request.headers()) {
Ok(value) => value,
Err(()) => return error(StatusCode::BAD_REQUEST, "malformed_authentication_envelope"),
};
let candidate = match single(request.headers(), "k1-public-key") {
Ok(value) => value,
Err(()) => return error(StatusCode::BAD_REQUEST, "malformed_authentication_envelope"),
};
let candidate = match std::str::from_utf8(candidate.as_bytes())
.ok()
.and_then(|value| fixed(value).ok())
{
Some(value) => value,
None => return error(StatusCode::UNAUTHORIZED, "authentication_failed"),
};
if let Err(response) = validate_epoch(&state, envelope.epoch).await {
return response;
}
let method = request.method().as_str().to_owned();
let target = request_target(&request);
if !valid_signature(&state, &envelope, &candidate, &method, &target) {
return error(StatusCode::UNAUTHORIZED, "authentication_failed");
}
if let Err(response) =
check_body(&mut request, envelope.body_sha256, state.max_body_bytes).await
{
return response;
}
request.extensions_mut().insert(RegistrationPrincipal {
username: envelope.username,
public_key: candidate,
});
next.run(request).await
}
async fn validate_epoch(state: &AppState, epoch: u64) -> Result<(), Response> {
match state.replay.validate_epoch(epoch).await {
Ok(_) => Ok(()),
Err(ReplayError::EpochOutsideWindow { .. }) => {
Err(error(StatusCode::UNAUTHORIZED, "stale_epoch"))
}
Err(_) => Err(error(StatusCode::SERVICE_UNAVAILABLE, "epoch_unavailable")),
}
}
fn valid_signature(
state: &AppState,
envelope: &Envelope,
public_key: &[u8; 32],
method: &str,
target: &str,
) -> bool {
verify(
&RequestBinding {
server_id: &state.server_id,
public_origin: &state.public_origin,
username: &envelope.username,
epoch: envelope.epoch,
nonce: envelope.nonce,
method,
target,
content_type: &envelope.content_type,
body_sha256: envelope.body_sha256,
},
public_key,
&envelope.signature,
)
.is_ok()
}
async fn check_body(
request: &mut Request,
expected: [u8; 32],
limit: usize,
) -> Result<(), Response> {
let body = std::mem::replace(request.body_mut(), Body::empty());
let bytes = to_bytes(body, limit)
.await
.map_err(|_| error(StatusCode::PAYLOAD_TOO_LARGE, "body_too_large"))?;
let actual: [u8; 32] = Sha256::digest(&bytes).into();
if actual != expected {
return Err(error(StatusCode::BAD_REQUEST, "body_digest_mismatch"));
}
*request.body_mut() = Body::from(bytes);
Ok(())
}
fn error(status: StatusCode, code: &'static str) -> Response {
let mut response = Response::new(Body::from(format!("{{\"error\":\"{code}\"}}")));
*response.status_mut() = status;
response
.headers_mut()
.insert(CONTENT_TYPE, HeaderValue::from_static("application/json"));
response
}