Manage keys in KME level: keys are now added for specific KMEs as pre…

…-init, and fully initialized for a SAE pair when SAE do the request. Add directory watching for automatically add new QKD keys

Cargo.toml

@@ -27,6 +27,9 @@ base64 = "0.21.5"
 async-trait = "0.1.75"
 rustls-pemfile = "2.0.0"
 log = "0.4.20"
+simple_logger = "4.3.3"
+notify = "6.1.1"
+clap = { version = "4.4.18", features = ["derive"] }
 
 [dev-dependencies]
 assert_cmd = "2.0.12"

README.md

@@ -52,7 +52,7 @@ Response example:
 ```json
 {
   "source_KME_ID": "1",
-  "target_KME_ID": "?? TODO",
+  "target_KME_ID": "2",
   "master_SAE_ID": "1",
   "slave_SAE_ID": "2",
   "key_size": 256,

src/lib.rs

@@ -29,6 +29,9 @@ fn io_err(e: &str) -> io::Error {
 /// The size of the QKD key in bytes. This is equal to QKD_MIN_KEY_SIZE_BITS and QKD_MAX_KEY_SIZE_BITS, as we offer no flexibility in key size
 pub const QKD_KEY_SIZE_BITS: usize = 256;
 
+/// The size of the QKD key in bytes
+pub const QKD_KEY_SIZE_BYTES: usize = QKD_KEY_SIZE_BITS / 8;
+
 /// The minimum size of the QKD key in bits, returned in HTTP responses
 pub const QKD_MIN_KEY_SIZE_BITS: usize = 256;
 
@@ -50,8 +53,14 @@ pub const CLIENT_CERT_SERIAL_SIZE_BYTES: usize = 20;
 /// Location of the SQLite database file used by the KME to store keys, use ":memory:" for in-memory database
 pub const MEMORY_SQLITE_DB_PATH: &'static str = ":memory:";
 
-/// The ID of this KME, used to identify the KME in the database and across the network
-pub const THIS_KME_ID: i64 = 1; // TODO: change
+/// The type of SAE ID
+pub type SaeId = i64;
+
+/// The type of KME ID
+pub type KmeId = i64;
+
+/// Type for QKD encryption key: basically a byte array
+pub type QkdEncKey = [u8; QKD_KEY_SIZE_BYTES];
 
 #[cfg(test)]
 mod test {

src/main.rs

@@ -1,66 +1,130 @@
+use std::io::{BufReader, Read};
+use std::path::Path;
+use std::sync::Arc;
 use log::error;
-use qkd_kme_server::qkd_manager::{QkdKey, QkdManager};
+use clap::Parser;
+use notify::{EventKind, RecursiveMode, Watcher};
+use notify::event::{AccessKind, AccessMode};
+use qkd_kme_server::qkd_manager::{PreInitQkdKeyWrapper, QkdManager};
 use qkd_kme_server::routes::QKDKMERoutesV1;
 
 #[tokio::main]
 async fn main() {
+    simple_logger::SimpleLogger::new().init().unwrap();
+    let args = Args::parse();
+
+    println!("{:?}", args);
+
+
     let server = qkd_kme_server::server::Server {
         listen_addr: "127.0.0.1:3000".to_string(),
         ca_client_cert_path: "certs/CA-zone1.crt".to_string(),
         server_cert_path: "certs/kme1.crt".to_string(),
         server_key_path: "certs/kme1.key".to_string(),
     };
 
-    let qkd_manager = QkdManager::new(qkd_kme_server::MEMORY_SQLITE_DB_PATH);
+    let qkd_manager = Arc::new(QkdManager::new(qkd_kme_server::MEMORY_SQLITE_DB_PATH, args.this_kme_id));
     if qkd_manager.add_sae(1,
-    &[0x70, 0xf4, 0x4f, 0x56, 0x0c, 0x3f, 0x27, 0xd4, 0xb2, 0x11, 0xa4, 0x78, 0x13, 0xaf, 0xd0, 0x3c, 0x03, 0x81, 0x3b, 0x8e]
+                           1,
+                           &Some([0x70, 0xf4, 0x4f, 0x56, 0x0c, 0x3f, 0x27, 0xd4, 0xb2, 0x11, 0xa4, 0x78, 0x13, 0xaf, 0xd0, 0x3c, 0x03, 0x81, 0x3b, 0x8e])
     ).is_err() {
         error!("Error adding SAE to QKD manager");
         return;
     }
     if qkd_manager.add_sae(2,
-                           &[0x70, 0xf4, 0x4f, 0x56, 0x0c, 0x3f, 0x27, 0xd4, 0xb2, 0x11, 0xa4, 0x78, 0x13, 0xaf, 0xd0, 0x3c, 0x03, 0x81, 0x3b, 0x92]
+                           1,
+                           &Some([0x70, 0xf4, 0x4f, 0x56, 0x0c, 0x3f, 0x27, 0xd4, 0xb2, 0x11, 0xa4, 0x78, 0x13, 0xaf, 0xd0, 0x3c, 0x03, 0x81, 0x3b, 0x92])
     ).is_err() {
         error!("Error adding SAE to QKD manager");
         return;
     }
 
-    let qkd_key_1 = match QkdKey::new(
-        1,
-        2,
-        b"this_is_secret_key_1_of_32_bytes",
-    ) {
-        Ok(qkd_key) => qkd_key,
-        Err(_) => {
-            error!("Error creating QKD key");
+    let mut watchers: Vec<notify::RecommendedWatcher> = Vec::new();
+
+    for kme_dir in args.dirs_to_watch_other_kme_ids {
+        extract_all_keys_from_dir(&kme_dir.dir, kme_dir.kme_id, &qkd_manager);
+        let kme_id = kme_dir.kme_id;
+        let qkd_manager = Arc::clone(&qkd_manager);
+        watchers.push(match notify::recommended_watcher(move |res: Result<notify::Event, notify::Error>| {
+            match res {
+                Ok(event) => {
+                    if let EventKind::Access(AccessKind::Close(AccessMode::Write)) = event.kind {
+                        extract_all_keys_from_file(&event.paths[0].to_str().unwrap(), kme_id, &qkd_manager);
+                    }
+                }
+                Err(e) => {
+                    println!("Watch error: {:?}", e);
+                    return;
+                }
+            }
+        }) {
+            Ok(watcher) => watcher,
+            Err(e) => {
+                error!("Error creating watcher: {:?}", e);
+                return;
+            }
+        });
+        if watchers.iter_mut().last().unwrap().watch(Path::new(&kme_dir.dir), RecursiveMode::NonRecursive).is_err() {
+            error!("Error watching directory: {:?}", kme_dir.dir);
             return;
         }
-    };
+    }
 
-    if qkd_manager.add_qkd_key(qkd_key_1).is_err() {
-        error!("Error adding key to QKD manager");
+    if server.run::<QKDKMERoutesV1>(&qkd_manager).await.is_err() {
+        error!("Error running HTTP server");
         return;
     }
+}
 
-    let qkd_key_2 = match QkdKey::new(
-        1,
-        1,
-        b"this_is_secret_key_1_of_32_bytes",
-    ) {
-        Ok(qkd_key) => qkd_key,
-        Err(_) => {
-            error!("Error creating QKD key");
-            return;
+#[derive(Parser, Debug)]
+struct Args {
+    #[arg(long)]
+    this_kme_id: i64,
+    #[arg(long("kme_id,dir_watch"))]
+    dirs_to_watch_other_kme_ids: Vec<DirWatchOtherKmesArgs>,
+}
+
+#[derive(Debug, Clone)]
+struct DirWatchOtherKmesArgs {
+    dir: String,
+    kme_id: i64,
+}
+
+impl std::str::FromStr for DirWatchOtherKmesArgs {
+    type Err = std::io::Error;
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        let kme_id_and_dir = s.split(",").collect::<Vec<&str>>();
+        if kme_id_and_dir.len() != 2 {
+            return Err(std::io::Error::other("Invalid format"));
         }
-    };
+        let kme_id = i64::from_str(kme_id_and_dir[0]).map_err(|_| std::io::Error::other("Invalid format"))?;
+        Ok(Self {
+            dir: kme_id_and_dir[1].to_string(),
+            kme_id,
+        })
+    }
+}
 
-    if qkd_manager.add_qkd_key(qkd_key_2).is_err() {
-        error!("Error adding key to QKD manager");
-        return;
+// TODO: move to QKD manager struct
+fn extract_all_keys_from_file(file_path: &str, other_kme_id: i64, qkd_manager: &QkdManager) {
+    let file = std::fs::File::open(file_path).unwrap();
+    let mut reader = BufReader::with_capacity(32, file);
+    let mut buffer = [0; 32];
+    while let Ok(_) = reader.read_exact(&mut buffer) {
+        let qkd_key = PreInitQkdKeyWrapper::new(
+            other_kme_id,
+            &buffer,
+        ).unwrap();
+        qkd_manager.add_pre_init_qkd_key(qkd_key).unwrap();
     }
+}
 
-    if server.run::<QKDKMERoutesV1>(&qkd_manager).await.is_err() {
-        error!("Error running HTTP server");
-        return;
+fn extract_all_keys_from_dir(dir_path: &str, other_kme_id: i64, qkd_manager: &QkdManager) {
+    let paths = std::fs::read_dir(dir_path).unwrap();
+    for path in paths {
+        let path = path.unwrap().path();
+        if path.is_file() {
+            extract_all_keys_from_file(path.to_str().unwrap(), other_kme_id, qkd_manager);
+        }
     }
 }

src/qkd_manager/http_response_obj.rs

@@ -1,6 +1,7 @@
 //! Objects serialized to HTTP response body
 
 use std::io;
+use crate::{KmeId, SaeId};
 
 /// Trait to be implemented by objects that can be serialized to JSON
 pub(crate) trait HttpResponseBody where Self: serde::Serialize {
@@ -82,8 +83,9 @@ pub(crate) struct ResponseQkdKey {
 #[allow(non_snake_case)]
 pub(crate) struct ResponseQkdSAEInfo {
     /// SAE ID of the SAE
-    pub(crate) SAE_ID: i64,
-    // TODO: KME ID ?
+    pub(crate) SAE_ID: SaeId,
+    /// KME ID SAE belongs to
+    pub(crate) KME_ID: KmeId,
 }
 
 impl HttpResponseBody for ResponseQkdSAEInfo {} // can't use Derive macro because of the generic constraint
@@ -144,8 +146,9 @@ mod test {
     fn test_serialize_response_qkd_sae_info() {
         let response_qkd_sae_info = super::ResponseQkdSAEInfo {
             SAE_ID: 1,
+            KME_ID: 1,
         };
         let response_qkd_sae_info_json = response_qkd_sae_info.to_json().unwrap();
-        assert_eq!(response_qkd_sae_info_json, "{\n  \"SAE_ID\": 1\n}");
+        assert_eq!(response_qkd_sae_info_json, "{\n  \"SAE_ID\": 1,\n  \"KME_ID\": 1\n}");
     }
 }

src/qkd_manager/init_qkd_database.sql

@@ -0,0 +1,24 @@
+/* Uninitialized keys, available for SAEs in this KME and other_kme */
+CREATE TABLE IF NOT EXISTS uninit_keys (
+    id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
+    key_uuid TEXT NOT NULL,
+    key BLOB NOT NULL,
+    other_kme_id INTEGER NOT NULL
+);
+
+/* Keys assigned to SAEs */
+CREATE TABLE IF NOT EXISTS keys (
+    id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
+    key_uuid TEXT NOT NULL,
+    key BLOB NOT NULL,
+    origin_sae_id INTEGER NOT NULL,
+    target_sae_id INTEGER NOT NULL,
+    FOREIGN KEY (origin_sae_id) REFERENCES saes(sae_id),
+    FOREIGN KEY (target_sae_id) REFERENCES saes(sae_id)
+);
+
+CREATE TABLE IF NOT EXISTS saes (
+    sae_id INTEGER PRIMARY KEY NOT NULL,
+    sae_certificate_serial BLOB,
+    kme_id INTEGER NOT NULL
+);