Call keystore.privkey() everytime we need it

Previously it was cached after the first call and it so happens that the key might sometimes change. Due to FFI layer issues, Libdrop instance cannot be destroyed and recreated thus leaving the only solution either a new API or the easier one - passing KeyStore in Arc<> and calling when needed. Signed-off-by: Lukas Pukenis <[email protected]>
NordSecurity · Sep 23, 2024 · 52a7f91 · 52a7f91
1 parent d5f4cba
commit 52a7f91
Show file tree

Hide file tree

Showing 8 changed files with 50 additions and 49 deletions.
diff --git a/changelog.md b/changelog.md
@@ -4,6 +4,7 @@
 * Add `base_dir` field to the `FilePending` event
 * Update moose tracker to v13.1.0
     * Adds the `transfer_intent_received` event
+* Added private key retrieval on demand instead of caching it
 
 ---
 <br>

diff --git a/drop-transfer/examples/udrop.rs b/drop-transfer/examples/udrop.rs
@@ -315,7 +315,12 @@ async fn main() -> anyhow::Result<()> {
 
     let auth = {
         let pubkey = drop_auth::PublicKey::from(PUB_KEY);
-        auth::Context::new(drop_auth::SecretKey::from(PRIV_KEY), move |_| Some(pubkey))
+        let privkey = move || {
+            let privkey = drop_auth::SecretKey::from(PRIV_KEY);
+            Some(privkey)
+        };
+
+        auth::Context::new(privkey, move |_| Some(pubkey))
     };
 
     let storage_file = matches.get_one::<String>("storage").unwrap();

diff --git a/drop-transfer/src/auth.rs b/drop-transfer/src/auth.rs
@@ -4,17 +4,17 @@ use drop_auth::{PublicKey, SecretKey};
 use hyper::{http::HeaderValue, Response};
 
 pub struct Context {
-    secret: SecretKey,
+    secret: Box<dyn Fn() -> Option<SecretKey> + Send + Sync>,
     public: Box<dyn Fn(IpAddr) -> Option<PublicKey> + Send + Sync>,
 }
 
 impl Context {
     pub fn new(
-        secret: SecretKey,
+        secret: impl Fn() -> Option<SecretKey> + Send + Sync + 'static,
         public: impl Fn(IpAddr) -> Option<PublicKey> + Send + Sync + 'static,
     ) -> Self {
         Self {
-            secret,
+            secret: Box::new(secret),
             public: Box::new(public),
         }
     }
@@ -28,7 +28,8 @@ impl Context {
         tokio::task::block_in_place(|| {
             let auth_req = drop_auth::http::Authorization::parse(auth_header_value)?;
             let pubkey = (self.public)(peer_ip)?;
-            drop_auth::authorize(nonce, &self.secret, &pubkey, &auth_req)
+            let secret = (self.secret)()?;
+            drop_auth::authorize(nonce, &secret, &pubkey, &auth_req)
         })
         .is_some()
     }
@@ -74,9 +75,10 @@ impl Context {
                 .context("Failed to parse 'www-authenticate' header")?;
 
             let public = (self.public)(peer_ip).context("Failed to fetch peer's public key")?;
+            let secret = (self.secret)().context("Failed to fetch private key")?;
 
             let ticket =
-                drop_auth::create_ticket_as_client(&self.secret, &public, resp, check_nonce_prefix)
+                drop_auth::create_ticket_as_client(&secret, &public, resp, check_nonce_prefix)
                     .context("Failed to create auth ticket")?;
 
             let value = HeaderValue::from_str(&ticket.to_string())?;
@@ -96,8 +98,9 @@ impl Context {
                 .context("Failed to parse 'www-authenticate' header")?;
 
             let public = (self.public)(peer_ip).context("Failed to fetch peer's public key")?;
+            let secret = (self.secret)().context("Failed to fetch private key")?;
 
-            let ticket = drop_auth::create_ticket_as_server(&self.secret, &public, resp)
+            let ticket = drop_auth::create_ticket_as_server(&secret, &public, resp)
                 .context("Failed to create auth ticket")?;
 
             let value = HeaderValue::from_str(&ticket.to_string())?;

diff --git a/drop-transfer/src/ws/server/v6.rs b/drop-transfer/src/ws/server/v6.rs
@@ -97,7 +97,7 @@ impl<'a> handler::HandlerInit for HandlerInit<'a> {
             .context("Failed to receive transfer request")?;
 
         // print msg as ascii
-        debug!(self.logger, "************** msg:\n\t{msg:?}");
+        debug!(self.logger, "msg:\n\t{msg:?}");
 
         let msg = msg.to_str().ok().context("Expected JSON message")?;
         debug!(self.logger, "Request received:\n\t{msg}");

diff --git a/norddrop/src/device.rs b/norddrop/src/device.rs
@@ -5,7 +5,7 @@ use std::{
 };
 
 use drop_analytics::DeveloperExceptionEventData;
-use drop_auth::{PublicKey, SecretKey};
+use drop_auth::{PublicKey, SecretKey, PUBLIC_KEY_LENGTH, SECRET_KEY_LENGTH};
 use drop_config::{Config, DropConfig, MooseConfig};
 use drop_storage::types::Transfer as TransferInfo;
 use drop_transfer::{auth, utils::Hidden, Event, FileToSend, OutgoingTransfer, Service, Transfer};
@@ -15,7 +15,7 @@ use tokio::{
     task::JoinHandle,
 };
 
-use crate::{event, TransferDescriptor};
+use crate::{event, KeyStore, TransferDescriptor};
 
 pub type Result<T = ()> = std::result::Result<T, crate::LibdropError>;
 
@@ -52,15 +52,11 @@ impl EventDispatcher {
 impl NordDropFFI {
     pub(super) fn new(
         event_cb: impl Fn(crate::Event) + Send + Sync + 'static,
-        pubkey_cb: impl Fn(IpAddr) -> Option<PublicKey> + Send + 'static,
-        privkey: SecretKey,
+        key_store: Arc<dyn KeyStore>,
         logger: Logger,
     ) -> Result<Self> {
         trace!(logger, "norddrop_new()");
 
-        // It's a debug print. Not visible in the production build
-        debug!(logger, "Private key: {:02X?}", privkey.to_bytes());
-
         Ok(NordDropFFI {
             instance: Arc::default(),
             logger: logger.clone(),
@@ -69,7 +65,7 @@ impl NordDropFFI {
                 cb: Arc::new(event_cb) as _,
             },
             config: DropConfig::default(),
-            keys: Arc::new(crate_key_context(logger, privkey, pubkey_cb)),
+            keys: Arc::new(create_key_context(logger, key_store)),
             #[cfg(unix)]
             fdresolv: None,
         })
@@ -522,22 +518,30 @@ impl NordDropFFI {
     }
 }
 
-fn crate_key_context(
-    logger: slog::Logger,
-    privkey: SecretKey,
-    pubkey_cb: impl Fn(IpAddr) -> Option<PublicKey> + Send + 'static,
-) -> auth::Context {
-    let pubkey_cb = std::sync::Mutex::new(pubkey_cb);
-    let public = move |ip: IpAddr| {
+fn create_key_context(logger: slog::Logger, key_store: Arc<dyn KeyStore>) -> auth::Context {
+    let privkey = {
+        let key_store = key_store.clone();
+        let logger = logger.clone();
+
+        move || {
+            let privkey: [u8; SECRET_KEY_LENGTH] = key_store.privkey().try_into().ok()?;
+            debug!(logger, "Retrieved private key: {:?}", privkey);
+            Some(SecretKey::from(privkey))
+        }
+    };
+
+    let pubkey_cb = std::sync::Mutex::new(key_store);
+    let pubkey = move |ip: IpAddr| {
         let guard = pubkey_cb.lock().expect("Failed to lock pubkey callback");
-        let key = guard(ip)?;
+        let pubkey = guard.on_pubkey(ip.to_string())?;
         drop(guard);
 
-        debug!(logger, "Public key for {ip:?}: {key:?}");
-        Some(key)
+        let pubkey: [u8; PUBLIC_KEY_LENGTH] = pubkey.try_into().ok()?;
+        debug!(logger, "Retrieved public key for: {} key: {:?}", ip, pubkey);
+        Some(PublicKey::from(pubkey))
     };
 
-    auth::Context::new(privkey, public)
+    auth::Context::new(privkey, pubkey)
 }
 
 fn open_database(

diff --git a/norddrop/src/uni.rs b/norddrop/src/uni.rs
@@ -1,6 +1,4 @@
-use std::sync::Mutex;
-
-use drop_auth::{PublicKey, SecretKey, PUBLIC_KEY_LENGTH};
+use std::sync::{Arc, Mutex};
 
 use crate::{device::NordDropFFI, Event, TransferDescriptor, TransferInfo};
 
@@ -36,20 +34,9 @@ impl NordDrop {
     ) -> Result<Self> {
         let logger = super::log::create(logger);
 
-        let privkey = key_store.privkey();
-        let privkey: [u8; PUBLIC_KEY_LENGTH] = privkey
-            .try_into()
-            .map_err(|_| crate::LibdropError::InvalidPrivkey)?;
-        let privkey = SecretKey::from(privkey);
-
         let dev = NordDropFFI::new(
             move |ev| event_callback.on_event(ev),
-            move |peer_ip| {
-                let pubkey = key_store.on_pubkey(peer_ip.to_string())?;
-                let pubkey: [u8; PUBLIC_KEY_LENGTH] = pubkey.try_into().ok()?;
-                Some(PublicKey::from(pubkey))
-            },
-            privkey,
+            key_store.into(),
             logger,
         )?;
 

diff --git a/test/drop_test/ffi.py b/test/drop_test/ffi.py
@@ -68,7 +68,7 @@ def __init__(self):
     def on_event(self, event: norddrop.Event):
         event = new_event(event)
         if DEBUG_PRINT_EVENT:
-            tprint(bcolors.HEADER + "--- event: ", event, bcolors.ENDC, flush=True)
+            tprint(bcolors.HEADER + "event: ", event, bcolors.ENDC, flush=True)
 
         with self._lock:
             self._events.append(event)

diff --git a/test/runner.py b/test/runner.py
@@ -4,6 +4,8 @@
 import os
 import re
 import time
+import sys
+from datetime import timedelta
 from threading import Semaphore
 from typing import Tuple
 
@@ -132,6 +134,7 @@ def run():
     for s in scenarios:
         total_containers += len(s.runners())
 
+    start_time = time.time()
     while True:
         if len(already_done) == len(scenarios):
             break
@@ -202,8 +205,6 @@ def run():
                     info = ContainerHolder(container, scenario.id(), TESTCASE_TIMEOUT)
                     scenario_results[scenario.id()].append(info)
 
-        curr_time = time.strftime("%H:%M:%S", time.localtime())
-
         done_containers = 0
         failed_container_count = 0
         for scenario in scenarios:
@@ -214,11 +215,11 @@ def run():
                         success, reason = container.success()
                         if not success:
                             failed_container_count += 1
-
-        print(
-            f"*** Test suite progress: {curr_time}: {done_containers}/{total_containers} containers finished, {failed_container_count} failed",
-            flush=True,
+        elapsed_time = str(timedelta(seconds=(round(time.time() - start_time))))
+        sys.stdout.write(
+            f"Testing in progress. Elapsed time: {elapsed_time}. Container exit stats: {done_containers}/{total_containers}, of those failed: {failed_container_count}\r"
         )
+        sys.stdout.flush()
 
         time.sleep(1)