Update example code for PQ PSK exchange

1 files changed, 16 insertions, 9 deletions

diff --git a/talpid-relay-config-client/examples/psk-exchange.rs b/talpid-relay-config-client/examples/psk-exchange.rs
index de768c6964..51cb138e53 100644
--- a/talpid-relay-config-client/examples/psk-exchange.rs
+++ b/talpid-relay-config-client/examples/psk-exchange.rs
@@ -1,17 +1,24 @@
-use std::net::{IpAddr, Ipv4Addr};
+use std::{
+    io,
+    net::{IpAddr, Ipv4Addr},
+};
 
-use talpid_types::net::wireguard::PrivateKey;
+use talpid_types::net::wireguard::PublicKey;
 
 #[tokio::main]
 async fn main() {
-    let current_private_key = PrivateKey::new_from_random();
+    println!("Make sure you're connected to a WireGuard peer and enter your public key: ");
 
-    let (private_key, psk) = talpid_relay_config_client::push_pq_key(
-        IpAddr::V4(Ipv4Addr::new(10, 64, 0, 1)),
-        current_private_key.public_key(),
-    )
-    .await
-    .unwrap();
+    let mut pubkey_s = String::new();
+    io::stdin()
+        .read_line(&mut pubkey_s)
+        .expect("Failed to read from stdin");
+    let pubkey = PublicKey::from_base64(pubkey_s.trim()).expect("Invalid public key");
+
+    let (private_key, psk) =
+        talpid_relay_config_client::push_pq_key(IpAddr::V4(Ipv4Addr::new(10, 64, 0, 1)), pubkey)
+            .await
+            .unwrap();
 
     println!("private key: {:?}", private_key);
     println!("psk: {:?}", psk);