From 6e307fc70e19093869311a04060cddf120b3c31e Mon Sep 17 00:00:00 2001 From: Mathias Hall-Andersen Date: Fri, 19 Jun 2020 23:45:56 +0200 Subject: Replace RwLock with DashMap in handshake --- Cargo.lock | 47 +++++++++++++++++++++++++++++++ Cargo.toml | 1 + src/wireguard/handshake/device.rs | 58 +++++++++++++++++++++------------------ src/wireguard/peer.rs | 6 ++-- src/wireguard/router/device.rs | 1 - 5 files changed, 84 insertions(+), 29 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 7bf08ab..662120d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -16,6 +16,14 @@ dependencies = [ "generic-array 0.12.3 (registry+https://github.com/rust-lang/crates.io-index)", ] +[[package]] +name = "ahash" +version = "0.3.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +dependencies = [ + "const-random 0.1.8 (registry+https://github.com/rust-lang/crates.io-index)", +] + [[package]] name = "aho-corasick" version = "0.7.10" @@ -157,6 +165,24 @@ dependencies = [ "bitflags 1.2.1 (registry+https://github.com/rust-lang/crates.io-index)", ] +[[package]] +name = "const-random" +version = "0.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +dependencies = [ + "const-random-macro 0.1.8 (registry+https://github.com/rust-lang/crates.io-index)", + "proc-macro-hack 0.5.16 (registry+https://github.com/rust-lang/crates.io-index)", +] + +[[package]] +name = "const-random-macro" +version = "0.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +dependencies = [ + "getrandom 0.1.14 (registry+https://github.com/rust-lang/crates.io-index)", + "proc-macro-hack 0.5.16 (registry+https://github.com/rust-lang/crates.io-index)", +] + [[package]] name = "cpuprofiler" version = "0.0.4" @@ -207,6 +233,16 @@ dependencies = [ "zeroize 1.1.0 (registry+https://github.com/rust-lang/crates.io-index)", ] +[[package]] +name = "dashmap" +version = "3.11.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +dependencies = [ + "ahash 0.3.8 (registry+https://github.com/rust-lang/crates.io-index)", + "cfg-if 0.1.10 (registry+https://github.com/rust-lang/crates.io-index)", + "num_cpus 1.13.0 (registry+https://github.com/rust-lang/crates.io-index)", +] + [[package]] name = "digest" version = "0.8.1" @@ -593,6 +629,11 @@ name = "ppv-lite86" version = "0.2.8" source = "registry+https://github.com/rust-lang/crates.io-index" +[[package]] +name = "proc-macro-hack" +version = "0.5.16" +source = "registry+https://github.com/rust-lang/crates.io-index" + [[package]] name = "proc-macro2" version = "1.0.14" @@ -1134,6 +1175,7 @@ dependencies = [ "clear_on_drop 0.2.3 (registry+https://github.com/rust-lang/crates.io-index)", "cpuprofiler 0.0.4 (registry+https://github.com/rust-lang/crates.io-index)", "crossbeam-channel 0.4.2 (registry+https://github.com/rust-lang/crates.io-index)", + "dashmap 3.11.4 (registry+https://github.com/rust-lang/crates.io-index)", "digest 0.8.1 (registry+https://github.com/rust-lang/crates.io-index)", "env_logger 0.7.1 (registry+https://github.com/rust-lang/crates.io-index)", "generic-array 0.12.3 (registry+https://github.com/rust-lang/crates.io-index)", @@ -1216,6 +1258,7 @@ dependencies = [ [metadata] "checksum addr2line 0.12.1 (registry+https://github.com/rust-lang/crates.io-index)" = "a49806b9dadc843c61e7c97e72490ad7f7220ae249012fbda9ad0609457c0543" "checksum aead 0.2.0 (registry+https://github.com/rust-lang/crates.io-index)" = "4cf01b9b56e767bb57b94ebf91a58b338002963785cdd7013e21c0d4679471e4" +"checksum ahash 0.3.8 (registry+https://github.com/rust-lang/crates.io-index)" = "e8fd72866655d1904d6b0997d0b07ba561047d070fbe29de039031c641b61217" "checksum aho-corasick 0.7.10 (registry+https://github.com/rust-lang/crates.io-index)" = "8716408b8bc624ed7f65d223ddb9ac2d044c0547b6fa4b0d554f3a9540496ada" "checksum arraydeque 0.4.5 (registry+https://github.com/rust-lang/crates.io-index)" = "f0ffd3d69bd89910509a5d31d1f1353f38ccffdd116dd0099bbd6627f7bd8ad8" "checksum atty 0.2.14 (registry+https://github.com/rust-lang/crates.io-index)" = "d9b39be18770d11421cdb1b9947a45dd3f37e93092cbf377614828a319d5fee8" @@ -1236,11 +1279,14 @@ dependencies = [ "checksum chacha20poly1305 0.4.1 (registry+https://github.com/rust-lang/crates.io-index)" = "48901293601228db2131606f741db33561f7576b5d19c99cd66222380a7dc863" "checksum clear_on_drop 0.2.3 (registry+https://github.com/rust-lang/crates.io-index)" = "97276801e127ffb46b66ce23f35cc96bd454fa311294bced4bbace7baa8b1d17" "checksum cloudabi 0.0.3 (registry+https://github.com/rust-lang/crates.io-index)" = "ddfc5b9aa5d4507acaf872de71051dfd0e309860e88966e1051e462a077aac4f" +"checksum const-random 0.1.8 (registry+https://github.com/rust-lang/crates.io-index)" = "2f1af9ac737b2dd2d577701e59fd09ba34822f6f2ebdb30a7647405d9e55e16a" +"checksum const-random-macro 0.1.8 (registry+https://github.com/rust-lang/crates.io-index)" = "25e4c606eb459dd29f7c57b2e0879f2b6f14ee130918c2b78ccb58a9624e6c7a" "checksum cpuprofiler 0.0.4 (registry+https://github.com/rust-lang/crates.io-index)" = "43f8479dbcfd2bbaa0c0c26779b913052b375981cdf533091f2127ea3d42e52b" "checksum crossbeam-channel 0.4.2 (registry+https://github.com/rust-lang/crates.io-index)" = "cced8691919c02aac3cb0a1bc2e9b73d89e832bf9a06fc579d4e71b68a2da061" "checksum crossbeam-utils 0.7.2 (registry+https://github.com/rust-lang/crates.io-index)" = "c3c7c73a2d1e9fc0886a08b93e98eb643461230d5f1925e4036204d5f2e261a8" "checksum crypto-mac 0.7.0 (registry+https://github.com/rust-lang/crates.io-index)" = "4434400df11d95d556bac068ddfedd482915eb18fe8bea89bc80b6e4b1c179e5" "checksum curve25519-dalek 2.0.0 (registry+https://github.com/rust-lang/crates.io-index)" = "26778518a7f6cffa1d25a44b602b62b979bd88adb9e99ffec546998cf3404839" +"checksum dashmap 3.11.4 (registry+https://github.com/rust-lang/crates.io-index)" = "8cfcd41ae02d60edded204341d2798ba519c336c51a37330aa4b98a1128def32" "checksum digest 0.8.1 (registry+https://github.com/rust-lang/crates.io-index)" = "f3d0c8c8752312f9713efd397ff63acb9f85585afbf179282e720e7704954dd5" "checksum env_logger 0.7.1 (registry+https://github.com/rust-lang/crates.io-index)" = "44533bbbb3bb3c1fa17d9f2e4e38bbbaf8396ba82193c4cb1b6445d711445d36" "checksum error-chain 0.12.2 (registry+https://github.com/rust-lang/crates.io-index)" = "d371106cc88ffdfb1eabd7111e432da544f16f3e2d7bf1dfe8bf575f1df045cd" @@ -1289,6 +1335,7 @@ dependencies = [ "checksum pnet_transport 0.25.0 (registry+https://github.com/rust-lang/crates.io-index)" = "1b75ccaee7b5daba9f9a7d47bceeb73cc32edde9952dc5409460d6621ec667b6" "checksum poly1305 0.5.2 (registry+https://github.com/rust-lang/crates.io-index)" = "b5829f50f48e9ddb79f3f7c3097029d0caee30f8286accb241416df603b080b8" "checksum ppv-lite86 0.2.8 (registry+https://github.com/rust-lang/crates.io-index)" = "237a5ed80e274dbc66f86bd59c1e25edc039660be53194b5fe0a482e0f2612ea" +"checksum proc-macro-hack 0.5.16 (registry+https://github.com/rust-lang/crates.io-index)" = "7e0456befd48169b9f13ef0f0ad46d492cf9d2dbb918bcf38e01eed4ce3ec5e4" "checksum proc-macro2 1.0.14 (registry+https://github.com/rust-lang/crates.io-index)" = "de40dd4ff82d9c9bab6dae29dbab1167e515f8df9ed17d2987cb6012db206933" "checksum proptest 0.9.6 (registry+https://github.com/rust-lang/crates.io-index)" = "01c477819b845fe023d33583ebf10c9f62518c8d79a0960ba5c36d6ac8a55a5b" "checksum quick-error 1.2.3 (registry+https://github.com/rust-lang/crates.io-index)" = "a1d01941d82fa2ab50be1e79e6714289dd7cde78eba4c074bc5a4374f650dfe0" diff --git a/Cargo.toml b/Cargo.toml index 625e830..8da2f84 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -25,6 +25,7 @@ clear_on_drop = "0.2.3" env_logger = "0.7" num_cpus = "^1.10" crossbeam-channel = "0.4" +dashmap = "3.11" cpuprofiler = { version = "*", optional = true } [dependencies.treebitmap] diff --git a/src/wireguard/handshake/device.rs b/src/wireguard/handshake/device.rs index 3a3d023..5e69921 100644 --- a/src/wireguard/handshake/device.rs +++ b/src/wireguard/handshake/device.rs @@ -1,11 +1,12 @@ -use spin::RwLock; use std::collections::hash_map; use std::collections::HashMap; use std::net::SocketAddr; use std::sync::Mutex; -use zerocopy::AsBytes; use byteorder::{ByteOrder, LittleEndian}; +use dashmap::mapref::entry::Entry; +use dashmap::DashMap; +use zerocopy::AsBytes; use rand::prelude::{CryptoRng, RngCore}; use rand::Rng; @@ -36,7 +37,7 @@ pub struct KeyState { /// (the instance is a Peer object in the parent module) pub struct Device { keyst: Option, - id_map: RwLock>, + id_map: DashMap, // concurrent map pk_map: HashMap<[u8; 32], Peer>, limiter: Mutex, } @@ -62,7 +63,7 @@ impl<'a, O> Iterator for Iter<'a, O> { */ impl Device { pub fn clear(&mut self) { - self.id_map.write().clear(); + self.id_map.clear(); self.pk_map.clear(); } @@ -96,7 +97,7 @@ impl Device { pub fn new() -> Device { Device { keyst: None, - id_map: RwLock::new(HashMap::new()), + id_map: DashMap::new(), pk_map: HashMap::new(), limiter: Mutex::new(RateLimiter::new()), } @@ -208,16 +209,14 @@ impl Device { /// /// The call might fail if the public key is not found pub fn remove(&mut self, pk: &PublicKey) -> Result<(), ConfigError> { - // take write-lock on receive id table - let mut id_map = self.id_map.write(); - // remove the peer self.pk_map .remove(pk.as_bytes()) .ok_or(ConfigError::new("Public key not in device"))?; - // purge the id map (linear scan) - id_map.retain(|_, v| v != pk.as_bytes()); + // remove every id entry for the peer in the public key map + // O(n) operations, however it is rare: only when removing peers. + self.id_map.retain(|_, v| v != pk.as_bytes()); Ok(()) } @@ -265,9 +264,8 @@ impl Device { /// /// * `id` - The (sender) id to release pub fn release(&self, id: u32) { - let mut m = self.id_map.write(); - debug_assert!(m.contains_key(&id), "Releasing id not allocated"); - m.remove(&id); + let old = self.id_map.remove(&id); + assert!(old.is_some(), "released id not allocated"); } /// Begin a new handshake @@ -446,32 +444,40 @@ impl Device { // // Return the peer currently associated with the receiver identifier pub(super) fn lookup_id(&self, id: u32) -> Result<(&Peer, PublicKey), HandshakeError> { - let im = self.id_map.read(); - let pk = im.get(&id).ok_or(HandshakeError::UnknownReceiverId)?; - match self.pk_map.get(pk) { + // obtain a read reference to entry in the id_map + let pk = self + .id_map + .get(&id) + .ok_or(HandshakeError::UnknownReceiverId)?; + + // lookup the public key from the pk map + match self.pk_map.get(&*pk) { Some(peer) => Ok((peer, PublicKey::from(*pk))), - _ => unreachable!(), // if the id-lookup succeeded, the peer should exist + _ => unreachable!(), } } // Internal function // - // Allocated a new receiver identifier for the peer + // Allocated a new receiver identifier for the peer. + // Implemented via rejection sampling. fn allocate(&self, rng: &mut R, pk: &PublicKey) -> u32 { loop { let id = rng.gen(); - // check membership with read lock - if self.id_map.read().contains_key(&id) { + // read lock the shard and do quick check + if self.id_map.contains_key(&id) { continue; } - // take write lock and add index - let mut m = self.id_map.write(); - if !m.contains_key(&id) { - m.insert(id, *pk.as_bytes()); - return id; - } + // write lock the shard and insert + match self.id_map.entry(id) { + Entry::Vacant(entry) => { + entry.insert(*pk.as_bytes()); + return id; + } + _ => (), + }; } } } diff --git a/src/wireguard/peer.rs b/src/wireguard/peer.rs index 27d39bd..37b4684 100644 --- a/src/wireguard/peer.rs +++ b/src/wireguard/peer.rs @@ -22,13 +22,15 @@ pub struct PeerInner { // wireguard device state pub wg: WireGuard, + // TODO: eliminate + pub pk: PublicKey, + // handshake state pub walltime_last_handshake: Mutex>, // walltime for last handshake (for UAPI status) pub last_handshake_sent: Mutex, // instant for last handshake - pub handshake_queued: AtomicBool, // is a handshake job currently queued for the peer? + pub handshake_queued: AtomicBool, // is a handshake job currently queued? // stats and configuration - pub pk: PublicKey, // public key (TODO: there has to be a way to remove this) pub rx_bytes: AtomicU64, // received bytes pub tx_bytes: AtomicU64, // transmitted bytes diff --git a/src/wireguard/router/device.rs b/src/wireguard/router/device.rs index 1a12abb..62ef932 100644 --- a/src/wireguard/router/device.rs +++ b/src/wireguard/router/device.rs @@ -3,7 +3,6 @@ use std::ops::Deref; use std::sync::atomic::AtomicBool; use std::sync::Arc; use std::thread; -use std::time::Instant; use log; use spin::{Mutex, RwLock}; -- cgit v1.2.3-59-g8ed1b