From 12051336827bfb3befa1dc4ab18a82f37b3941a0 Mon Sep 17 00:00:00 2001 From: Matt Gibson Date: Thu, 15 Jan 2026 10:18:01 -0800 Subject: [PATCH] Reader API implementation --- akd/Cargo.lock | 15 ++++ akd/crates/akd_storage/src/vrf_key_config.rs | 2 +- akd/crates/common/src/lib.rs | 4 +- akd/crates/reader/Cargo.toml | 3 +- akd/crates/reader/src/lib.rs | 12 +-- akd/crates/reader/src/routes/audit.rs | 36 +++++++++ akd/crates/reader/src/routes/batch_lookup.rs | 47 +++++++++++ .../reader/src/routes/get_epoch_hash.rs | 40 ++++++++++ .../reader/src/routes/get_public_key.rs | 26 ++++++ akd/crates/reader/src/routes/health.rs | 10 ++- akd/crates/reader/src/routes/key_history.rs | 70 ++++++++++++++++ akd/crates/reader/src/routes/lookup.rs | 52 ++++++++++++ akd/crates/reader/src/routes/mod.rs | 79 ++++++++++++++++++- .../reader/src/routes/response_types.rs | 9 +++ 14 files changed, 392 insertions(+), 13 deletions(-) create mode 100644 akd/crates/reader/src/routes/audit.rs create mode 100644 akd/crates/reader/src/routes/batch_lookup.rs create mode 100644 akd/crates/reader/src/routes/get_epoch_hash.rs create mode 100644 akd/crates/reader/src/routes/get_public_key.rs create mode 100644 akd/crates/reader/src/routes/key_history.rs create mode 100644 akd/crates/reader/src/routes/lookup.rs create mode 100644 akd/crates/reader/src/routes/response_types.rs diff --git a/akd/Cargo.lock b/akd/Cargo.lock index 8c2ebd99b1..a8a6d76752 100644 --- a/akd/Cargo.lock +++ b/akd/Cargo.lock @@ -62,6 +62,7 @@ dependencies = [ "hex", "log", "protobuf", + "serde", "tokio", ] @@ -79,6 +80,8 @@ dependencies = [ "protobuf", "protobuf-codegen", "protobuf-parse", + "serde", + "serde_bytes", "tokio", "zeroize", ] @@ -953,6 +956,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "115531babc129696a58c64a4fef0a8bf9e9698629fb97e9e40767d235cfbcd53" dependencies = [ "pkcs8", + "serde", "signature", ] @@ -2255,6 +2259,7 @@ dependencies = [ "anyhow", "axum", "bitwarden-akd-configuration", + "bitwarden-encoding", "chrono", "common", "config", @@ -2531,6 +2536,16 @@ dependencies = [ "typeid", ] +[[package]] +name = "serde_bytes" +version = "0.11.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a5d440709e79d88e51ac01c4b72fc6cb7314017bb7da9eeff678aa94c10e3ea8" +dependencies = [ + "serde", + "serde_core", +] + [[package]] name = "serde_core" version = "1.0.228" diff --git a/akd/crates/akd_storage/src/vrf_key_config.rs b/akd/crates/akd_storage/src/vrf_key_config.rs index e7a34dd0bd..03bf3e02a7 100644 --- a/akd/crates/akd_storage/src/vrf_key_config.rs +++ b/akd/crates/akd_storage/src/vrf_key_config.rs @@ -91,7 +91,7 @@ impl VrfKeyConfig { Ok(blake3::hash(&root_key_bytes).as_bytes().to_vec()) } - pub fn root_key_type(&self) -> VrfRootKeyType { + pub(crate) fn root_key_type(&self) -> VrfRootKeyType { match self { #[cfg(test)] VrfKeyConfig::ConstantVrfKey => VrfRootKeyType::None, diff --git a/akd/crates/common/src/lib.rs b/akd/crates/common/src/lib.rs index 7c3089bb70..f5ef419aae 100644 --- a/akd/crates/common/src/lib.rs +++ b/akd/crates/common/src/lib.rs @@ -1,5 +1,7 @@ -use akd::Directory; +use akd::{directory::ReadOnlyDirectory, Directory}; use akd_storage::{AkdDatabase, VrfKeyDatabase}; use bitwarden_akd_configuration::BitwardenV1Configuration; pub type BitAkdDirectory = Directory; +pub type ReadOnlyBitAkdDirectory = + ReadOnlyDirectory; diff --git a/akd/crates/reader/Cargo.toml b/akd/crates/reader/Cargo.toml index 55708ff662..4e846d7a72 100644 --- a/akd/crates/reader/Cargo.toml +++ b/akd/crates/reader/Cargo.toml @@ -10,7 +10,7 @@ keywords.workspace = true tokio = { workspace = true } tracing = { workspace = true } tracing-subscriber = { workspace = true, features = ["env-filter"] } -akd = { workspace = true } +akd = { workspace = true, features = ["serde_serialization"] } akd_storage = { workspace = true } anyhow = { workspace = true } bitwarden-akd-configuration = { workspace = true } @@ -20,6 +20,7 @@ config = { workspace = true } uuid = { workspace = true } axum = { workspace = true } chrono = { workspace = true } +bitwarden-encoding = { workspace = true } [lints] workspace = true diff --git a/akd/crates/reader/src/lib.rs b/akd/crates/reader/src/lib.rs index f079c29329..0555996675 100644 --- a/akd/crates/reader/src/lib.rs +++ b/akd/crates/reader/src/lib.rs @@ -1,5 +1,5 @@ use akd::directory::ReadOnlyDirectory; -use akd_storage::{AkdDatabase, ReadOnlyPublishQueueType, VrfKeyDatabase}; +use akd_storage::{AkdDatabase, VrfKeyDatabase}; use anyhow::{Context, Result}; use axum::Router; use bitwarden_akd_configuration::BitwardenV1Configuration; @@ -10,12 +10,14 @@ mod config; mod routes; pub use crate::config::ApplicationConfig; +pub use routes::response_types; #[derive(Clone)] struct AppState { // Add any shared state here, e.g., database connections directory: ReadOnlyDirectory, - publish_queue: ReadOnlyPublishQueueType, + // TODO: use this to allow for unique failures for lookup and key history requests that have pending updates + // publish_queue: ReadOnlyPublishQueueType, } #[instrument(skip_all, name = "reader_start")] @@ -23,7 +25,7 @@ pub async fn start( config: ApplicationConfig, shutdown_rx: &Receiver<()>, ) -> Result>> { - let (directory, _, publish_queue) = config + let (directory, _, _) = config .storage .initialize_readonly_directory::() .await @@ -34,14 +36,14 @@ pub async fn start( let handle = tokio::spawn(async move { let app_state = AppState { directory: directory, - publish_queue: publish_queue, + // publish_queue: publish_queue, }; let app = Router::new() .merge(crate::routes::api_routes()) .with_state(app_state); - let listener = TcpListener::bind((&config.socket_address())) + let listener = TcpListener::bind(&config.socket_address()) .await .context("Socket bind failed")?; info!( diff --git a/akd/crates/reader/src/routes/audit.rs b/akd/crates/reader/src/routes/audit.rs new file mode 100644 index 0000000000..82f14dc4dd --- /dev/null +++ b/akd/crates/reader/src/routes/audit.rs @@ -0,0 +1,36 @@ +use axum::{extract::State, http::StatusCode, Json}; +use serde::{Deserialize, Serialize}; +use tracing::{error, info, instrument}; + +use crate::{routes::Response, AppState}; + +#[derive(Debug, Serialize, Deserialize)] +pub struct AuditRequest { + pub start_epoch: u64, + pub end_epoch: u64, +} + +pub type AuditData = akd::AppendOnlyProof; + +#[instrument(skip_all)] +pub async fn audit_handler( + State(AppState { directory, .. }): State, + Json(AuditRequest { + start_epoch, + end_epoch, + }): Json, +) -> (StatusCode, Json>) { + info!( + start_epoch, + end_epoch, "Handling epoch audit request request" + ); + let audit_proof = directory.audit(start_epoch, end_epoch).await; + + match audit_proof { + Ok(audit_proof) => (StatusCode::OK, Json(Response::success(audit_proof))), + Err(e) => { + error!(err = ?e, "Failed to perform epoch audit"); + (StatusCode::INTERNAL_SERVER_ERROR, Json(Response::fail(e))) + } + } +} diff --git a/akd/crates/reader/src/routes/batch_lookup.rs b/akd/crates/reader/src/routes/batch_lookup.rs new file mode 100644 index 0000000000..3b54f56412 --- /dev/null +++ b/akd/crates/reader/src/routes/batch_lookup.rs @@ -0,0 +1,47 @@ +use axum::{extract::State, http::StatusCode, Json}; +use serde::{Deserialize, Serialize}; +use tracing::{error, info, instrument}; + +use crate::{ + routes::{get_epoch_hash::EpochData, lookup::AkdLabelB64, Response}, + AppState, +}; + +#[derive(Debug, Serialize, Deserialize)] +pub struct BatchLookupRequest { + /// An array of labels to look up. Each label is encoded as base64. + pub labels_b64: Vec, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct BatchLookupData { + pub lookup_proofs: Vec, + pub epoch_data: EpochData, +} + +#[instrument(skip_all)] +pub async fn batch_lookup_handler( + State(AppState { directory, .. }): State, + Json(BatchLookupRequest { labels_b64 }): Json, +) -> (StatusCode, Json>) { + info!("Handling get public key request"); + let labels = labels_b64 + .into_iter() + .map(|label_b64| label_b64.into()) + .collect::>(); + let lookup_proofs = directory.batch_lookup(&labels).await; + + match lookup_proofs { + Ok((lookup_proofs, epoch_hash)) => ( + StatusCode::OK, + Json(Response::success(BatchLookupData { + lookup_proofs, + epoch_data: epoch_hash.into(), + })), + ), + Err(e) => { + error!(err = ?e, "Failed to get AKD public key"); + (StatusCode::INTERNAL_SERVER_ERROR, Json(Response::fail(e))) + } + } +} diff --git a/akd/crates/reader/src/routes/get_epoch_hash.rs b/akd/crates/reader/src/routes/get_epoch_hash.rs new file mode 100644 index 0000000000..97bfac060a --- /dev/null +++ b/akd/crates/reader/src/routes/get_epoch_hash.rs @@ -0,0 +1,40 @@ +use akd::EpochHash; +use axum::{extract::State, http::StatusCode, Json}; +use serde::{Deserialize, Serialize}; +use tracing::{error, info, instrument}; + +use crate::{routes::Response, AppState}; + +#[derive(Debug, Serialize, Deserialize)] +pub struct EpochData { + pub epoch: u64, + pub epoch_hash_b64: bitwarden_encoding::B64, +} + +impl From for EpochData { + fn from(epoch_hash: EpochHash) -> Self { + EpochData { + epoch: epoch_hash.0, + epoch_hash_b64: bitwarden_encoding::B64::from(epoch_hash.1.as_ref()), + } + } +} + +#[instrument(skip_all)] +pub async fn get_epoch_hash_handler( + State(AppState { directory, .. }): State, +) -> (StatusCode, Json>) { + info!("Handling get epoch hash request"); + let epoch_hash = directory.get_epoch_hash().await; + + match epoch_hash { + Ok(epoch_hash) => ( + StatusCode::OK, + Json(Response::::success(epoch_hash.into())), + ), + Err(e) => { + error!(err = ?e, "Failed to get current AKD epoch hash"); + (StatusCode::INTERNAL_SERVER_ERROR, Json(Response::fail(e))) + } + } +} diff --git a/akd/crates/reader/src/routes/get_public_key.rs b/akd/crates/reader/src/routes/get_public_key.rs new file mode 100644 index 0000000000..1b83c0f085 --- /dev/null +++ b/akd/crates/reader/src/routes/get_public_key.rs @@ -0,0 +1,26 @@ +use axum::{extract::State, http::StatusCode, Json}; +use tracing::{error, info, instrument}; + +use crate::{routes::Response, AppState}; + +/// Public key encoded as a base64 string +pub type PublicKeyData = bitwarden_encoding::B64; + +#[instrument(skip_all)] +pub async fn get_public_key_handler( + State(AppState { directory, .. }): State, +) -> (StatusCode, Json>) { + info!("Handling get public key request"); + let public_key = directory.get_public_key().await; + + match public_key { + Ok(public_key) => ( + StatusCode::OK, + Json(Response::success(public_key.as_ref().into())), + ), + Err(e) => { + error!(err = ?e, "Failed to get AKD public key"); + (StatusCode::INTERNAL_SERVER_ERROR, Json(Response::fail(e))) + } + } +} diff --git a/akd/crates/reader/src/routes/health.rs b/akd/crates/reader/src/routes/health.rs index a299442bc7..b6c9686e63 100644 --- a/akd/crates/reader/src/routes/health.rs +++ b/akd/crates/reader/src/routes/health.rs @@ -1,17 +1,19 @@ -use axum::Json; +use axum::{http::StatusCode, Json}; use serde::{Deserialize, Serialize}; use tracing::{info, instrument}; +use crate::routes::Response; + #[derive(Debug, Serialize, Deserialize)] -pub struct ServerHealth { +pub struct HealthData { time: String, } #[instrument(skip_all)] -pub async fn health_handler() -> Json { +pub async fn health_handler() -> (StatusCode, Json>) { info!("Handling server health request"); let time = chrono::Utc::now().to_rfc3339(); - Json(ServerHealth { time }) + (StatusCode::OK, Json(Response::success(HealthData { time }))) } diff --git a/akd/crates/reader/src/routes/key_history.rs b/akd/crates/reader/src/routes/key_history.rs new file mode 100644 index 0000000000..e9ef082579 --- /dev/null +++ b/akd/crates/reader/src/routes/key_history.rs @@ -0,0 +1,70 @@ +use axum::{extract::State, http::StatusCode, Json}; +use serde::{Deserialize, Serialize}; +use tracing::{error, info, instrument}; + +use crate::{ + routes::{get_epoch_hash::EpochData, Response}, + AppState, +}; + +#[derive(Debug, Serialize, Deserialize)] +pub struct KeyHistoryRequest { + /// the label to look up encoded as an uppercase hex string + pub label: akd::AkdLabel, + pub history_params: HistoryParams, +} + +/// The parameters that dictate how much of the history proof to return to the consumer +/// (either a complete history, or some limited form). +#[derive(Copy, Clone, Serialize, Deserialize, Debug)] +#[serde(tag = "type")] +pub enum HistoryParams { + /// Returns a complete history for a label + Complete, + /// Returns up to the most recent N updates for a label + MostRecent(usize), + /// Returns all updates since a specified epoch (inclusive) + SinceEpoch(u64), +} + +impl From for akd::HistoryParams { + fn from(params: HistoryParams) -> Self { + match params { + HistoryParams::Complete => akd::HistoryParams::Complete, + HistoryParams::MostRecent(n) => akd::HistoryParams::MostRecent(n), + HistoryParams::SinceEpoch(epoch) => akd::HistoryParams::SinceEpoch(epoch), + } + } +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct HistoryData { + pub history_proof: akd::HistoryProof, + pub epoch_data: EpochData, +} + +#[instrument(skip_all)] +pub async fn key_history_handler( + State(AppState { directory, .. }): State, + Json(KeyHistoryRequest { + label, + history_params, + }): Json, +) -> (StatusCode, Json>) { + info!("Handling get public key request"); + let history_proof = directory.key_history(&label, history_params.into()).await; + + match history_proof { + Ok((history_proof, epoch_hash)) => ( + StatusCode::OK, + Json(Response::success(HistoryData { + history_proof: history_proof, + epoch_data: epoch_hash.into(), + })), + ), + Err(e) => { + error!(err = ?e, "Failed to get AKD public key"); + (StatusCode::INTERNAL_SERVER_ERROR, Json(Response::fail(e))) + } + } +} diff --git a/akd/crates/reader/src/routes/lookup.rs b/akd/crates/reader/src/routes/lookup.rs new file mode 100644 index 0000000000..e8cc3d11f0 --- /dev/null +++ b/akd/crates/reader/src/routes/lookup.rs @@ -0,0 +1,52 @@ +use axum::{extract::State, http::StatusCode, Json}; +use serde::{Deserialize, Serialize}; +use tracing::{error, info, instrument}; + +use crate::{ + routes::{get_epoch_hash::EpochData, Response}, + AppState, +}; + +#[derive(Debug, Serialize, Deserialize)] +pub struct AkdLabelB64(bitwarden_encoding::B64); + +impl From for akd::AkdLabel { + fn from(label_b64: AkdLabelB64) -> Self { + akd::AkdLabel(label_b64.0.into_bytes()) + } +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct LookupRequest { + /// the label to look up encoded as base64 + pub label_b64: AkdLabelB64, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct LookupData { + pub lookup_proof: akd::LookupProof, + pub epoch_data: EpochData, +} + +#[instrument(skip_all)] +pub async fn lookup_handler( + State(AppState { directory, .. }): State, + Json(LookupRequest { label_b64 }): Json, +) -> (StatusCode, Json>) { + info!("Handling get public key request"); + let lookup_proof = directory.lookup(label_b64.into()).await; + + match lookup_proof { + Ok((lookup_proof, epoch_hash)) => ( + StatusCode::OK, + Json(Response::success(LookupData { + lookup_proof, + epoch_data: epoch_hash.into(), + })), + ), + Err(e) => { + error!(err = ?e, "Failed to get AKD public key"); + (StatusCode::INTERNAL_SERVER_ERROR, Json(Response::fail(e))) + } + } +} diff --git a/akd/crates/reader/src/routes/mod.rs b/akd/crates/reader/src/routes/mod.rs index a1ef4e49e1..ced262c649 100644 --- a/akd/crates/reader/src/routes/mod.rs +++ b/akd/crates/reader/src/routes/mod.rs @@ -1,9 +1,86 @@ +use akd::errors::AkdError; use axum::routing::get; +use serde::{Deserialize, Serialize}; +mod audit; +mod batch_lookup; +mod get_epoch_hash; +mod get_public_key; mod health; +mod key_history; +mod lookup; +pub mod response_types; use crate::AppState; pub fn api_routes() -> axum::Router { - axum::Router::new().route("/health", get(health::health_handler)) + axum::Router::new() + .route("/health", get(health::health_handler)) + .route("/public_key", get(get_public_key::get_public_key_handler)) + .route("/epoch_hash", get(get_epoch_hash::get_epoch_hash_handler)) + .route("/lookup", get(lookup::lookup_handler)) + .route("/key_history", get(key_history::key_history_handler)) + .route("/batch_lookup", get(batch_lookup::batch_lookup_handler)) + .route("/audit", get(audit::audit_handler)) +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct Response { + success: bool, + data: Option, + error: Option, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct ResponseError { + pub akd_error_type: String, + pub message: String, +} + +impl From for ResponseError { + fn from(err: AkdError) -> Self { + ResponseError { + akd_error_type: match &err { + AkdError::TreeNode(_) => "TreeNode".to_string(), + AkdError::Directory(err) => match err { + akd::errors::DirectoryError::Verification(_) => "VerificationError".to_string(), + akd::errors::DirectoryError::InvalidEpoch(_) => "InvalidEpoch".to_string(), + akd::errors::DirectoryError::ReadOnlyDirectory(_) => { + "ReadOnlyDirectory".to_string() + } + akd::errors::DirectoryError::Publish(_) => "Publish".to_string(), + }, + AkdError::AzksErr(_) => "AzksErr".to_string(), + AkdError::Vrf(_) => "Vrf".to_string(), + AkdError::Storage(err) => match err { + akd::errors::StorageError::NotFound(_) => "NotFound".to_string(), + akd::errors::StorageError::Transaction(_) => "Transaction".to_string(), + akd::errors::StorageError::Connection(_) => "Connection".to_string(), + akd::errors::StorageError::Other(_) => "Other".to_string(), + }, + AkdError::AuditErr(_) => "AuditErr".to_string(), + AkdError::Parallelism(_) => "Parallelism".to_string(), + AkdError::TestErr(_) => "TestErr".to_string(), + }, + message: err.to_string(), + } + } +} + +impl<'a, T: Serialize + Deserialize<'a>> Response { + fn success(data: T) -> Self { + Self { + success: true, + data: Some(data), + error: None, + } + } + + fn fail(err: AkdError) -> Self { + Self { + success: false, + data: None, + error: Some(err.into()), + } + } } diff --git a/akd/crates/reader/src/routes/response_types.rs b/akd/crates/reader/src/routes/response_types.rs new file mode 100644 index 0000000000..6880170c31 --- /dev/null +++ b/akd/crates/reader/src/routes/response_types.rs @@ -0,0 +1,9 @@ +use crate::routes::Response; + +pub type AuditResponse = Response; +pub type BatchLookupData = Response; +pub type HealthData = Response; +pub type HistoryData = Response; +pub type LookupData = Response; +pub type EpochData = Response; +pub type PublicKeyData = Response;