Refactor post_quantum_secrets to not take an enable-bool

1 files changed, 30 insertions, 32 deletions

diff --git a/talpid-tunnel-config-client/src/lib.rs b/talpid-tunnel-config-client/src/lib.rs
index ab47c13be7..7a7b428fc6 100644
--- a/talpid-tunnel-config-client/src/lib.rs
+++ b/talpid-tunnel-config-client/src/lib.rs
@@ -89,20 +89,24 @@ pub async fn request_ephemeral_peer_with(
     mut client: RelayConfigService,
     parent_pubkey: PublicKey,
     ephemeral_pubkey: PublicKey,
-    enable_post_quantum: bool,
+    enable_quantum_resistant: bool,
     enable_daita: bool,
 ) -> Result<EphemeralPeer, Error> {
-    let (pq_request, kem_secrets) = post_quantum_secrets(enable_post_quantum).await;
-    let daita = Some(proto::DaitaRequestV1 {
-        activate_daita: enable_daita,
-    });
+    let (pq_request, kem_secrets) = if enable_quantum_resistant {
+        let (pq_request, kem_secrets) = post_quantum_secrets().await;
+        (Some(pq_request), Some(kem_secrets))
+    } else {
+        (None, None)
+    };
 
     let response = client
         .register_peer_v1(proto::EphemeralPeerRequestV1 {
             wg_parent_pubkey: parent_pubkey.as_bytes().to_vec(),
             wg_ephemeral_peer_pubkey: ephemeral_pubkey.as_bytes().to_vec(),
             post_quantum: pq_request,
-            daita,
+            daita: Some(proto::DaitaRequestV1 {
+                activate_daita: enable_daita,
+            }),
         })
         .await
         .map_err(Error::GrpcError)?;
@@ -175,34 +179,28 @@ pub async fn request_ephemeral_peer(
     .await
 }
 
-async fn post_quantum_secrets(
-    enable_post_quantum: bool,
-) -> (
-    Option<PostQuantumRequestV1>,
-    Option<(classic_mceliece_rust::SecretKey<'static>, ml_kem::Keypair)>,
+async fn post_quantum_secrets() -> (
+    PostQuantumRequestV1,
+    (classic_mceliece_rust::SecretKey<'static>, ml_kem::Keypair),
 ) {
-    if enable_post_quantum {
-        let (cme_kem_pubkey, cme_kem_secret) = classic_mceliece::generate_keys().await;
-        let ml_kem_keypair = ml_kem::keypair();
+    let (cme_kem_pubkey, cme_kem_secret) = classic_mceliece::generate_keys().await;
+    let ml_kem_keypair = ml_kem::keypair();
 
-        (
-            Some(proto::PostQuantumRequestV1 {
-                kem_pubkeys: vec![
-                    proto::KemPubkeyV1 {
-                        algorithm_name: classic_mceliece::ALGORITHM_NAME.to_owned(),
-                        key_data: cme_kem_pubkey.as_array().to_vec(),
-                    },
-                    proto::KemPubkeyV1 {
-                        algorithm_name: ml_kem::ALGORITHM_NAME.to_owned(),
-                        key_data: ml_kem_keypair.encapsulation_key(),
-                    },
-                ],
-            }),
-            Some((cme_kem_secret, ml_kem_keypair)),
-        )
-    } else {
-        (None, None)
-    }
+    (
+        proto::PostQuantumRequestV1 {
+            kem_pubkeys: vec![
+                proto::KemPubkeyV1 {
+                    algorithm_name: classic_mceliece::ALGORITHM_NAME.to_owned(),
+                    key_data: cme_kem_pubkey.as_array().to_vec(),
+                },
+                proto::KemPubkeyV1 {
+                    algorithm_name: ml_kem::ALGORITHM_NAME.to_owned(),
+                    key_data: ml_kem_keypair.encapsulation_key(),
+                },
+            ],
+        },
+        (cme_kem_secret, ml_kem_keypair),
+    )
 }
 
 /// Performs `dst = dst ^ src`.