Secret memory with memfd_secret (#321)

Implements:
- An additional allocator to use memfd_secret(2) and guard pages using mmap(2), implemented in quininer/memsec#16
- An allocator that abstracts away underlying allocators, and uses specified allocator set by rosenpass_secret_memory::policy functions (or a function that sets rosenpass_secret_memory::alloc::ALLOC_INIT
- Updates to tests- integration, fuzz, bench: some tests use procspawn to spawn multiple processes with different allocator policies

.github/workflows/qc.yaml

@@ -176,8 +176,12 @@ jobs:
           cargo fuzz run fuzz_handle_msg -- -max_total_time=5
           ulimit -s 8192000 && RUST_MIN_STACK=33554432000 && cargo fuzz run fuzz_kyber_encaps -- -max_total_time=5
           cargo fuzz run fuzz_mceliece_encaps -- -max_total_time=5
-          cargo fuzz run fuzz_box_secret_alloc -- -max_total_time=5
-          cargo fuzz run fuzz_vec_secret_alloc -- -max_total_time=5
+          cargo fuzz run fuzz_box_secret_alloc_malloc -- -max_total_time=5
+          cargo fuzz run fuzz_box_secret_alloc_memfdsec -- -max_total_time=5
+          cargo fuzz run fuzz_box_secret_alloc_memfdsec_mallocfb -- -max_total_time=5
+          cargo fuzz run fuzz_vec_secret_alloc_malloc -- -max_total_time=5
+          cargo fuzz run fuzz_vec_secret_alloc_memfdsec -- -max_total_time=5
+          cargo fuzz run fuzz_vec_secret_alloc_memfdsec_mallocfb -- -max_total_time=5
 
   codecov:
     runs-on: ubuntu-latest

Cargo.toml

@@ -12,15 +12,11 @@ members = [
     "fuzz",
     "secret-memory",
     "rp",
-    "wireguard-broker"
-]
-
-default-members = [
-    "rosenpass",
-    "rp",
     "wireguard-broker",
 ]
 
+default-members = ["rosenpass", "rp", "wireguard-broker"]
+
 [workspace.metadata.release]
 # ensure that adding `--package` as argument to `cargo release` still creates version tags in the form of `vx.y.z`
 tag-prefix = ""
@@ -45,7 +41,7 @@ env_logger = "0.10.2"
 toml = "0.7.8"
 static_assertions = "1.1.0"
 allocator-api2 = "0.2.14"
-memsec = "0.6.3"
+memsec = { version="0.7.0", features = [ "alloc_ext", ] }
 rand = "0.8.5"
 typenum = "1.17.0"
 log = { version = "0.4.21" }
@@ -54,9 +50,15 @@ serde = { version = "1.0.203", features = ["derive"] }
 arbitrary = { version = "1.3.2", features = ["derive"] }
 anyhow = { version = "1.0.86", features = ["backtrace", "std"] }
 mio = { version = "0.8.11", features = ["net", "os-poll"] }
-oqs-sys = { version = "0.9.1", default-features = false, features = ['classic_mceliece', 'kyber']  }
+oqs-sys = { version = "0.9.1", default-features = false, features = [
+    'classic_mceliece',
+    'kyber',
+] }
 blake2 = "0.10.6"
-chacha20poly1305 = { version = "0.10.1", default-features = false, features = [ "std", "heapless" ] }
+chacha20poly1305 = { version = "0.10.1", default-features = false, features = [
+    "std",
+    "heapless",
+] }
 zerocopy = { version = "0.7.34", features = ["derive"] }
 home = "0.5.9"
 derive_builder = "0.20.0"
@@ -65,14 +67,16 @@ postcard= {version = "1.0.8", features = ["alloc"]}
 
 #Dev dependencies
 serial_test = "3.1.1"
-tempfile="3"
+tempfile = "3"
 stacker = "0.1.15"
 libfuzzer-sys = "0.4"
 test_bin = "0.4.0"
 criterion = "0.4.0"
 allocator-api2-tests = "0.2.15"
+procspawn = {version = "1.0.0", features= ["test-support"]}
+
 
 #Broker dependencies (might need cleanup or changes)
 wireguard-uapi = "3.0.0"
 command-fds = "0.2.3"
-rustix = { version = "0.38.27", features = ["net"] }
+rustix = { version = "0.38.27", features = ["net"] }

fuzz/Cargo.toml

@@ -48,13 +48,37 @@ test = false
 doc = false
 
 [[bin]]
-name = "fuzz_box_secret_alloc"
-path = "fuzz_targets/box_secret_alloc.rs"
+name = "fuzz_box_secret_alloc_malloc"
+path = "fuzz_targets/box_secret_alloc_malloc.rs"
 test = false
 doc = false
 
 [[bin]]
-name = "fuzz_vec_secret_alloc"
-path = "fuzz_targets/vec_secret_alloc.rs"
+name = "fuzz_vec_secret_alloc_malloc"
+path = "fuzz_targets/vec_secret_alloc_malloc.rs"
 test = false
 doc = false
+
+[[bin]]
+name = "fuzz_box_secret_alloc_memfdsec"
+path = "fuzz_targets/box_secret_alloc_memfdsec.rs"
+test = false
+doc = false
+
+[[bin]]
+name = "fuzz_vec_secret_alloc_memfdsec"
+path = "fuzz_targets/vec_secret_alloc_memfdsec.rs"
+test = false
+doc = false
+
+[[bin]]
+name = "fuzz_box_secret_alloc_memfdsec_mallocfb"
+path = "fuzz_targets/box_secret_alloc_memfdsec_mallocfb.rs"
+test = false
+doc = false
+
+[[bin]]
+name = "fuzz_vec_secret_alloc_memfdsec_mallocfb"
+path = "fuzz_targets/vec_secret_alloc_memfdsec_mallocfb.rs"
+test = false
+doc = false

fuzz/fuzz_targets/box_secret_alloc_malloc.rs

@@ -0,0 +1,12 @@
+#![no_main]
+
+use libfuzzer_sys::fuzz_target;
+use rosenpass_secret_memory::alloc::secret_box;
+use rosenpass_secret_memory::policy::*;
+use std::sync::Once;
+static ONCE: Once = Once::new();
+
+fuzz_target!(|data: &[u8]| {
+    ONCE.call_once(secret_policy_use_only_malloc_secrets);
+    let _ = secret_box(data);
+});

fuzz/fuzz_targets/box_secret_alloc_memfdsec.rs

@@ -0,0 +1,13 @@
+#![no_main]
+
+use libfuzzer_sys::fuzz_target;
+use rosenpass_secret_memory::alloc::secret_box;
+use rosenpass_secret_memory::policy::*;
+use std::sync::Once;
+
+static ONCE: Once = Once::new();
+
+fuzz_target!(|data: &[u8]| {
+    ONCE.call_once(secret_policy_use_only_memfd_secrets);
+    let _ = secret_box(data);
+});

fuzz/fuzz_targets/box_secret_alloc_memfdsec_mallocfb.rs

@@ -2,7 +2,12 @@
 
 use libfuzzer_sys::fuzz_target;
 use rosenpass_secret_memory::alloc::secret_box;
+use rosenpass_secret_memory::policy::*;
+use std::sync::Once;
+
+static ONCE: Once = Once::new();
 
 fuzz_target!(|data: &[u8]| {
+    ONCE.call_once(secret_policy_try_use_memfd_secrets);
     let _ = secret_box(data);
 });

fuzz/fuzz_targets/handle_msg.rs

@@ -6,9 +6,13 @@ use libfuzzer_sys::fuzz_target;
 use rosenpass::protocol::CryptoServer;
 use rosenpass_cipher_traits::Kem;
 use rosenpass_ciphers::kem::StaticKem;
+use rosenpass_secret_memory::policy::*;
 use rosenpass_secret_memory::Secret;
+use std::sync::Once;
 
+static ONCE: Once = Once::new();
 fuzz_target!(|rx_buf: &[u8]| {
+    ONCE.call_once(secret_policy_use_only_malloc_secrets);
     let sk = Secret::from_slice(&[0; StaticKem::SK_LEN]);
     let pk = Secret::from_slice(&[0; StaticKem::PK_LEN]);
 

fuzz/fuzz_targets/vec_secret_alloc_malloc.rs

@@ -0,0 +1,15 @@
+#![no_main]
+
+use std::sync::Once;
+
+use libfuzzer_sys::fuzz_target;
+use rosenpass_secret_memory::alloc::secret_vec;
+use rosenpass_secret_memory::policy::*;
+
+static ONCE: Once = Once::new();
+
+fuzz_target!(|data: &[u8]| {
+    ONCE.call_once(secret_policy_use_only_malloc_secrets);
+    let mut vec = secret_vec();
+    vec.extend_from_slice(data);
+});

fuzz/fuzz_targets/vec_secret_alloc_memfdsec.rs

@@ -0,0 +1,15 @@
+#![no_main]
+
+use std::sync::Once;
+
+use libfuzzer_sys::fuzz_target;
+use rosenpass_secret_memory::alloc::secret_vec;
+use rosenpass_secret_memory::policy::*;
+
+static ONCE: Once = Once::new();
+
+fuzz_target!(|data: &[u8]| {
+    ONCE.call_once(secret_policy_use_only_memfd_secrets);
+    let mut vec = secret_vec();
+    vec.extend_from_slice(data);
+});

fuzz/fuzz_targets/vec_secret_alloc_memfdsec_mallocfb.rs

@@ -1,9 +1,15 @@
 #![no_main]
 
+use std::sync::Once;
+
 use libfuzzer_sys::fuzz_target;
 use rosenpass_secret_memory::alloc::secret_vec;
+use rosenpass_secret_memory::policy::*;
+
+static ONCE: Once = Once::new();
 
 fuzz_target!(|data: &[u8]| {
+    ONCE.call_once(secret_policy_try_use_memfd_secrets);
     let mut vec = secret_vec();
     vec.extend_from_slice(data);
 });

rosenpass/Cargo.toml

@@ -49,6 +49,7 @@ criterion = { workspace = true }
 test_bin = { workspace = true }
 stacker = { workspace = true }
 serial_test = {workspace = true}
+procspawn = {workspace = true}
 
 [features]
 enable_broker_api = ["rosenpass-wireguard-broker/enable_broker_api"]

rosenpass/benches/handshake.rs

@@ -5,6 +5,7 @@ use rosenpass_cipher_traits::Kem;
 use rosenpass_ciphers::kem::StaticKem;
 
 use criterion::{black_box, criterion_group, criterion_main, Criterion};
+use rosenpass_secret_memory::secret_policy_try_use_memfd_secrets;
 
 fn handle(
     tx: &mut CryptoServer,
@@ -56,6 +57,7 @@ fn make_server_pair() -> Result<(CryptoServer, CryptoServer)> {
 }
 
 fn criterion_benchmark(c: &mut Criterion) {
+    secret_policy_try_use_memfd_secrets();
     let (mut a, mut b) = make_server_pair().unwrap();
     c.bench_function("cca_secret_alloc", |bench| {
         bench.iter(|| {

rosenpass/src/cli.rs

@@ -3,6 +3,7 @@ use clap::{Parser, Subcommand};
 use rosenpass_cipher_traits::Kem;
 use rosenpass_ciphers::kem::StaticKem;
 use rosenpass_secret_memory::file::StoreSecret;
+use rosenpass_secret_memory::secret_policy_try_use_memfd_secrets;
 use rosenpass_util::file::{LoadValue, LoadValueB64};
 use rosenpass_wireguard_broker::brokers::native_unix::{
     NativeUnixBroker, NativeUnixBrokerConfigBaseBuilder, NativeUnixBrokerConfigBaseBuilderError,
@@ -154,6 +155,9 @@ impl CliCommand {
     /// ## TODO
     /// - This method consumes the [`CliCommand`] value. It might be wise to use a reference...
     pub fn run(self, test_helpers: Option<AppServerTest>) -> anyhow::Result<()> {
+        //Specify secret policy
+        secret_policy_try_use_memfd_secrets();
+
         use CliCommand::*;
         match self {
             Man => {

rosenpass/src/protocol.rs

@@ -19,12 +19,16 @@
 //! [CryptoServer].
 //!
 //! ```
+//! use rosenpass_secret_memory::policy::*;
 //! use rosenpass_cipher_traits::Kem;
 //! use rosenpass_ciphers::kem::StaticKem;
 //! use rosenpass::{
 //!     protocol::{SSk, SPk, MsgBuf, PeerPtr, CryptoServer, SymKey},
 //! };
 //! # fn main() -> anyhow::Result<()> {
+//! // Set security policy for storing secrets
+//!
+//! secret_policy_try_use_memfd_secrets();
 //!
 //! // initialize secret and public key for peer a ...
 //! let (mut peer_a_sk, mut peer_a_pk) = (SSk::zero(), SPk::zero());
@@ -2145,6 +2149,7 @@ mod test {
     use std::{net::SocketAddrV4, thread::sleep, time::Duration};
 
     use super::*;
+    use serial_test::serial;
 
     struct VecHostIdentifier(Vec<u8>);
 
@@ -2166,7 +2171,21 @@ mod test {
         }
     }
 
+    fn setup_logging() {
+        use std::io::Write;
+        let mut log_builder = env_logger::Builder::from_default_env(); // sets log level filter from environment (or defaults)
+        log_builder.filter_level(log::LevelFilter::Info);
+        log_builder.format_timestamp_nanos();
+        log_builder.format(|buf, record| {
+            let ts_format = buf.timestamp_nanos().to_string();
+            writeln!(buf, "{}: {}", &ts_format[14..], record.args())
+        });
+
+        let _ = log_builder.try_init();
+    }
+
     #[test]
+    #[serial]
     /// Ensure that the protocol implementation can deal with truncated
     /// messages and with overlong messages.
     ///
@@ -2182,6 +2201,8 @@ mod test {
     /// Through all this, the handshake should still successfully terminate;
     /// i.e. an exchanged key must be produced in both servers.
     fn handles_incorrect_size_messages() {
+        setup_logging();
+        rosenpass_secret_memory::secret_policy_try_use_memfd_secrets();
         stacker::grow(8 * 1024 * 1024, || {
             const OVERSIZED_MESSAGE: usize = ((MAX_MESSAGE_LEN as f32) * 1.2) as usize;
             type MsgBufPlus = Public<OVERSIZED_MESSAGE>;
@@ -2252,7 +2273,10 @@ mod test {
     }
 
     #[test]
+    #[serial]
     fn test_regular_exchange() {
+        setup_logging();
+        rosenpass_secret_memory::secret_policy_try_use_memfd_secrets();
         stacker::grow(8 * 1024 * 1024, || {
             type MsgBufPlus = Public<MAX_MESSAGE_LEN>;
             let (mut a, mut b) = make_server_pair().unwrap();
@@ -2296,7 +2320,7 @@ mod test {
 
             //B handles InitConf, sends EmptyData
             let HandleMsgResult {
-                resp,
+                resp: _,
                 exchanged_with,
             } = b
                 .handle_msg(&a_to_b_buf.as_slice()[..init_conf_len], &mut *b_to_a_buf)
@@ -2310,7 +2334,10 @@ mod test {
     }
 
     #[test]
+    #[serial]
     fn test_regular_init_conf_retransmit() {
+        setup_logging();
+        rosenpass_secret_memory::secret_policy_try_use_memfd_secrets();
         stacker::grow(8 * 1024 * 1024, || {
             type MsgBufPlus = Public<MAX_MESSAGE_LEN>;
             let (mut a, mut b) = make_server_pair().unwrap();
@@ -2355,7 +2382,7 @@ mod test {
 
             //B handles InitConf, sends EmptyData
             let HandleMsgResult {
-                resp,
+                resp: _,
                 exchanged_with,
             } = b
                 .handle_msg(&a_to_b_buf.as_slice()[..init_conf_len], &mut *b_to_a_buf)
@@ -2368,7 +2395,7 @@ mod test {
 
             //B handles InitConf again, sends EmptyData
             let HandleMsgResult {
-                resp,
+                resp: _,
                 exchanged_with,
             } = b
                 .handle_msg(&a_to_b_buf.as_slice()[..init_conf_len], &mut *b_to_a_buf)
@@ -2382,7 +2409,10 @@ mod test {
     }
 
     #[test]
+    #[serial]
     fn cookie_reply_mechanism_responder_under_load() {
+        setup_logging();
+        rosenpass_secret_memory::secret_policy_try_use_memfd_secrets();
         stacker::grow(8 * 1024 * 1024, || {
             type MsgBufPlus = Public<MAX_MESSAGE_LEN>;
             let (mut a, mut b) = make_server_pair().unwrap();
@@ -2476,7 +2506,10 @@ mod test {
     }
 
     #[test]
+    #[serial]
     fn cookie_reply_mechanism_initiator_bails_on_message_under_load() {
+        setup_logging();
+        rosenpass_secret_memory::secret_policy_try_use_memfd_secrets();
         stacker::grow(8 * 1024 * 1024, || {
             type MsgBufPlus = Public<MAX_MESSAGE_LEN>;
             let (mut a, mut b) = make_server_pair().unwrap();

rosenpass/tests/integration_test.rs

@@ -6,7 +6,7 @@ use std::{
     time::Duration,
 };
 
-use clap::{builder::Str, Parser};
+use clap::Parser;
 use rosenpass::{app_server::AppServerTestBuilder, cli::CliArgs};
 use rosenpass_secret_memory::{Public, Secret};
 use rosenpass_wireguard_broker::{WireguardBrokerMio, WG_KEY_LEN, WG_PEER_LEN};
@@ -275,6 +275,7 @@ fn check_exchange_under_dos() {
     fs::remove_dir_all(&tmpdir).unwrap();
 }
 
+#[allow(dead_code)]
 #[derive(Debug, Default)]
 struct MockBrokerInner {
     psk: Option<Secret<WG_KEY_LEN>>,

rp/src/key.rs

@@ -102,6 +102,7 @@ mod tests {
     use std::fs;
 
     use rosenpass::protocol::{SPk, SSk};
+    use rosenpass_secret_memory::secret_policy_try_use_memfd_secrets;
     use rosenpass_secret_memory::Secret;
     use rosenpass_util::file::LoadValue;
     use rosenpass_util::file::LoadValueB64;
@@ -110,7 +111,8 @@ mod tests {
     use crate::key::{genkey, pubkey, WG_B64_LEN};
 
     #[test]
-    fn it_works() {
+    fn test_key_loopback() {
+        secret_policy_try_use_memfd_secrets();
         let private_keys_dir = tempdir().unwrap();
         fs::remove_dir(private_keys_dir.path()).unwrap();
 

rp/src/main.rs

@@ -3,6 +3,7 @@ use std::process::exit;
 use cli::{Cli, Command};
 use exchange::exchange;
 use key::{genkey, pubkey};
+use rosenpass_secret_memory::policy;
 
 mod cli;
 mod exchange;
@@ -10,6 +11,8 @@ mod key;
 
 #[tokio::main]
 async fn main() {
+    policy::secret_policy_try_use_memfd_secrets();
+
     let cli = match Cli::parse(std::env::args().peekable()) {
         Ok(cli) => cli,
         Err(err) => {

secret-memory/Cargo.toml

@@ -23,3 +23,4 @@ log = { workspace = true }
 allocator-api2-tests = { workspace = true }
 tempfile = {workspace = true}
 base64ct = {workspace = true}
+procspawn = {workspace = true}

secret-memory/src/alloc/memsec/malloc.rs

@@ -4,37 +4,41 @@ use std::ptr::NonNull;
 use allocator_api2::alloc::{AllocError, Allocator, Layout};
 
 #[derive(Copy, Clone, Default)]
-struct MemsecAllocatorContents;
+struct MallocAllocatorContents;
 
 /// Memory allocation using using the memsec crate
 #[derive(Copy, Clone, Default)]
-pub struct MemsecAllocator {
-    _dummy_private_data: MemsecAllocatorContents,
+pub struct MallocAllocator {
+    _dummy_private_data: MallocAllocatorContents,
 }
 
 /// A box backed by the memsec allocator
-pub type MemsecBox<T> = allocator_api2::boxed::Box<T, MemsecAllocator>;
+pub type MallocBox<T> = allocator_api2::boxed::Box<T, MallocAllocator>;
 
 /// A vector backed by the memsec allocator
-pub type MemsecVec<T> = allocator_api2::vec::Vec<T, MemsecAllocator>;
+pub type MallocVec<T> = allocator_api2::vec::Vec<T, MallocAllocator>;
 
-pub fn memsec_box<T>(x: T) -> MemsecBox<T> {
-    MemsecBox::<T>::new_in(x, MemsecAllocator::new())
+pub fn malloc_box_try<T>(x: T) -> Result<MallocBox<T>, AllocError> {
+    MallocBox::<T>::try_new_in(x, MallocAllocator::new())
 }
 
-pub fn memsec_vec<T>() -> MemsecVec<T> {
-    MemsecVec::<T>::new_in(MemsecAllocator::new())
+pub fn malloc_box<T>(x: T) -> MallocBox<T> {
+    MallocBox::<T>::new_in(x, MallocAllocator::new())
 }
 
-impl MemsecAllocator {
+pub fn malloc_vec<T>() -> MallocVec<T> {
+    MallocVec::<T>::new_in(MallocAllocator::new())
+}
+
+impl MallocAllocator {
     pub fn new() -> Self {
         Self {
-            _dummy_private_data: MemsecAllocatorContents,
+            _dummy_private_data: MallocAllocatorContents,
         }
     }
 }
 
-unsafe impl Allocator for MemsecAllocator {
+unsafe impl Allocator for MallocAllocator {
     fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
         // Call memsec allocator
         let mem: Option<NonNull<[u8]>> = unsafe { memsec::malloc_sized(layout.size()) };
@@ -48,8 +52,8 @@ unsafe impl Allocator for MemsecAllocator {
         // Ensure the right alignment is used
         let off = (mem.as_ptr() as *const u8).align_offset(layout.align());
         if off != 0 {
-            log::error!("Allocation {layout:?} was requested but memsec returned allocation \
-                with offset {off} from the requested alignment. Memsec always allocates values \
+            log::error!("Allocation {layout:?} was requested but malloc-based memsec returned allocation \
+                with offset {off} from the requested alignment. Malloc always allocates values \
                 at the end of a memory page for security reasons, custom alignments are not supported. \
                 You could try allocating an oversized value.");
             unsafe { memsec::free(mem) };
@@ -66,7 +70,7 @@ unsafe impl Allocator for MemsecAllocator {
     }
 }
 
-impl fmt::Debug for MemsecAllocator {
+impl fmt::Debug for MallocAllocator {
     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
         fmt.write_str("<memsec based Rust allocator>")
     }
@@ -78,21 +82,21 @@ mod test {
 
     use super::*;
 
-    make_test! { test_sizes(MemsecAllocator::new()) }
-    make_test! { test_vec(MemsecAllocator::new()) }
-    make_test! { test_many_boxes(MemsecAllocator::new()) }
+    make_test! { test_sizes(MallocAllocator::new()) }
+    make_test! { test_vec(MallocAllocator::new()) }
+    make_test! { test_many_boxes(MallocAllocator::new()) }
 
     #[test]
-    fn memsec_allocation() {
-        let alloc = MemsecAllocator::new();
-        memsec_allocation_impl::<0>(&alloc);
-        memsec_allocation_impl::<7>(&alloc);
-        memsec_allocation_impl::<8>(&alloc);
-        memsec_allocation_impl::<64>(&alloc);
-        memsec_allocation_impl::<999>(&alloc);
+    fn malloc_allocation() {
+        let alloc = MallocAllocator::new();
+        malloc_allocation_impl::<0>(&alloc);
+        malloc_allocation_impl::<7>(&alloc);
+        malloc_allocation_impl::<8>(&alloc);
+        malloc_allocation_impl::<64>(&alloc);
+        malloc_allocation_impl::<999>(&alloc);
     }
 
-    fn memsec_allocation_impl<const N: usize>(alloc: &MemsecAllocator) {
+    fn malloc_allocation_impl<const N: usize>(alloc: &MallocAllocator) {
         let layout = Layout::new::<[u8; N]>();
         let mem = alloc.allocate(layout).unwrap();
 

secret-memory/src/alloc/memsec/memfdsec.rs

@@ -0,0 +1,112 @@
+#![cfg(target_os = "linux")]
+use std::fmt;
+use std::ptr::NonNull;
+
+use allocator_api2::alloc::{AllocError, Allocator, Layout};
+
+#[derive(Copy, Clone, Default)]
+struct MemfdSecAllocatorContents;
+
+/// Memory allocation using using the memsec crate
+#[derive(Copy, Clone, Default)]
+pub struct MemfdSecAllocator {
+    _dummy_private_data: MemfdSecAllocatorContents,
+}
+
+/// A box backed by the memsec allocator
+pub type MemfdSecBox<T> = allocator_api2::boxed::Box<T, MemfdSecAllocator>;
+
+/// A vector backed by the memsec allocator
+pub type MemfdSecVec<T> = allocator_api2::vec::Vec<T, MemfdSecAllocator>;
+
+pub fn memfdsec_box_try<T>(x: T) -> Result<MemfdSecBox<T>, AllocError> {
+    MemfdSecBox::<T>::try_new_in(x, MemfdSecAllocator::new())
+}
+
+pub fn memfdsec_box<T>(x: T) -> MemfdSecBox<T> {
+    MemfdSecBox::<T>::new_in(x, MemfdSecAllocator::new())
+}
+
+pub fn memfdsec_vec<T>() -> MemfdSecVec<T> {
+    MemfdSecVec::<T>::new_in(MemfdSecAllocator::new())
+}
+
+impl MemfdSecAllocator {
+    pub fn new() -> Self {
+        Self {
+            _dummy_private_data: MemfdSecAllocatorContents,
+        }
+    }
+}
+
+unsafe impl Allocator for MemfdSecAllocator {
+    fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
+        // Call memsec allocator
+        let mem: Option<NonNull<[u8]>> = unsafe { memsec::memfd_secret_sized(layout.size()) };
+
+        // Unwrap the option
+        let Some(mem) = mem else {
+            log::error!("Allocation {layout:?} was requested but memfd-based memsec returned a null pointer");
+            return Err(AllocError);
+        };
+
+        // Ensure the right alignment is used
+        let off = (mem.as_ptr() as *const u8).align_offset(layout.align());
+        if off != 0 {
+            log::error!("Allocation {layout:?} was requested but memfd-based memsec returned allocation \
+                with offset {off} from the requested alignment. Memfd always allocates values \
+                at the end of a memory page for security reasons, custom alignments are not supported. \
+                You could try allocating an oversized value.");
+            unsafe { memsec::free_memfd_secret(mem) };
+            return Err(AllocError);
+        };
+
+        Ok(mem)
+    }
+
+    unsafe fn deallocate(&self, ptr: NonNull<u8>, _layout: Layout) {
+        unsafe {
+            memsec::free_memfd_secret(ptr);
+        }
+    }
+}
+
+impl fmt::Debug for MemfdSecAllocator {
+    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
+        fmt.write_str("<memsec based Rust allocator>")
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use allocator_api2_tests::make_test;
+
+    use super::*;
+
+    make_test! { test_sizes(MemfdSecAllocator::new()) }
+    make_test! { test_vec(MemfdSecAllocator::new()) }
+    make_test! { test_many_boxes(MemfdSecAllocator::new()) }
+
+    #[test]
+    fn memfdsec_allocation() {
+        let alloc = MemfdSecAllocator::new();
+        memfdsec_allocation_impl::<0>(&alloc);
+        memfdsec_allocation_impl::<7>(&alloc);
+        memfdsec_allocation_impl::<8>(&alloc);
+        memfdsec_allocation_impl::<64>(&alloc);
+        memfdsec_allocation_impl::<999>(&alloc);
+    }
+
+    fn memfdsec_allocation_impl<const N: usize>(alloc: &MemfdSecAllocator) {
+        let layout = Layout::new::<[u8; N]>();
+        let mem = alloc.allocate(layout).unwrap();
+
+        // https://libsodium.gitbook.io/doc/memory_management#guarded-heap-allocations
+        // promises us that allocated memory is initialized with the magic byte 0xDB
+        // and memsec promises to provide a reimplementation of the libsodium mechanism;
+        // it uses the magic value 0xD0 though
+        assert_eq!(unsafe { mem.as_ref() }, &[0xD0u8; N]);
+        let mem = NonNull::new(mem.as_ptr() as *mut u8).unwrap();
+        unsafe { alloc.deallocate(mem, layout) };
+    }
+}

secret-memory/src/alloc/memsec/mod.rs

@@ -0,0 +1,2 @@
+pub mod malloc;
+pub mod memfdsec;

secret-memory/src/alloc/mod.rs

@@ -1,6 +1,86 @@
 pub mod memsec;
 
-pub use crate::alloc::memsec::{
-    memsec_box as secret_box, memsec_vec as secret_vec, MemsecAllocator as SecretAllocator,
-    MemsecBox as SecretBox, MemsecVec as SecretVec,
-};
+use std::sync::OnceLock;
+
+use allocator_api2::alloc::{AllocError, Allocator};
+use memsec::malloc::MallocAllocator;
+
+#[cfg(target_os = "linux")]
+use memsec::memfdsec::MemfdSecAllocator;
+
+static ALLOC_TYPE: OnceLock<SecretAllocType> = OnceLock::new();
+
+/// Sets the secret allocation type to use.
+/// Intended usage at startup before secret allocation
+/// takes place
+pub fn set_secret_alloc_type(alloc_type: SecretAllocType) {
+    ALLOC_TYPE.set(alloc_type).unwrap();
+}
+
+pub fn get_or_init_secret_alloc_type(alloc_type: SecretAllocType) -> SecretAllocType {
+    *ALLOC_TYPE.get_or_init(|| alloc_type)
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum SecretAllocType {
+    MemsecMalloc,
+    #[cfg(target_os = "linux")]
+    MemsecMemfdSec,
+}
+
+pub struct SecretAlloc {
+    alloc_type: SecretAllocType,
+}
+
+impl Default for SecretAlloc {
+    fn default() -> Self {
+        Self {
+            alloc_type: *ALLOC_TYPE.get().expect(
+                "Secret security policy not specified. \
+                Run the specifying policy function in \
+                 rosenpass_secret_memory::policy or set a \
+                 custom policy by initializing \
+                 rosenpass_secret_memory::alloc::ALLOC_TYPE \
+                 before using secrets",
+            ),
+        }
+    }
+}
+
+unsafe impl Allocator for SecretAlloc {
+    fn allocate(
+        &self,
+        layout: std::alloc::Layout,
+    ) -> Result<std::ptr::NonNull<[u8]>, allocator_api2::alloc::AllocError> {
+        match self.alloc_type {
+            SecretAllocType::MemsecMalloc => MallocAllocator::default().allocate(layout),
+            #[cfg(target_os = "linux")]
+            SecretAllocType::MemsecMemfdSec => MemfdSecAllocator::default().allocate(layout),
+        }
+    }
+
+    unsafe fn deallocate(&self, ptr: std::ptr::NonNull<u8>, layout: std::alloc::Layout) {
+        match self.alloc_type {
+            SecretAllocType::MemsecMalloc => MallocAllocator::default().deallocate(ptr, layout),
+            #[cfg(target_os = "linux")]
+            SecretAllocType::MemsecMemfdSec => MemfdSecAllocator::default().deallocate(ptr, layout),
+        }
+    }
+}
+
+pub type SecretBox<T> = allocator_api2::boxed::Box<T, SecretAlloc>;
+
+/// A vector backed by the memsec allocator
+pub type SecretVec<T> = allocator_api2::vec::Vec<T, SecretAlloc>;
+
+pub fn secret_box_try<T>(x: T) -> Result<SecretBox<T>, AllocError> {
+    SecretBox::<T>::try_new_in(x, SecretAlloc::default())
+}
+
+pub fn secret_box<T>(x: T) -> SecretBox<T> {
+    SecretBox::<T>::new_in(x, SecretAlloc::default())
+}
+
+pub fn secret_vec<T>() -> SecretVec<T> {
+    SecretVec::<T>::new_in(SecretAlloc::default())
+}

secret-memory/src/lib.rs

@@ -9,3 +9,6 @@ pub use crate::public::Public;
 
 mod secret;
 pub use crate::secret::Secret;
+
+pub mod policy;
+pub use crate::policy::*;

secret-memory/src/policy.rs

@@ -0,0 +1,82 @@
+pub fn secret_policy_try_use_memfd_secrets() {
+    let alloc_type = {
+        #[cfg(target_os = "linux")]
+        {
+            if crate::alloc::memsec::memfdsec::memfdsec_box_try(0u8).is_ok() {
+                crate::alloc::SecretAllocType::MemsecMemfdSec
+            } else {
+                crate::alloc::SecretAllocType::MemsecMalloc
+            }
+        }
+
+        #[cfg(not(target_os = "linux"))]
+        {
+            crate::alloc::SecretAllocType::MemsecMalloc
+        }
+    };
+    assert_eq!(
+        alloc_type,
+        crate::alloc::get_or_init_secret_alloc_type(alloc_type)
+    );
+
+    log::info!("Secrets will be allocated using {:?}", alloc_type);
+}
+
+#[cfg(target_os = "linux")]
+pub fn secret_policy_use_only_memfd_secrets() {
+    let alloc_type = crate::alloc::SecretAllocType::MemsecMemfdSec;
+
+    assert_eq!(
+        alloc_type,
+        crate::alloc::get_or_init_secret_alloc_type(alloc_type)
+    );
+
+    log::info!("Secrets will be allocated using {:?}", alloc_type);
+}
+
+pub fn secret_policy_use_only_malloc_secrets() {
+    let alloc_type = crate::alloc::SecretAllocType::MemsecMalloc;
+    assert_eq!(
+        alloc_type,
+        crate::alloc::get_or_init_secret_alloc_type(alloc_type)
+    );
+
+    log::info!("Secrets will be allocated using {:?}", alloc_type);
+}
+
+pub mod test {
+    #[macro_export]
+    macro_rules! test_spawn_process_with_policies {
+        ($body:block, $($f: expr),*) => {
+            $(
+                let handle = procspawn::spawn((), |_| {
+
+                $f();
+
+                $body
+
+                });
+                handle.join().unwrap();
+            )*
+            };
+        }
+
+    #[macro_export]
+    macro_rules! test_spawn_process_provided_policies {
+        ($body: block) => {
+            $crate::test_spawn_process_with_policies!(
+                $body,
+                $crate::policy::secret_policy_try_use_memfd_secrets,
+                $crate::secret_policy_use_only_malloc_secrets
+            );
+
+            #[cfg(target_os = "linux")]
+            {
+                $crate::test_spawn_process_with_policies!(
+                    $body,
+                    $crate::policy::secret_policy_use_only_memfd_secrets
+                );
+            }
+        };
+    }
+}

secret-memory/src/secret.rs

@@ -321,133 +321,147 @@ impl<const N: usize> StoreSecret for Secret<N> {
 
 #[cfg(test)]
 mod test {
+    use crate::test_spawn_process_provided_policies;
+
     use super::*;
     use std::{fs, os::unix::fs::PermissionsExt};
     use tempfile::tempdir;
 
+    procspawn::enable_test_support!();
+
     /// check that we can alloc using the magic pool
     #[test]
     fn secret_memory_pool_take() {
-        const N: usize = 0x100;
-        let mut pool = SecretMemoryPool::new();
-        let secret: ZeroizingSecretBox<[u8; N]> = pool.take();
-        assert_eq!(secret.as_ref(), &[0; N]);
+        test_spawn_process_provided_policies!({
+            const N: usize = 0x100;
+            let mut pool = SecretMemoryPool::new();
+            let secret: ZeroizingSecretBox<[u8; N]> = pool.take();
+            assert_eq!(secret.as_ref(), &[0; N]);
+        });
     }
 
     /// check that a secret lives, even if its [SecretMemoryPool] is deleted
     #[test]
     fn secret_memory_pool_drop() {
-        const N: usize = 0x100;
-        let mut pool = SecretMemoryPool::new();
-        let secret: ZeroizingSecretBox<[u8; N]> = pool.take();
-        std::mem::drop(pool);
-        assert_eq!(secret.as_ref(), &[0; N]);
+        test_spawn_process_provided_policies!({
+            const N: usize = 0x100;
+            let mut pool = SecretMemoryPool::new();
+            let secret: ZeroizingSecretBox<[u8; N]> = pool.take();
+            std::mem::drop(pool);
+            assert_eq!(secret.as_ref(), &[0; N]);
+        });
     }
 
     /// check that a secret can be reborn, freshly initialized with zero
     #[test]
     fn secret_memory_pool_release() {
-        const N: usize = 1;
-        let mut pool = SecretMemoryPool::new();
-        let mut secret: ZeroizingSecretBox<[u8; N]> = pool.take();
-        let old_secret_ptr = secret.as_ref().as_ptr();
+        test_spawn_process_provided_policies!({
+            const N: usize = 1;
+            let mut pool = SecretMemoryPool::new();
+            let mut secret: ZeroizingSecretBox<[u8; N]> = pool.take();
+            let old_secret_ptr = secret.as_ref().as_ptr();
 
-        secret.as_mut()[0] = 0x13;
-        pool.release(secret);
+            secret.as_mut()[0] = 0x13;
+            pool.release(secret);
 
-        // now check that we get the same ptr
-        let new_secret: ZeroizingSecretBox<[u8; N]> = pool.take();
-        assert_eq!(old_secret_ptr, new_secret.as_ref().as_ptr());
+            // now check that we get the same ptr
+            let new_secret: ZeroizingSecretBox<[u8; N]> = pool.take();
+            assert_eq!(old_secret_ptr, new_secret.as_ref().as_ptr());
 
-        // and that the secret was zeroized
-        assert_eq!(new_secret.as_ref(), &[0; N]);
+            // and that the secret was zeroized
+            assert_eq!(new_secret.as_ref(), &[0; N]);
+        });
     }
 
     /// test loading a secret from an example file, and then storing it again in a different file
     #[test]
     fn test_secret_load_store() {
-        const N: usize = 100;
+        test_spawn_process_provided_policies!({
+            const N: usize = 100;
 
-        // Generate original random bytes
-        let original_bytes: [u8; N] = [rand::random(); N];
+            // Generate original random bytes
+            let original_bytes: [u8; N] = [rand::random(); N];
 
-        // Create a temporary directory
-        let temp_dir = tempdir().unwrap();
+            // Create a temporary directory
+            let temp_dir = tempdir().unwrap();
 
-        // Store the original secret to an example file in the temporary directory
-        let example_file = temp_dir.path().join("example_file");
-        std::fs::write(example_file.clone(), &original_bytes).unwrap();
+            // Store the original secret to an example file in the temporary directory
+            let example_file = temp_dir.path().join("example_file");
+            std::fs::write(example_file.clone(), &original_bytes).unwrap();
 
-        // Load the secret from the example file
-        let loaded_secret = Secret::load(&example_file).unwrap();
+            // Load the secret from the example file
+            let loaded_secret = Secret::load(&example_file).unwrap();
 
-        // Check that the loaded secret matches the original bytes
-        assert_eq!(loaded_secret.secret(), &original_bytes);
+            // Check that the loaded secret matches the original bytes
+            assert_eq!(loaded_secret.secret(), &original_bytes);
 
-        // Store the loaded secret to a different file in the temporary directory
-        let new_file = temp_dir.path().join("new_file");
-        loaded_secret.store(&new_file).unwrap();
+            // Store the loaded secret to a different file in the temporary directory
+            let new_file = temp_dir.path().join("new_file");
+            loaded_secret.store(&new_file).unwrap();
 
-        // Read the contents of the new file
-        let new_file_contents = fs::read(&new_file).unwrap();
+            // Read the contents of the new file
+            let new_file_contents = fs::read(&new_file).unwrap();
 
-        // Read the contents of the original file
-        let original_file_contents = fs::read(&example_file).unwrap();
+            // Read the contents of the original file
+            let original_file_contents = fs::read(&example_file).unwrap();
 
-        // Check that the contents of the new file match the original file
-        assert_eq!(new_file_contents, original_file_contents);
+            // Check that the contents of the new file match the original file
+            assert_eq!(new_file_contents, original_file_contents);
+        });
     }
 
     /// test loading a base64 encoded secret from an example file, and then storing it again in a different file
     #[test]
     fn test_secret_load_store_base64() {
-        const N: usize = 100;
-        // Generate original random bytes
-        let original_bytes: [u8; N] = [rand::random(); N];
-        // Create a temporary directory
-        let temp_dir = tempdir().unwrap();
-        let example_file = temp_dir.path().join("example_file");
-        let mut encoded_secret = [0u8; N * 2];
-        let encoded_secret = b64_encode(&original_bytes, &mut encoded_secret).unwrap();
+        test_spawn_process_provided_policies!({
+            const N: usize = 100;
+            // Generate original random bytes
+            let original_bytes: [u8; N] = [rand::random(); N];
+            // Create a temporary directory
+            let temp_dir = tempdir().unwrap();
+            let example_file = temp_dir.path().join("example_file");
+            let mut encoded_secret = [0u8; N * 2];
+            let encoded_secret = b64_encode(&original_bytes, &mut encoded_secret).unwrap();
 
-        std::fs::write(&example_file, encoded_secret).unwrap();
+            std::fs::write(&example_file, encoded_secret).unwrap();
 
-        // Load the secret from the example file
-        let loaded_secret = Secret::load_b64::<{ N * 2 }, _>(&example_file).unwrap();
-        // Check that the loaded secret matches the original bytes
-        assert_eq!(loaded_secret.secret(), &original_bytes);
+            // Load the secret from the example file
+            let loaded_secret = Secret::load_b64::<{ N * 2 }, _>(&example_file).unwrap();
+            // Check that the loaded secret matches the original bytes
+            assert_eq!(loaded_secret.secret(), &original_bytes);
 
-        // Store the loaded secret to a different file in the temporary directory
-        let new_file = temp_dir.path().join("new_file");
-        loaded_secret.store_b64::<{ N * 2 }, _>(&new_file).unwrap();
+            // Store the loaded secret to a different file in the temporary directory
+            let new_file = temp_dir.path().join("new_file");
+            loaded_secret.store_b64::<{ N * 2 }, _>(&new_file).unwrap();
 
-        // Read the contents of the new file
-        let new_file_contents = fs::read(&new_file).unwrap();
-        // Read the contents of the original file
-        let original_file_contents = fs::read(&example_file).unwrap();
-        // Check that the contents of the new file match the original file
-        assert_eq!(new_file_contents, original_file_contents);
+            // Read the contents of the new file
+            let new_file_contents = fs::read(&new_file).unwrap();
+            // Read the contents of the original file
+            let original_file_contents = fs::read(&example_file).unwrap();
+            // Check that the contents of the new file match the original file
+            assert_eq!(new_file_contents, original_file_contents);
 
-        //Check new file permissions are secret
-        let metadata = fs::metadata(&new_file).unwrap();
-        assert_eq!(metadata.permissions().mode() & 0o000777, 0o600);
+            //Check new file permissions are secret
+            let metadata = fs::metadata(&new_file).unwrap();
+            assert_eq!(metadata.permissions().mode() & 0o000777, 0o600);
 
-        // Store the loaded secret to a different file in the temporary directory for a second time
-        let new_file = temp_dir.path().join("new_file_writer");
-        let new_file_writer = fopen_w(new_file.clone(), Visibility::Secret).unwrap();
-        loaded_secret
-            .store_b64_writer::<{ N * 2 }, _>(&new_file_writer)
-            .unwrap();
+            // Store the loaded secret to a different file in the temporary directory for a second time
+            let new_file = temp_dir.path().join("new_file_writer");
+            let new_file_writer = fopen_w(new_file.clone(), Visibility::Secret).unwrap();
+            loaded_secret
+                .store_b64_writer::<{ N * 2 }, _>(&new_file_writer)
+                .unwrap();
 
-        // Read the contents of the new file
-        let new_file_contents = fs::read(&new_file).unwrap();
-        // Read the contents of the original file
-        let original_file_contents = fs::read(&example_file).unwrap();
-        // Check that the contents of the new file match the original file
-        assert_eq!(new_file_contents, original_file_contents);
+            // Read the contents of the new file
+            let new_file_contents = fs::read(&new_file).unwrap();
+            // Read the contents of the original file
+            let original_file_contents = fs::read(&example_file).unwrap();
+            // Check that the contents of the new file match the original file
+            assert_eq!(new_file_contents, original_file_contents);
 
-        //Check new file permissions are secret
-        let metadata = fs::metadata(&new_file).unwrap();
-        assert_eq!(metadata.permissions().mode() & 0o000777, 0o600);
+            //Check new file permissions are secret
+            let metadata = fs::metadata(&new_file).unwrap();
+            assert_eq!(metadata.permissions().mode() & 0o000777, 0o600);
+        });
     }
 }

wireguard-broker/Cargo.toml

@@ -33,6 +33,7 @@ rosenpass-util = { workspace = true }
 
 [dev-dependencies]
 rand = {workspace = true}
+procspawn = {workspace = true}
 
 [features]
 enable_broker_api=[]

wireguard-broker/src/brokers/netlink.rs

@@ -77,7 +77,7 @@ impl WireGuardBroker for NetlinkWireGuardBroker {
     fn set_psk(&mut self, config: SerializedBrokerConfig) -> Result<(), Self::Error> {
         let config: NetworkBrokerConfig = config
             .try_into()
-            .map_err(|e| SetPskError::NoSuchInterface)?;
+            .map_err(|_e| SetPskError::NoSuchInterface)?;
         // Ensure that the peer exists by querying the device configuration
         // TODO: Use InvalidInterfaceError
 

wireguard-broker/tests/integration.rs

@@ -53,82 +53,88 @@ mod integration_tests {
         }
     }
 
+    procspawn::enable_test_support!();
+
     #[test]
     fn test_psk_exchanges() {
         const TEST_RUNS: usize = 100;
 
-        let server_broker_inner = Arc::new(Mutex::new(MockServerBrokerInner::default()));
-        // Create a mock BrokerServer
-        let server_broker = MockServerBroker::new(server_broker_inner.clone());
+        use rosenpass_secret_memory::test_spawn_process_provided_policies;
 
-        let mut server = BrokerServer::<SetPskError, MockServerBroker>::new(server_broker);
+        test_spawn_process_provided_policies!({
+            let server_broker_inner = Arc::new(Mutex::new(MockServerBrokerInner::default()));
+            // Create a mock BrokerServer
+            let server_broker = MockServerBroker::new(server_broker_inner.clone());
 
-        let (client_socket, mut server_socket) = mio::net::UnixStream::pair().unwrap();
+            let mut server = BrokerServer::<SetPskError, MockServerBroker>::new(server_broker);
+
+            let (client_socket, mut server_socket) = mio::net::UnixStream::pair().unwrap();
+
+            // Spawn a new thread to connect to the unix socket
+            let handle = std::thread::spawn(move || {
+                for _ in 0..TEST_RUNS {
+                    // Wait for 8 bytes of length to come in
+                    let mut length_buffer = [0; 8];
+
+                    while let Err(_err) = server_socket.read_exact(&mut length_buffer) {}
+
+                    let length = u64::from_le_bytes(length_buffer) as usize;
+
+                    // Read the amount of length bytes into a buffer
+                    let mut data_buffer = [0; REQUEST_MSG_BUFFER_SIZE];
+                    while let Err(_err) = server_socket.read_exact(&mut data_buffer[0..length]) {}
+
+                    let mut response = [0; RESPONSE_MSG_BUFFER_SIZE];
+                    server.handle_message(&data_buffer[0..length], &mut response)?;
+                }
+                Ok::<(), BrokerServerError>(())
+            });
+
+            // Create a MioBrokerClient and send a psk
+            let mut client = MioBrokerClient::new(client_socket);
 
-        // Spawn a new thread to connect to the unix socket
-        let handle = std::thread::spawn(move || {
             for _ in 0..TEST_RUNS {
-                // Wait for 8 bytes of length to come in
-                let mut length_buffer = [0; 8];
+                //Create psk of random 32 bytes
+                let psk = Secret::random();
+                let peer_id = Public::random();
+                let interface = "test";
+                let config = SerializedBrokerConfig {
+                    psk: &psk,
+                    peer_id: &peer_id,
+                    interface: interface.as_bytes(),
+                    additional_params: &[],
+                };
+                client.set_psk(config).unwrap();
 
-                while let Err(_err) = server_socket.read_exact(&mut length_buffer) {}
+                //Sleep for a while to allow the server to process the message
+                std::thread::sleep(std::time::Duration::from_millis(
+                    rand::thread_rng().gen_range(100..500),
+                ));
 
-                let length = u64::from_le_bytes(length_buffer) as usize;
+                let psk = psk.secret().to_owned();
 
-                // Read the amount of length bytes into a buffer
-                let mut data_buffer = [0; REQUEST_MSG_BUFFER_SIZE];
-                while let Err(_err) = server_socket.read_exact(&mut data_buffer[0..length]) {}
+                loop {
+                    let mut lock = server_broker_inner.try_lock();
 
-                let mut response = [0; RESPONSE_MSG_BUFFER_SIZE];
-                server.handle_message(&data_buffer[0..length], &mut response)?;
-            }
-            Ok::<(), BrokerServerError>(())
-        });
+                    if let Ok(ref mut inner) = lock {
+                        // Check if the psk is received by the server
+                        let received_psk = &inner.psk;
+                        assert_eq!(
+                            received_psk.as_ref().map(|psk| psk.secret().to_owned()),
+                            Some(psk)
+                        );
 
-        // Create a MioBrokerClient and send a psk
-        let mut client = MioBrokerClient::new(client_socket);
+                        let recieved_peer_id = inner.peer_id;
+                        assert_eq!(recieved_peer_id, Some(peer_id));
 
-        for _ in 0..TEST_RUNS {
-            //Create psk of random 32 bytes
-            let psk = Secret::random();
-            let peer_id = Public::random();
-            let interface = "test";
-            let config = SerializedBrokerConfig {
-                psk: &psk,
-                peer_id: &peer_id,
-                interface: interface.as_bytes(),
-                additional_params: &[],
-            };
-            client.set_psk(config).unwrap();
+                        let target_interface = &inner.interface;
+                        assert_eq!(target_interface.as_deref(), Some(interface));
 
-            //Sleep for a while to allow the server to process the message
-            std::thread::sleep(std::time::Duration::from_millis(
-                rand::thread_rng().gen_range(100..500),
-            ));
-
-            let psk = psk.secret().to_owned();
-
-            loop {
-                let mut lock = server_broker_inner.try_lock();
-
-                if let Ok(ref mut inner) = lock {
-                    // Check if the psk is received by the server
-                    let received_psk = &inner.psk;
-                    assert_eq!(
-                        received_psk.as_ref().map(|psk| psk.secret().to_owned()),
-                        Some(psk)
-                    );
-
-                    let recieved_peer_id = inner.peer_id;
-                    assert_eq!(recieved_peer_id, Some(peer_id));
-
-                    let target_interface = &inner.interface;
-                    assert_eq!(target_interface.as_deref(), Some(interface));
-
-                    break;
+                        break;
+                    }
                 }
             }
-        }
-        handle.join().unwrap().unwrap();
+            handle.join().unwrap().unwrap();
+        });
     }
 }