Add signature verification

Cargo.toml

@@ -8,9 +8,13 @@ edition = "2021"
 [dependencies]
 anyhow = "1.0.69"
 axum = { version = "0.6.11", features = ["json"] }
+p384 = { version = "0.13.0", features = ["ecdsa", "serde"] }
 serde = { version = "1.0.155", features = ["derive"] }
 serde_json = "1.0.94"
 tokio = { version = "1.26.0", features = ["full"] }
 tracing = "0.1.37"
 tracing-subscriber = "0.3.16"
 uuid = { version = "1.3.0", features = ["v4", "serde"] }
+
+[dev-dependencies]
+uuid = { version = "1.3.0", features = ["v4"] }

src/lib.rs

@@ -1,8 +1,11 @@
 use axum::extract::State;
 use axum::{http::StatusCode, Json};
+use p384::ecdsa::signature::Verifier;
+use p384::ecdsa::{Signature, VerifyingKey};
 use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
 use std::net::{IpAddr, SocketAddr};
+use std::str::FromStr;
 use std::sync::{Arc, Mutex};
 use tokio::io::{self, AsyncWriteExt};
 use tokio::net::{TcpListener, TcpStream};
@@ -12,7 +15,16 @@ use uuid::Uuid;
 #[derive(Deserialize, Serialize, Debug)]
 pub struct ProxyCommand {
     command: Command,
-    signature: [u8; 32],
+    signature: Signature,
+}
+
+impl ProxyCommand {
+    fn verify_signature(&self, verifying_key: &VerifyingKey) -> bool {
+        let message = serde_json::to_string(&self.command).unwrap();
+        verifying_key
+            .verify(message.as_bytes(), &self.signature)
+            .is_ok()
+    }
 }
 
 #[derive(Deserialize, Serialize, Debug)]
@@ -41,12 +53,14 @@ pub struct ProxyResponse {
 #[derive(Debug)]
 pub struct GlobalState {
     proxies: Mutex<HashMap<Uuid, ProxyState>>,
+    verifying_key: VerifyingKey,
 }
 
 impl GlobalState {
-    pub fn new() -> Self {
+    pub fn new(verifying_key: &str) -> Self {
         Self {
             proxies: Mutex::new(HashMap::new()),
+            verifying_key: VerifyingKey::from_str(verifying_key).unwrap(),
         }
     }
 }
@@ -66,7 +80,15 @@ pub async fn process_command(
     Json(payload): Json<ProxyCommand>,
 ) -> (StatusCode, Json<ProxyResponse>) {
     tracing::info!("Received payload: {:?}", payload);
-    // TODO: verify signature
+
+    if !payload.verify_signature(&state.verifying_key) {
+        return (
+            StatusCode::UNAUTHORIZED,
+            Json(ProxyResponse {
+                message: "Invalid signature".to_string(),
+            }),
+        );
+    }
     match payload.command {
         Command::New {
             incoming_port,
@@ -108,7 +130,7 @@ pub async fn process_command(
         }
     }
     (
-        StatusCode::CREATED,
+        StatusCode::ACCEPTED,
         Json(ProxyResponse {
             message: "Success".to_string(),
         }),
@@ -224,10 +246,16 @@ mod tests {
     use std::net::{IpAddr, Ipv4Addr};
 
     use crate::{Command, ProxyCommand};
+    use p384::{
+        ecdsa::{signature::Signer, Signature, SigningKey, VerifyingKey},
+        elliptic_curve::rand_core::OsRng,
+    };
     use uuid::uuid;
 
     #[test]
     fn serialize_proxy_command_new() {
+        let key = SigningKey::from_slice(&[1; 48]).unwrap();
+        let signature = key.sign(&[]); // Not a valid signature
         let proxy_command = ProxyCommand {
             command: Command::New {
                 incoming_port: 5555,
@@ -235,24 +263,62 @@ mod tests {
                 destination_ip: IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)),
                 id: uuid!("67e55044-10b1-426f-9247-bb680e5fe0c8"),
             },
-            signature: [0u8; 32],
+            signature,
         };
         let expected = "{\"command\":{\"New\":{\"incoming_port\":5555,\"destination_port\":6666,\"\
                         destination_ip\":\"127.0.0.1\",\"id\":\"67e55044-10b1-426f-9247-bb680e5fe0c8\"}},\
-                        \"signature\":[0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]}";
+                        \"signature\":\"\
+                            5C912C4B3BFF2ADB49885DCBDB53D6D3041D0632E498CDFF\
+                            2114CD2DCAC936AB0901B47C411E5BB57FE77BEF96044940\
+                            81680ADAD0775CD144E2D2678537F621ED587E13EB430126\
+                            C7A757AEC99CE08A2D0F3A5C9FB45E9349F36408DFD7BA17\"}";
+
         assert_eq!(serde_json::to_string(&proxy_command).unwrap(), expected);
     }
 
     #[test]
     fn serialize_proxy_command_delete() {
+        let key = SigningKey::from_slice(&[1; 48]).unwrap();
+        let signature = key.sign(&[]); // Not a valid signature
         let proxy_command = ProxyCommand {
             command: Command::Delete {
                 id: uuid!("67e55044-10b1-426f-9247-bb680e5fe0c8"),
             },
-            signature: [0u8; 32],
+            signature,
         };
-        let expected = "{\"command\":{\"Delete\":{\"id\":\"67e55044-10b1-426f-9247-bb680e5fe0c8\"}},\
-                        \"signature\":[0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]}";
+        let expected =
+            "{\"command\":{\"Delete\":{\"id\":\"67e55044-10b1-426f-9247-bb680e5fe0c8\"}},\
+                        \"signature\":\"\
+                            5C912C4B3BFF2ADB49885DCBDB53D6D3041D0632E498CDFF\
+                            2114CD2DCAC936AB0901B47C411E5BB57FE77BEF96044940\
+                            81680ADAD0775CD144E2D2678537F621ED587E13EB430126\
+                            C7A757AEC99CE08A2D0F3A5C9FB45E9349F36408DFD7BA17\"}";
+
         assert_eq!(serde_json::to_string(&proxy_command).unwrap(), expected);
     }
+
+    #[test]
+    fn verify_signature() {
+        let command = Command::New {
+            incoming_port: 4567,
+            destination_port: 7654,
+            destination_ip: IpAddr::V4(Ipv4Addr::new(123, 23, 76, 21)),
+            id: uuid::Uuid::new_v4(),
+        };
+
+        // Create signed message
+        let signing_key = SigningKey::random(&mut OsRng);
+        let message = serde_json::to_string(&command).unwrap();
+        let signature: Signature = signing_key.sign(message.as_bytes());
+        let bytes = signature.to_bytes();
+        assert_eq!(bytes.len(), 96);
+        let proxy_command = ProxyCommand {
+            command,
+            signature: bytes.as_slice().try_into().unwrap(),
+        };
+
+        // Verify signed message
+        let verifying_key = VerifyingKey::from(&signing_key);
+        assert!(proxy_command.verify_signature(&verifying_key));
+    }
 }

src/main.rs

@@ -15,7 +15,9 @@ async fn main() {
 
     tracing::subscriber::set_global_default(subscriber).unwrap();
 
-    let shared_state = Arc::new(GlobalState::new());
+    let verifying_key = std::env::args().nth(1).expect("No verifying key provided");
+
+    let shared_state = Arc::new(GlobalState::new(&verifying_key));
     // build our application with a route
     let app = Router::new()
         // `GET /` goes to `root`