Do explicit zeroization of key material, clean up PQ code

5 files changed, 80 insertions, 42 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 6dcabc7113..e4a510edc8 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2308,6 +2308,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "32b79004a05337e54e8ffc0ec7470e40fa26eca6fe182968ec2b803247f2283c"
 dependencies = [
  "rand_core 0.6.4",
+ "zeroize",
 ]
 
 [[package]]
@@ -3304,6 +3305,7 @@ dependencies = [
  "tonic",
  "tonic-build",
  "tower",
+ "zeroize",
 ]
 
 [[package]]
diff --git a/talpid-tunnel-config-client/Cargo.toml b/talpid-tunnel-config-client/Cargo.toml
index 4e78b7ad0c..62542da683 100644
--- a/talpid-tunnel-config-client/Cargo.toml
+++ b/talpid-tunnel-config-client/Cargo.toml
@@ -15,8 +15,9 @@ tonic = "0.8"
 prost = "0.11"
 tower = "0.4"
 tokio = "1"
-classic-mceliece-rust = { version = "2.0.0", features = ["mceliece460896f"] }
-pqc_kyber = { version = "0.4.0", features = ["std", "kyber1024"] }
+classic-mceliece-rust = { version = "2.0.0", features = ["mceliece460896f", "zeroize"] }
+pqc_kyber = { version = "0.4.0", features = ["std", "kyber1024", "zeroize"] }
+zeroize = "1.5.7"
 
 [dev-dependencies]
 tokio = { version = "1", features = ["rt-multi-thread"] }
diff --git a/talpid-tunnel-config-client/src/classic_mceliece.rs b/talpid-tunnel-config-client/src/classic_mceliece.rs
index dbba95f067..780566958d 100644
--- a/talpid-tunnel-config-client/src/classic_mceliece.rs
+++ b/talpid-tunnel-config-client/src/classic_mceliece.rs
@@ -1,4 +1,6 @@
-use classic_mceliece_rust::{keypair_boxed, SharedSecret};
+use classic_mceliece_rust::{
+    keypair_boxed, Ciphertext, PublicKey, SecretKey, SharedSecret, CRYPTO_CIPHERTEXTBYTES,
+};
 
 /// The `keypair_boxed` function needs just under 1 MiB of stack in debug
 /// builds. Even though it probably works to run it directly on the main
@@ -6,7 +8,9 @@ use classic_mceliece_rust::{keypair_boxed, SharedSecret};
 /// keys on a separate thread with a large enough stack.
 const STACK_SIZE: usize = 2 * 1024 * 1024;
 
-pub use classic_mceliece_rust::{Ciphertext, PublicKey, SecretKey, CRYPTO_CIPHERTEXTBYTES};
+/// Use the smallest CME variant with NIST security level 3. This variant has significantly smaller
+/// keys than the larger variants, and is considered safe.
+pub const ALGORITHM_NAME: &str = "Classic-McEliece-460896f-round3";
 
 pub async fn generate_keys() -> (PublicKey<'static>, SecretKey<'static>) {
     let (tx, rx) = tokio::sync::oneshot::channel();
@@ -22,6 +26,21 @@ pub async fn generate_keys() -> (PublicKey<'static>, SecretKey<'static>) {
     rx.await.unwrap()
 }
 
-pub fn decapsulate(secret: &SecretKey, ciphertext: &Ciphertext) -> SharedSecret<'static> {
-    classic_mceliece_rust::decapsulate_boxed(ciphertext, secret)
+pub fn decapsulate(
+    secret: &SecretKey,
+    ciphertext_slice: &[u8],
+) -> Result<SharedSecret<'static>, super::Error> {
+    let ciphertext_array =
+        <[u8; CRYPTO_CIPHERTEXTBYTES]>::try_from(ciphertext_slice).map_err(|_| {
+            super::Error::InvalidCiphertextLength {
+                algorithm: ALGORITHM_NAME,
+                actual: ciphertext_slice.len(),
+                expected: CRYPTO_CIPHERTEXTBYTES,
+            }
+        })?;
+    let ciphertext = Ciphertext::from(ciphertext_array);
+    Ok(classic_mceliece_rust::decapsulate_boxed(
+        &ciphertext,
+        secret,
+    ))
 }
diff --git a/talpid-tunnel-config-client/src/kyber.rs b/talpid-tunnel-config-client/src/kyber.rs
index 19cc3f338c..0946e7f5c2 100644
--- a/talpid-tunnel-config-client/src/kyber.rs
+++ b/talpid-tunnel-config-client/src/kyber.rs
@@ -1,10 +1,28 @@
-use pqc_kyber::SecretKey;
+use pqc_kyber::{SecretKey, KYBER_CIPHERTEXTBYTES};
 
-pub use pqc_kyber::{keypair, KyberError, KYBER_CIPHERTEXTBYTES};
+pub use pqc_kyber::{keypair, KyberError};
 
+/// Use the strongest variant of Kyber. It is fast and the keys are small, so there is no practical
+/// benefit of going with anything lower.
+pub const ALGORITHM_NAME: &str = "Kyber1024";
+
+// Always inline in order to try to avoid potential copies of `shared_secret` to multiple places on the stack.
+#[inline(always)]
 pub fn decapsulate(
     secret_key: SecretKey,
-    ciphertext: [u8; KYBER_CIPHERTEXTBYTES],
-) -> Result<[u8; 32], KyberError> {
-    pqc_kyber::decapsulate(ciphertext.as_slice(), secret_key.as_slice())
+    ciphertext_slice: &[u8],
+) -> Result<[u8; 32], super::Error> {
+    // The `pqc_kyber` library takes a byte slice. But we convert it into an array
+    // in order to catch the length mismatch error and report it better than `pqc_kyber` would.
+    let ciphertext_array =
+        <[u8; KYBER_CIPHERTEXTBYTES]>::try_from(ciphertext_slice).map_err(|_| {
+            super::Error::InvalidCiphertextLength {
+                algorithm: ALGORITHM_NAME,
+                actual: ciphertext_slice.len(),
+                expected: KYBER_CIPHERTEXTBYTES,
+            }
+        })?;
+    let shared_secret = pqc_kyber::decapsulate(ciphertext_array.as_slice(), secret_key.as_slice())
+        .map_err(super::Error::FailedDecapsulateKyber)?;
+    Ok(shared_secret)
 }
diff --git a/talpid-tunnel-config-client/src/lib.rs b/talpid-tunnel-config-client/src/lib.rs
index e7d8e766ca..a4beee9be3 100644
--- a/talpid-tunnel-config-client/src/lib.rs
+++ b/talpid-tunnel-config-client/src/lib.rs
@@ -1,6 +1,7 @@
 use std::{fmt, net::IpAddr};
 use talpid_types::net::wireguard::{PresharedKey, PrivateKey, PublicKey};
 use tonic::transport::Channel;
+use zeroize::Zeroize;
 
 mod classic_mceliece;
 mod kyber;
@@ -14,8 +15,14 @@ mod proto {
 pub enum Error {
     GrpcConnectError(tonic::transport::Error),
     GrpcError(tonic::Status),
-    InvalidCiphertextLength { actual: usize, expected: usize },
-    InvalidCiphertextCount { actual: usize },
+    InvalidCiphertextLength {
+        algorithm: &'static str,
+        actual: usize,
+        expected: usize,
+    },
+    InvalidCiphertextCount {
+        actual: usize,
+    },
     FailedDecapsulateKyber(kyber::KyberError),
 }
 
@@ -25,9 +32,13 @@ impl std::fmt::Display for Error {
         match self {
             GrpcConnectError(_) => "Failed to connect to config service".fmt(f),
             GrpcError(status) => write!(f, "RPC failed: {status}"),
-            InvalidCiphertextLength { actual, expected } => write!(
+            InvalidCiphertextLength {
+                algorithm,
+                actual,
+                expected,
+            } => write!(
                 f,
-                "Expected a ciphertext of length {expected}, got {actual} bytes"
+                "Expected a {expected} bytes ciphertext for {algorithm}, got {actual} bytes"
             ),
             InvalidCiphertextCount { actual } => {
                 write!(f, "Expected 2 ciphertext in the response, got {actual}")
@@ -52,14 +63,6 @@ type RelayConfigService = proto::post_quantum_secure_client::PostQuantumSecureCl
 /// Port used by the tunnel config service.
 pub const CONFIG_SERVICE_PORT: u16 = 1337;
 
-/// Use the smallest CME variant with NIST security level 3. This variant has significantly smaller
-/// keys than the larger variants, and is considered safe.
-const CLASSIC_MCELIECE_VARIANT: &str = "Classic-McEliece-460896f-round3";
-
-/// Use the strongest variant of Kyber. It is fast and the keys are small, so there is no practical
-/// benefit of going with anything lower.
-const KYBER_VARIANT: &str = "Kyber1024";
-
 /// Generates a new WireGuard key pair and negotiates a PSK with the relay in a PQ-safe
 /// manner. This creates a peer on the relay with the new WireGuard pubkey and PSK,
 /// which can then be used to establish a PQ-safe tunnel to the relay.
@@ -79,11 +82,11 @@ pub async fn push_pq_key(
             wg_psk_pubkey: wg_psk_privkey.public_key().as_bytes().to_vec(),
             kem_pubkeys: vec![
                 proto::KemPubkeyV1 {
-                    algorithm_name: CLASSIC_MCELIECE_VARIANT.to_owned(),
+                    algorithm_name: classic_mceliece::ALGORITHM_NAME.to_owned(),
                     key_data: cme_kem_pubkey.as_array().to_vec(),
                 },
                 proto::KemPubkeyV1 {
-                    algorithm_name: KYBER_VARIANT.to_owned(),
+                    algorithm_name: kyber::ALGORITHM_NAME.to_owned(),
                     key_data: kyber_keypair.public.to_vec(),
                 },
             ],
@@ -105,28 +108,23 @@ pub async fn push_pq_key(
 
     // Decapsulate Classic McEliece and mix into PSK
     {
-        let ciphertext_array =
-            <[u8; classic_mceliece::CRYPTO_CIPHERTEXTBYTES]>::try_from(cme_ciphertext.as_slice())
-                .map_err(|_| Error::InvalidCiphertextLength {
-                actual: cme_ciphertext.len(),
-                expected: classic_mceliece::CRYPTO_CIPHERTEXTBYTES,
-            })?;
-        let ciphertext = classic_mceliece::Ciphertext::from(ciphertext_array);
-        let shared_secret = classic_mceliece::decapsulate(&cme_kem_secret, &ciphertext);
+        let mut shared_secret = classic_mceliece::decapsulate(&cme_kem_secret, cme_ciphertext)?;
         xor_assign(&mut psk_data, shared_secret.as_array());
+
+        // This should happen automatically due to `SharedSecret` implementing ZeroizeOnDrop. But doing it explicitly
+        // provides a stronger guarantee that it's not accidentally removed.
+        shared_secret.zeroize();
     }
     // Decapsulate Kyber and mix into PSK
     {
-        let ciphertext_array = <[u8; kyber::KYBER_CIPHERTEXTBYTES]>::try_from(
-            kyber_ciphertext.as_slice(),
-        )
-        .map_err(|_| Error::InvalidCiphertextLength {
-            actual: kyber_ciphertext.len(),
-            expected: kyber::KYBER_CIPHERTEXTBYTES,
-        })?;
-        let shared_secret = kyber::decapsulate(kyber_keypair.secret, ciphertext_array)
-            .map_err(Error::FailedDecapsulateKyber)?;
+        let mut shared_secret = kyber::decapsulate(kyber_keypair.secret, kyber_ciphertext)?;
         xor_assign(&mut psk_data, &shared_secret);
+
+        // The shared secret is sadly stored in an array on the stack. So we can't get any
+        // guarantees that it's not copied around on the stack. The best we can do here
+        // is to zero out the version we have and hope the compiler optimizes out copies.
+        // https://github.com/Argyle-Software/kyber/issues/59
+        shared_secret.zeroize();
     }
 
     Ok((wg_psk_privkey, PresharedKey::from(psk_data)))