From 653c314409c017e47f3338552e7998196a3b6de9 Mon Sep 17 00:00:00 2001
From: George Hopkins <george-hopkins@null.net>
Date: Mon, 2 Oct 2023 16:24:37 +0200
Subject: [PATCH] Support pre-shared key

---
 Cargo.toml    |  1 +
 src/config.rs | 24 ++++++++++++++++++++++++
 src/wg.rs     |  2 +-
 3 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/Cargo.toml b/Cargo.toml
index 6ec0b5e..4dee3d7 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -22,6 +22,7 @@ async-trait = "0.1"
 priority-queue = "1.3.0"
 smoltcp = { version = "0.8.2", default-features = false, features = ["std", "log", "medium-ip", "proto-ipv4", "proto-ipv6", "socket-udp", "socket-tcp"] }
 bytes = "1"
+base64 = "0.21"
 
 # forward boringtuns tracing events to log
 tracing = { version = "0.1", default-features = false, features = ["log"] }
diff --git a/src/config.rs b/src/config.rs
index 93d9129..54839cf 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -6,6 +6,7 @@ use std::net::{IpAddr, SocketAddr, ToSocketAddrs};
 use std::sync::Arc;
 
 use anyhow::Context;
+use base64::prelude::{Engine as _, BASE64_STANDARD};
 pub use boringtun::crypto::{X25519PublicKey, X25519SecretKey};
 
 const DEFAULT_PORT_FORWARD_SOURCE: &str = "127.0.0.1";
@@ -17,6 +18,7 @@ pub struct Config {
     pub remote_port_forwards: Vec<PortForwardConfig>,
     pub private_key: Arc<X25519SecretKey>,
     pub endpoint_public_key: Arc<X25519PublicKey>,
+    pub endpoint_preshared_key: Option<[u8; 32]>,
     pub endpoint_addr: SocketAddr,
     pub endpoint_bind_addr: SocketAddr,
     pub source_peer_ip: IpAddr,
@@ -73,6 +75,12 @@ impl Config {
                     .long("endpoint-public-key")
                     .env("ONETUN_ENDPOINT_PUBLIC_KEY")
                     .help("The public key of the WireGuard endpoint (remote)."),
+                Arg::with_name("endpoint-preshared-key")
+                    .required(false)
+                    .takes_value(true)
+                    .long("endpoint-preshared-key")
+                    .env("ONETUN_ENDPOINT_PRESHARED_KEY")
+                    .help("The pre-shared key of the WireGuard endpoint (remote)."),
                 Arg::with_name("endpoint-addr")
                     .required(true)
                     .takes_value(true)
@@ -264,6 +272,9 @@ impl Config {
                 parse_public_key(matches.value_of("endpoint-public-key"))
                     .with_context(|| "Invalid endpoint public key")?,
             ),
+            endpoint_preshared_key: parse_preshared_key(
+                matches.value_of("endpoint-preshared-key"),
+            )?,
             endpoint_addr,
             endpoint_bind_addr,
             source_peer_ip,
@@ -304,6 +315,19 @@ fn parse_public_key(s: Option<&str>) -> anyhow::Result<X25519PublicKey> {
         .with_context(|| "Invalid public key")
 }
 
+fn parse_preshared_key(s: Option<&str>) -> anyhow::Result<Option<[u8; 32]>> {
+    if let Some(s) = s {
+        let psk = BASE64_STANDARD
+            .decode(s)
+            .with_context(|| "Invalid pre-shared key")?;
+        Ok(Some(psk.try_into().map_err(|_| {
+            anyhow::anyhow!("Unsupported pre-shared key")
+        })?))
+    } else {
+        Ok(None)
+    }
+}
+
 fn parse_keep_alive(s: Option<&str>) -> anyhow::Result<Option<u16>> {
     if let Some(s) = s {
         let parsed: u16 = s.parse().with_context(|| {
diff --git a/src/wg.rs b/src/wg.rs
index 0646606..0f2278a 100644
--- a/src/wg.rs
+++ b/src/wg.rs
@@ -221,7 +221,7 @@ impl WireGuardTunnel {
         Tunn::new(
             config.private_key.clone(),
             config.endpoint_public_key.clone(),
-            None,
+            config.endpoint_preshared_key,
             config.keepalive_seconds,
             0,
             None,