Can retrieve Shannon's entropy of all stored keys

src/entropy/mod.rs

@@ -0,0 +1,63 @@
+pub(crate) trait EntropyAccumulator {
+    fn add_bytes(&mut self, bytes: &[u8]);
+    fn get_entropy(&self) -> f64;
+}
+
+pub(crate) struct ShannonEntropyAccumulator {
+    /// Counter for each byte value
+    bytes_counter: [u64; 256],
+    /// Total received bytes
+    total_bytes: u64,
+}
+
+impl ShannonEntropyAccumulator {
+    pub(crate) fn new() -> Self {
+        Self {
+            bytes_counter: [0; 256],
+            total_bytes: 0,
+        }
+    }
+}
+
+impl EntropyAccumulator for ShannonEntropyAccumulator {
+    fn add_bytes(&mut self, bytes: &[u8]) {
+        for byte in bytes {
+            self.bytes_counter[*byte as usize] += 1;
+        }
+        self.total_bytes += bytes.len() as u64;
+    }
+
+    fn get_entropy(&self) -> f64 {
+        let mut entropy = 0.0;
+        for count in self.bytes_counter.iter() {
+            if *count == 0 {
+                continue;
+            }
+            let symbol_probability = *count as f64 / self.total_bytes as f64;
+            entropy -= symbol_probability * symbol_probability.log2();
+        }
+        entropy
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::entropy::EntropyAccumulator;
+
+    #[test]
+    fn test_shannon_entropy_accumulator() {
+        let mut entropy_accumulator_1 = super::ShannonEntropyAccumulator::new();
+        entropy_accumulator_1.add_bytes(&[0, 0, 0, 0, 0, 0]);
+        assert_eq!(entropy_accumulator_1.get_entropy(), 0.0);
+
+        let mut entropy_accumulator_2 = super::ShannonEntropyAccumulator::new();
+        entropy_accumulator_2.add_bytes(&[0x00, 0x00, 0x01, 0x01, 0x02]);
+        assert_eq!(entropy_accumulator_2.get_entropy(), 1.5219280948873621);
+
+        let mut entropy_accumulator_3 = super::ShannonEntropyAccumulator::new();
+        entropy_accumulator_3.add_bytes("Souvent sur la montagne, à l’ombre du vieux chêne,\n".as_bytes());
+        assert_eq!(entropy_accumulator_3.get_entropy(), 4.465641023041018);
+        entropy_accumulator_3.add_bytes("Au coucher du soleil, tristement je m’assieds ;\n".as_bytes());
+        assert_eq!(entropy_accumulator_3.get_entropy(), 4.507894683096287);
+    }
+}

src/lib.rs

@@ -15,6 +15,7 @@ pub mod server;
 pub mod routes;
 pub mod qkd_manager;
 pub mod config;
+pub(crate) mod entropy;
 
 
 /// Cast a string to an io::Error
@@ -82,6 +83,4 @@ mod test {
         assert_eq!(err.kind(), std::io::ErrorKind::Other);
         assert_eq!(err.to_string(), "test");
     }
-}
-
-// TODO shannon entropy
+}

src/qkd_manager/http_response_obj.rs

@@ -98,6 +98,13 @@ pub(crate) struct ResponseQkdKeysList {
 }
 impl HttpResponseBody for ResponseQkdKeysList {}
 
+#[derive(serde::Serialize, Debug, PartialEq)]
+pub(crate) struct ResponseTotalKeysEntropy {
+    pub(crate) total_entropy: f64,
+}
+
+impl HttpResponseBody for ResponseTotalKeysEntropy {}
+
 #[cfg(test)]
 mod test {
     use crate::qkd_manager::http_response_obj::HttpResponseBody;
@@ -151,4 +158,13 @@ mod test {
         let response_qkd_sae_info_json = response_qkd_sae_info.to_json().unwrap();
         assert_eq!(response_qkd_sae_info_json.replace("\r", ""), "{\n  \"SAE_ID\": 1,\n  \"KME_ID\": 1\n}");
     }
+
+    #[test]
+    fn test_serialize_response_total_keys_entropy() {
+        let response_total_keys_entropy = super::ResponseTotalKeysEntropy {
+            total_entropy: 1.0,
+        };
+        let response_total_keys_entropy_json = response_total_keys_entropy.to_json().unwrap();
+        assert_eq!(response_total_keys_entropy_json.replace("\r", ""), "{\n  \"total_entropy\": 1.0\n}");
+    }
 }

src/qkd_manager/mod.rs

@@ -12,7 +12,8 @@ use log::error;
 use sha1::Digest;
 use crate::qkd_manager::http_response_obj::ResponseQkdKeysList;
 use crate::qkd_manager::QkdManagerResponse::TransmissionError;
-use crate::{KmeId, QkdEncKey, SaeClientCertSerial, SaeId};
+use crate::{io_err, KmeId, QkdEncKey, SaeClientCertSerial, SaeId};
+use crate::entropy::{EntropyAccumulator, ShannonEntropyAccumulator};
 
 /// QKD manager interface, can be cloned for instance in each request handler task
 #[derive(Clone)]
@@ -25,6 +26,8 @@ pub struct QkdManager {
     pub(crate) dir_watcher: Arc<Mutex<Vec<notify::RecommendedWatcher>>>,
     /// The ID of the KME this QKD manager belongs to
     pub kme_id: KmeId,
+    /// Shannon's entropy calculator for keys stored in the database
+    shannon_entropy_calculator: Arc<Mutex<ShannonEntropyAccumulator>>,
 }
 
 impl QkdManager {
@@ -60,6 +63,7 @@ impl QkdManager {
             response_rx,
             dir_watcher,
             kme_id: this_kme_id,
+            shannon_entropy_calculator: Arc::new(Mutex::new(ShannonEntropyAccumulator::new())),
         }
     }
 
@@ -74,20 +78,27 @@ impl QkdManager {
     }
 
     /// Add a new QKD key to the database
+    /// Increases the total entropy of all keys in the database
     /// # Arguments
     /// * `key` - The QKD key to add (key + origin SAE ID + target SAE ID)
     /// # Returns
     /// Ok if the key was added successfully, an error otherwise
     pub fn add_pre_init_qkd_key(&self, key: PreInitQkdKeyWrapper) -> Result<QkdManagerResponse, QkdManagerResponse> {
-        self.command_tx.send(QkdManagerCommand::AddPreInitKey(key)).map_err(|_| {
+        const EXPECTED_QKD_MANAGER_RESPONSE: QkdManagerResponse = QkdManagerResponse::Ok;
+
+        self.command_tx.send(QkdManagerCommand::AddPreInitKey(key.to_owned())).map_err(|_| {
             TransmissionError
         })?;
-        match self.response_rx.recv().map_err(|_| {
+        let add_key_status = self.response_rx.recv().map_err(|_| {
             TransmissionError
-        })? {
-            QkdManagerResponse::Ok => Ok(QkdManagerResponse::Ok), // Ok is the QkdManagerResponse expected here
-            qkd_response_error => Err(qkd_response_error),
+        })?;
+        if add_key_status != EXPECTED_QKD_MANAGER_RESPONSE {
+            return Err(add_key_status);
         }
+        self.shannon_entropy_calculator.lock().map_err(|_| {
+            TransmissionError
+        })?.add_bytes(&key.key);
+        Ok(EXPECTED_QKD_MANAGER_RESPONSE)
     }
 
     /// Get a QKD key from the database (shall be called by the master SAE)
@@ -254,6 +265,20 @@ impl QkdManager {
             qkd_response_error => Err(qkd_response_error),
         }
     }
+
+    /// Get the Shannon entropy of all stored keys
+    /// # Returns
+    /// The total Shannon entropy of all stored keys, an error in case of concurrency issues
+    pub async fn get_total_keys_shannon_entropy(&self) -> Result<f64, io::Error> {
+        let entropy_calculator = Arc::clone(&self.shannon_entropy_calculator);
+        Ok(tokio::task::spawn_blocking(move || {
+            Ok::<f64, io::Error>(entropy_calculator
+                .lock()
+                .map_err(|_| io_err("Mutex locking error"))?
+                .get_entropy())
+        }).await
+            .map_err(|_| io_err("Async task error"))??)
+    }
 }
 
 /// A Pre-init QKD key, with its origin and target KME IDs
@@ -383,17 +408,34 @@ pub enum QkdManagerResponse {
 #[cfg(test)]
 mod test {
     use serial_test::serial;
+    use crate::QkdEncKey;
 
     const CLIENT_CERT_SERIAL_SIZE_BYTES: usize = 20;
 
-    #[test]
-    fn test_add_qkd_key() {
+    #[tokio::test]
+    async fn test_add_qkd_key() {
         const SQLITE_DB_PATH: &'static str = ":memory:";
         let qkd_manager = super::QkdManager::new(SQLITE_DB_PATH, 1);
         let key = super::PreInitQkdKeyWrapper::new(1, &[0; crate::QKD_KEY_SIZE_BYTES]).unwrap();
         let response = qkd_manager.add_pre_init_qkd_key(key);
         assert!(response.is_ok());
         assert_eq!(response.unwrap(), super::QkdManagerResponse::Ok);
+        assert_eq!(qkd_manager.get_total_keys_shannon_entropy().await.unwrap(), 0.0);
+    }
+
+    #[tokio::test]
+    async fn test_stored_keys_entropy() {
+        const SQLITE_DB_PATH: &'static str = ":memory:";
+        let first_key: QkdEncKey = <[u8; crate::QKD_KEY_SIZE_BYTES]>::try_from("ABCDEFGHIJKLMNOPQRSTUVWXYZ012345".as_bytes()).unwrap();
+        let second_key: QkdEncKey = <[u8; crate::QKD_KEY_SIZE_BYTES]>::try_from("6789+-abcdefghijklmnopqrstuvwxyz".as_bytes()).unwrap();
+
+        let qkd_manager = super::QkdManager::new(SQLITE_DB_PATH, 1);
+        let key = super::PreInitQkdKeyWrapper::new(1, &first_key).unwrap();
+        qkd_manager.add_pre_init_qkd_key(key).unwrap();
+        assert_eq!(qkd_manager.get_total_keys_shannon_entropy().await.unwrap(), 5.0);
+        let key = super::PreInitQkdKeyWrapper::new(1, &second_key).unwrap();
+        qkd_manager.add_pre_init_qkd_key(key).unwrap();
+        assert_eq!(qkd_manager.get_total_keys_shannon_entropy().await.unwrap(), 6.0);
     }
 
     #[test]

src/routes/keys/mod.rs

@@ -7,6 +7,7 @@ use hyper::body::Bytes;
 use crate::routes::{EtsiSaeQkdRoutesV1, RequestContext};
 
 mod get_key;
+mod route_entropy;
 
 /// Dispatches the request to the correct function
 pub(super) async fn key_handler(rcx: &RequestContext<'_>, req: Request<body::Incoming>, uri_segments: &[&str]) -> Result<Response<Full<Bytes>>, Infallible> {
@@ -17,6 +18,8 @@ pub(super) async fn key_handler(rcx: &RequestContext<'_>, req: Request<body::Inc
         ([slave_sae_id, "enc_keys"], &hyper::Method::POST) => get_key::route_get_key(rcx, req, slave_sae_id).await,
         // Get key(s) from a slave SAE, with ID provided by the master SAE
         ([slave_sae_id, "dec_keys"], &hyper::Method::POST) => get_key::route_get_key_with_id(rcx, req, slave_sae_id).await,
+        // Retrieve Shannon's entropy for all stored keys in KME database
+        (["entropy", "total"], &hyper::Method::GET) => route_entropy::route_get_entropy_total(rcx, req).await,
         // Route not found
         _ => EtsiSaeQkdRoutesV1::not_found(),
     }

src/routes/keys/route_entropy.rs

@@ -0,0 +1,30 @@
+use std::convert::Infallible;
+use http_body_util::Full;
+use hyper::{body, Request, Response};
+use hyper::body::Bytes;
+use log::error;
+use crate::qkd_manager::http_response_obj::HttpResponseBody;
+use crate::routes::request_context::RequestContext;
+
+pub(in crate::routes) async fn route_get_entropy_total(rcx: &RequestContext<'_>, _req: Request<body::Incoming>) -> Result<Response<Full<Bytes>>, Infallible> {
+    // Get the total entropy from stored keys
+    let entropy = match rcx.qkd_manager.get_total_keys_shannon_entropy().await {
+        Ok(entropy) => entropy,
+        Err(e) => {
+            error!("Error getting total entropy: {}", e.to_string());
+            return super::EtsiSaeQkdRoutesV1::internal_server_error();
+        }
+    };
+    let total_entropy_response_obj = crate::qkd_manager::http_response_obj::ResponseTotalKeysEntropy {
+        total_entropy: entropy,
+    };
+    // Return the total entropy as a JSON object response
+    let total_entropy_response_obj_json = match total_entropy_response_obj.to_json() {
+        Ok(json) => json,
+        Err(_) => {
+            error!("Error serializing total entropy object");
+            return super::EtsiSaeQkdRoutesV1::internal_server_error();
+        }
+    };
+    Ok(crate::routes::EtsiSaeQkdRoutesV1::json_response_from_str(&total_entropy_response_obj_json))
+}

tests/data/total_entropy.json

@@ -0,0 +1,3 @@
+{
+  "total_entropy": 3.6123008763572906
+}

tests/entropy.rs

@@ -0,0 +1,22 @@
+use const_format::concatcp;
+use serial_test::serial;
+use crate::common::util::assert_string_equal;
+
+mod common;
+
+#[tokio::test]
+#[serial]
+async fn get_total_entropy() {
+    const EXPECTED_BODY: &'static str = include_str!("data/total_entropy.json");
+    const REQUEST_URL: &'static str = concatcp!("https://", common::HOST_PORT ,"/api/v1/keys/entropy/total");
+
+    common::setup();
+    let reqwest_client = common::setup_cert_auth_reqwest_client();
+
+    let response = reqwest_client.get(REQUEST_URL).send().await;
+    assert!(response.is_ok());
+    let response = response.unwrap();
+    assert_eq!(response.status(), 200);
+    let response_body = response.text().await.unwrap();
+    assert_string_equal(&response_body, EXPECTED_BODY);
+}