/ client / src / key_exchange.rs
key_exchange.rs
  1  use base64::{engine::general_purpose, Engine};
  2  use oqs::kem::{Kem, Algorithm};
  3  use ed25519_dalek::{VerifyingKey as Ed25519PublicKey, SecretKey as Ed25519SecretKey};
  4  use rand::rngs::OsRng;
  5  use serde::{Deserialize, Serialize};
  6  use sha3::{Sha3_512, Digest};
  7  use x25519_dalek::{EphemeralSecret, PublicKey as X25519PublicKey};
  8  use std::{error::Error, time::Duration};
  9  
 10  use crate::{create_client_with_proxy, fetch_ciphertext, fetch_kyber_pubkey, send_ciphertext, send_kyber_pubkey, sign_data_with_dilithium, sign_data_with_eddsa, verify_signature_with_dilithium, verify_signature_with_eddsa};
 11  
 12  pub fn kyber_key_exchange(
 13      room_id: &str,
 14      dilithium_pks: &[oqs::sig::PublicKey],
 15      dilithium_sk: &oqs::sig::SecretKey,
 16      server_url: &str, // Added server URL parameter
 17  ) -> Result<String, Box<dyn Error>> {
 18      // Initialize KEM (Kyber1024)
 19      let kemalg = Kem::new(Algorithm::Kyber1024)?;
 20  
 21      // Generate the key pair once at the start (for both Alice and Bob)
 22      let (kem_pk, kem_sk) = kemalg.keypair()?;
 23      let kem_pk_hex = hex::encode(kem_pk.as_ref());
 24  
 25      let public_key = fetch_kyber_pubkey(room_id, server_url); // Pass server_url
 26      let is_alice = match public_key {
 27          Some(ref key) if !key.is_empty() => {
 28              println!("Fetched public key: {}", key);
 29              false
 30          }
 31          _ => {
 32              println!("No valid public key found. Sending own Kyber public key.");
 33              send_kyber_pubkey(room_id, &kem_pk_hex, server_url); // Pass server_url
 34              true
 35          }
 36      };
 37  
 38      let shared_secret_result = if is_alice {
 39          let ciphertext = fetch_ciphertext(room_id, server_url); // Pass server_url
 40          // Find the "-----BEGIN SIGNATURE-----" delimiter
 41          let start_pos = ciphertext.find("-----BEGIN SIGNATURE-----").ok_or("Signature start not found")?;
 42          // Extract the ciphertext before the signature part (i.e., before the "-----BEGIN SIGNATURE-----")
 43          let ciphertext_before_signature = &ciphertext[..start_pos].trim();
 44  
 45          // If the extracted ciphertext before the signature is hex-encoded, decode it
 46          let decoded_ct = hex::decode(ciphertext_before_signature)?;
 47  
 48          // Iterate over dilithium_pks to verify the signature
 49          let mut signature_verified = false;
 50          for dilithium_pk in dilithium_pks {
 51              if verify_signature_with_dilithium(ciphertext.as_bytes(), dilithium_pk).is_ok() {
 52                  println!("Signature verified with Dilithium public key.");
 53                  signature_verified = true;
 54                  break;
 55              }
 56          }
 57  
 58          if !signature_verified {
 59              return Err("Failed to verify signature with any Dilithium public key.".into());
 60          }
 61  
 62          let ciphertext_obj = kemalg
 63              .ciphertext_from_bytes(&decoded_ct)
 64              .ok_or("Invalid ciphertext bytes")?;
 65  
 66          let shared_secret = kemalg.decapsulate(&kem_sk, &ciphertext_obj)?;
 67          let mut hasher = Sha3_512::new();
 68          hasher.update(shared_secret.as_ref());
 69          let result = hasher.finalize();
 70          let shared_secret_result = hex::encode(result);
 71  
 72          shared_secret_result
 73      } else {
 74          let alice_pk_bytes = hex::decode(public_key.unwrap())?;
 75          let alice_pk_ref = kemalg
 76              .public_key_from_bytes(&alice_pk_bytes)
 77              .ok_or("Failed to convert Alice's public key")?;
 78  
 79          let (kem_ct, shared_secret) = kemalg.encapsulate(&alice_pk_ref)?;
 80  
 81          // Bob signs the ciphertext
 82          let ciphertext_signature = sign_data_with_dilithium(kem_ct.as_ref(), dilithium_sk)?;
 83          println!("Bob signed the ciphertext: {}", ciphertext_signature);
 84  
 85          send_ciphertext(room_id, &ciphertext_signature, server_url); // Pass server_url
 86  
 87          let mut hasher = Sha3_512::new();
 88          hasher.update(shared_secret.as_ref());
 89          let result = hasher.finalize();
 90          let shared_secret_result = hex::encode(result);
 91  
 92          shared_secret_result
 93      };
 94  
 95      Ok(shared_secret_result)
 96  }
 97  
 98  // Structures for public keys
 99  #[derive(Serialize, Deserialize)]
100  struct Message {
101      message: String,
102      room_id: String,
103  }
104  
105  pub fn perform_ecdh_key_exchange(
106      room_id: &str,
107      eddsa_sk: &Ed25519SecretKey,
108      eddsa_pk: &Ed25519PublicKey,
109      server_url: &str, // Added parameter for server URL
110  ) -> Result<String, Box<dyn std::error::Error>> {
111      // Generate ECDH key pair using X25519 once at the start
112      let secret_key = EphemeralSecret::random_from_rng(OsRng);
113      let public_key = X25519PublicKey::from(&secret_key);
114  
115      let public_key_bytes = public_key.as_bytes();
116  
117      // Sign the public key using EdDSA
118      let signed_public_key = sign_data_with_eddsa(public_key_bytes, eddsa_sk)?;
119  
120      // Format the signed public key with the proper markers
121      let formatted_signed_public_key = format!("ECDH_PUBLIC_KEY:{}[END DATA]", signed_public_key);
122  
123      // Check if the server URL contains `.i2p`
124      let proxy = if server_url.contains(".i2p") {
125          "http://127.0.0.1:4444" // I2P Proxy address
126      } else {
127          "socks5h://127.0.0.1:9050" // SOCKS5 Proxy address (Tor)
128      };
129      let client = create_client_with_proxy(proxy);
130  
131      loop {
132          // Send the formatted signed public key to the server
133          let message = Message {
134              message: formatted_signed_public_key.clone(),
135              room_id: room_id.to_string(),
136          };
137  
138          let send_url = format!("{}/send", server_url); // Use server_url for the send endpoint
139          if let Err(err) = client.post(&send_url).json(&message).timeout(Duration::from_secs(60)).send() {
140              eprintln!("Failed to send signed public key to the server: {}", err);
141              continue; // Retry
142          } else {
143              println!("Successfully sent signed public key to the server.");
144          }
145  
146          // Fetch the other party's signed public key
147          let fetch_url = format!("{}/messages?room_id={}", server_url, room_id); // Use server_url for the fetch endpoint
148          let res = match client.get(&fetch_url).timeout(Duration::from_secs(60)).send() {
149              Ok(response) => response,
150              Err(err) => {
151                  eprintln!("Failed to fetch the other party's public key: {}", err);
152                  continue; // Retry
153              }
154          };
155  
156          if !res.status().is_success() {
157              eprintln!("Non-success status code while fetching messages: {}", res.status());
158              continue; // Retry
159          }
160  
161          let html_response = match res.text() {
162              Ok(text) => text,
163              Err(err) => {
164                  eprintln!("Failed to read response text: {}", err);
165                  continue; // Retry
166              }
167          };
168  
169          // Look for all the signed public keys in the response
170          let start_tag = "ECDH_PUBLIC_KEY:";
171          let end_tag = "[END DATA]";
172          let mut keys_processed = false;
173  
174          let mut start = 0;
175          while let Some(start_pos) = html_response[start..].find(start_tag) {
176              start += start_pos + start_tag.len();
177              if let Some(end_pos) = html_response[start..].find(end_tag) {
178                  let extracted_signed_key = &html_response[start..start + end_pos].trim();
179  
180                  // Verify the other party's public key signature
181                  if let Err(err) = verify_signature_with_eddsa(extracted_signed_key, eddsa_pk) {
182                      eprintln!("Failed to verify the signature: {}", err);
183                      continue; // Retry if verification fails
184                  }
185  
186                  // If verification is successful, extract the public key and proceed
187                  let extracted_key = extracted_signed_key.split("-----BEGIN SIGNATURE-----").next().unwrap().trim();
188  
189                  // Ignore if it's the same as our own public key
190                  if extracted_key == formatted_signed_public_key {
191                      println!("Ignoring our own public key.");
192                      start += end_pos + end_tag.len(); // Move past the end tag to continue searching
193                      keys_processed = true;
194                      continue; // Skip processing if it's our own public key
195                  }
196  
197                  match hex::decode(extracted_key) {
198                      Ok(other_public_key_bytes) => {
199                          if other_public_key_bytes.len() != 32 {
200                              eprintln!("Invalid public key length: {}", other_public_key_bytes.len());
201                              continue;
202                          }
203  
204                          let other_public_key_bytes =
205                              match <[u8; 32]>::try_from(other_public_key_bytes.as_slice()) {
206                                  Ok(bytes) => bytes,
207                                  Err(err) => {
208                                      eprintln!("Failed to convert other public key bytes: {}", err);
209                                      continue;
210                                  }
211                              };
212  
213                          let other_public_key = X25519PublicKey::from(other_public_key_bytes);
214                          let shared_secret = secret_key.diffie_hellman(&other_public_key);
215                          let shared_secret_base64 = general_purpose::STANDARD.encode(shared_secret.as_bytes());
216                          return Ok(shared_secret_base64);
217                      }
218                      Err(err) => {
219                          eprintln!("Failed to decode other public key: {}", err);
220                      }
221                  }
222  
223                  // Move past the end tag
224                  start += end_pos + end_tag.len();
225                  keys_processed = true;
226              } else {
227                  eprintln!("End tag not found after start tag. Skipping.");
228                  break;
229              }
230          }
231  
232          if !keys_processed {
233              eprintln!("No valid other signed public keys found. Retrying...");
234          }
235      }
236  }