diff --git a/akd/crates/akd_storage/src/akd_database.rs b/akd/crates/akd_storage/src/akd_database.rs index d3ef929fe7..c747d5ff28 100644 --- a/akd/crates/akd_storage/src/akd_database.rs +++ b/akd/crates/akd_storage/src/akd_database.rs @@ -11,7 +11,9 @@ use akd::{ use async_trait::async_trait; use crate::{ - db_config::DatabaseType, vrf_key_config::VrfKeyConfig, vrf_key_database::VrfKeyDatabase, + db_config::DatabaseType, + vrf_key_config::VrfKeyConfig, + vrf_key_database::{VrfKeyConfigError, VrfKeyDatabase}, }; #[derive(Debug, Clone)] @@ -29,8 +31,8 @@ impl AkdDatabase { AkdDatabase { db, vrf_key_config } } - pub fn vrf_key_database(&self) -> VrfKeyDatabase { - VrfKeyDatabase::new(self.db.clone(), self.vrf_key_config.clone()) + pub async fn vrf_key_database(&self) -> Result { + VrfKeyDatabase::new(self.db.clone(), self.vrf_key_config.clone()).await } } diff --git a/akd/crates/akd_storage/src/db_config.rs b/akd/crates/akd_storage/src/db_config.rs index fd37853ba0..ef81e70463 100644 --- a/akd/crates/akd_storage/src/db_config.rs +++ b/akd/crates/akd_storage/src/db_config.rs @@ -1,7 +1,7 @@ use akd::errors::StorageError; use serde::{Deserialize, Serialize}; -use crate::ms_sql::MsSql; +use crate::{ms_sql::MsSql, vrf_key_database::VrfKeyStorageError}; #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(tag = "type")] @@ -37,3 +37,13 @@ impl DbConfig { Ok(db) } } + +impl DatabaseType { + pub async fn get_existing_vrf_root_key_hash( + &self, + ) -> Result>, VrfKeyStorageError> { + match self { + DatabaseType::MsSql(db) => db.get_existing_vrf_root_key_hash().await, + } + } +} diff --git a/akd/crates/akd_storage/src/ms_sql/mod.rs b/akd/crates/akd_storage/src/ms_sql/mod.rs index c75a56aa96..0e903f6591 100644 --- a/akd/crates/akd_storage/src/ms_sql/mod.rs +++ b/akd/crates/akd_storage/src/ms_sql/mod.rs @@ -173,6 +173,41 @@ impl MsSql { } impl MsSql { + pub async fn get_existing_vrf_root_key_hash( + &self, + ) -> Result>, VrfKeyStorageError> { + debug!("Checking for existence of any VRF key in database"); + + let mut conn = self.get_connection().await.map_err(|err| { + error!(%err, "Failed to get DB connection for VRF key existence check"); + VrfKeyStorageError + })?; + + let sql = vrf_key::get_first_root_key_hash(); + let query_stream = conn.query(sql.sql(), &sql.params()).await.map_err(|err| { + error!(%err, "Failed to execute VRF key count query"); + VrfKeyStorageError + })?; + let row = query_stream.into_row().await.map_err(|err| { + error!(%err, "Failed to fetch VRF key count row"); + VrfKeyStorageError + })?; + match row { + Some(row) => { + let hash = sql.parse(&row).map_err(|err| { + error!(%err, "Failed to parse VRF key count from row"); + VrfKeyStorageError + })?; + debug!("Existing VRF key found in database"); + Ok(Some(hash)) + } + None => { + debug!("No VRF keys found in database"); + Ok(None) + } + } + } + #[instrument(skip(self, config), level = "debug")] pub async fn get_vrf_key( &self, diff --git a/akd/crates/akd_storage/src/ms_sql/tables/vrf_key.rs b/akd/crates/akd_storage/src/ms_sql/tables/vrf_key.rs index cf6e8b4413..95198962f6 100644 --- a/akd/crates/akd_storage/src/ms_sql/tables/vrf_key.rs +++ b/akd/crates/akd_storage/src/ms_sql/tables/vrf_key.rs @@ -11,6 +11,23 @@ use crate::{ vrf_key_database::VrfKeyTableData, }; +pub fn get_first_root_key_hash() -> QueryStatement> { + debug!("Building has_vrf_key statement"); + let sql = format!( + r#" + SELECT TOP 1 root_key_hash + FROM {} + "#, + TABLE_VRF_KEYS + ); + QueryStatement::new(sql, SqlParams::new(), |row: &ms_database::Row| { + let hash: &[u8] = row + .get("root_key_hash") + .ok_or_else(|| StorageError::Other("root_key_hash is NULL or missing".to_string()))?; + Ok(hash.to_vec()) + }) +} + pub fn get_statement( config: &VrfKeyConfig, ) -> Result, VrfRootKeyError> { diff --git a/akd/crates/akd_storage/src/vrf_key_database.rs b/akd/crates/akd_storage/src/vrf_key_database.rs index 6bf0dded27..87c2070a53 100644 --- a/akd/crates/akd_storage/src/vrf_key_database.rs +++ b/akd/crates/akd_storage/src/vrf_key_database.rs @@ -32,6 +32,12 @@ pub struct VrfKeyCreationError; #[error("Internal VRF key storage error")] pub struct VrfKeyStorageError; +#[derive(Debug, Error)] +#[error( + "VRF key configuration error. Check root key settings are valid for the existing AKD database." +)] +pub struct VrfKeyConfigError; + #[derive(Debug, Clone)] pub struct VrfKeyDatabase { db: DatabaseType, @@ -40,13 +46,44 @@ pub struct VrfKeyDatabase { } impl VrfKeyDatabase { - pub fn new(db: DatabaseType, config: VrfKeyConfig) -> VrfKeyDatabase { - VrfKeyDatabase { - db, - vrf_key_config: config, + pub async fn new( + db: DatabaseType, + config: VrfKeyConfig, + ) -> Result { + let new_key_database = VrfKeyDatabase { + db: db.clone(), + vrf_key_config: config.clone(), cached_vrf_key: None, + }; + + let expected_key_hash = config.root_key_hash().map_err(|_| { + error!("Failed to compute root key hash from configuration. Likely due to invalid key material."); + VrfKeyConfigError + })?; + + let existing_key_hash = db.get_existing_vrf_root_key_hash().await.map_err(|err| { + error!(%err, "Error checking for existing VRF key in database"); + VrfKeyConfigError + })?; + + match existing_key_hash { + Some(hash) => { + if hash != expected_key_hash { + error!( + "Existing VRF root key hash in database does not match configured root key hash. \ + This indicates a misconfiguration or potential security issue." + ); + return Err(VrfKeyConfigError); + } + Ok(new_key_database) + } + None => { + info!("No existing VRF key found in database. Proceeding to create a new key if needed."); + Ok(new_key_database) + } } } + async fn get_vrf_key(&self) -> Result { match &self.db { DatabaseType::MsSql(db) => db.get_vrf_key(&self.vrf_key_config).await, diff --git a/akd/crates/akd_test_utility/src/main.rs b/akd/crates/akd_test_utility/src/main.rs index bdceecdfcb..9773ec367a 100644 --- a/akd/crates/akd_test_utility/src/main.rs +++ b/akd/crates/akd_test_utility/src/main.rs @@ -180,7 +180,9 @@ async fn main() -> Result<()> { return Ok(()); } - let mut directory = Directory::::new(storage_manager, state.vrf_key_database()) + let vrf_key_database = state.vrf_key_database().await?; + + let mut directory = Directory::::new(storage_manager, vrf_key_database) .await .context("Failed to create AKD directory")?;