allow only one vrf key

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, VrfKeyConfigError> {
+        VrfKeyDatabase::new(self.db.clone(), self.vrf_key_config.clone()).await
     }
 }
 

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<Option<Vec<u8>>, VrfKeyStorageError> {
+        match self {
+            DatabaseType::MsSql(db) => db.get_existing_vrf_root_key_hash().await,
+        }
+    }
+}

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<Option<Vec<u8>>, 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,

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<Vec<u8>> {
+    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<QueryStatement<VrfKeyTableData>, VrfRootKeyError> {

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<VrfKeyDatabase, VrfKeyConfigError> {
+        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<VrfKeyTableData, VrfKeyRetrievalError> {
         match &self.db {
             DatabaseType::MsSql(db) => db.get_vrf_key(&self.vrf_key_config).await,

akd/crates/akd_test_utility/src/main.rs

@@ -180,7 +180,9 @@ async fn main() -> Result<()> {
         return Ok(());
     }
 
-    let mut directory = Directory::<TC, _, _>::new(storage_manager, state.vrf_key_database())
+    let vrf_key_database = state.vrf_key_database().await?;
+
+    let mut directory = Directory::<TC, _, _>::new(storage_manager, vrf_key_database)
         .await
         .context("Failed to create AKD directory")?;