Merge pull request #656 from DefGuard/dev

* prevent from adding duplicate public keys * fix tests and cleanup
DefGuard · Jul 1, 2024 · 7e6ad6f · 7e6ad6f
2 parents f5f1056 + df77a98
commit 7e6ad6f
Show file tree

Hide file tree

Showing 6 changed files with 52 additions and 10 deletions.
diff --git a/src/error.rs b/src/error.rs
@@ -40,6 +40,8 @@ pub enum WebError {
     ModelError(String),
     #[error("Public key invalid {0}")]
     PubkeyValidation(String),
+    #[error("Public key already exists {0}")]
+    PubkeyExists(String),
     #[error("HTTP error: {0}")]
     Http(StatusCode),
     #[error(transparent)]

diff --git a/src/grpc/enrollment.rs b/src/grpc/enrollment.rs
@@ -327,6 +327,25 @@ impl EnrollmentServer {
             error!("Invalid pubkey {}", request.pubkey);
             Status::invalid_argument("invalid pubkey")
         })?;
+
+        // Make sure there is no device with the same pubkey, such state may lead to unexpected issues
+        if let Some(device) = Device::find_by_pubkey(&self.pool, &request.pubkey)
+            .await
+            .map_err(|_| {
+                error!("Failed to get device by its pubkey: {}", request.pubkey);
+                Status::internal("unexpected error")
+            })?
+        {
+            warn!(
+                "User {} failed to add device {}, identical pubkey ({}) already exists for device {}",
+                user.username,
+                request.name,
+                request.pubkey,
+                device.name
+            );
+            return Err(Status::invalid_argument("invalid key"));
+        };
+
         let mut device = Device::new(request.name, request.pubkey, enrollment.user_id);
 
         let mut transaction = self.pool.begin().await.map_err(|_| {

diff --git a/src/handlers/mod.rs b/src/handlers/mod.rs
@@ -92,6 +92,7 @@ impl From<WebError> for ApiResponse {
             ),
             WebError::IncorrectUsername(msg)
             | WebError::PubkeyValidation(msg)
+            | WebError::PubkeyExists(msg)
             | WebError::BadRequest(msg) => {
                 error!(msg);
                 ApiResponse::new(json!({ "msg": msg }), StatusCode::BAD_REQUEST)

diff --git a/src/handlers/wireguard.rs b/src/handlers/wireguard.rs
@@ -490,6 +490,17 @@ pub async fn add_device(
 
     Device::validate_pubkey(&add_device.wireguard_pubkey).map_err(WebError::PubkeyValidation)?;
 
+    // Make sure there is no device with the same pubkey, such state may lead to unexpected issues
+    if Device::find_by_pubkey(&appstate.pool, &add_device.wireguard_pubkey)
+        .await?
+        .is_some()
+    {
+        return Err(WebError::PubkeyExists(format!(
+            "Failed to add device {device_name}, identical pubkey ({}) already exists",
+            add_device.wireguard_pubkey
+        )));
+    }
+
     // save device
     let Some(user_id) = user.id else {
         error!(

diff --git a/tests/user.rs b/tests/user.rs
@@ -631,7 +631,7 @@ async fn test_user_add_device() {
     // add device for themselves
     let device_data = AddDevice {
         name: "TestDevice2".into(),
-        wireguard_pubkey: "mgVXE8WcfStoD8mRatHcX5aaQ0DlcpjvPXibHEOr9y8=".into(),
+        wireguard_pubkey: "hNuapt7lOxF93KUqZGUY00oKJxH8LYwwsUVB1uUa0y4=".into(),
     };
     let response = client
         .post("/api/v1/device/admin")
@@ -676,10 +676,19 @@ async fn test_user_add_device() {
         .content
         .contains("Device type:</span> iPhone, OS: iOS 17.1, Mobile Safari"));
 
+    // a device with duplicate pubkey cannot be added
+    let response = client
+        .post("/api/v1/device/hpotter")
+        .header(USER_AGENT, user_agent_header)
+        .json(&device_data)
+        .send()
+        .await;
+    assert_eq!(response.status(), StatusCode::BAD_REQUEST);
+
     // normal user cannot add a device for other users
     let device_data = AddDevice {
         name: "TestDevice3".into(),
-        wireguard_pubkey: "mgVXE8WcfStoD8mRatHcX5aaQ0DlcpjvPXibHEOr9y8=".into(),
+        wireguard_pubkey: "fF9K0tgatZTEJRvzpNUswr0h8HqCIi+v39B45+QZZzE=".into(),
     };
     let response = client
         .post("/api/v1/device/admin")

diff --git a/tests/wireguard.rs b/tests/wireguard.rs
@@ -310,7 +310,7 @@ async fn test_device_permissions() {
     let device = json!({"devices": [{
         "name": "device_2",
         "wireguard_ip": "10.0.0.3",
-        "wireguard_pubkey": "LQKsT6/3HWKuJmMulH63R8iK+5sI8FyYEL6WDIi6lQU=",
+        "wireguard_pubkey": "TJgN9JzUF5zdZAPYD96G/Wys2M3TvaT5TIrErUl20nI=",
         "user_id": 1,
         "created": "2023-05-05T23:56:04"
     }]});
@@ -323,7 +323,7 @@ async fn test_device_permissions() {
 
     let device = json!({
         "name": "device_3",
-        "wireguard_pubkey": "LQKsT6/3HWKuJmMulH63R8iK+5sI8FyYEL6WDIi6lQU=",
+        "wireguard_pubkey": "PKY3zg5/ecNyMjqLi6yJ3jwb4PvC/SGzjhJ3jrn2vVQ=",
     });
     let response = client
         .post("/api/v1/device/hpotter")
@@ -334,7 +334,7 @@ async fn test_device_permissions() {
     let device = json!({"devices": [{
         "name": "device_4",
         "wireguard_ip": "10.0.0.5",
-        "wireguard_pubkey": "LQKsT6/3HWKuJmMulH63R8iK+5sI8FyYEL6WDIi6lQU=",
+        "wireguard_pubkey": "gTMFF29nNLkJR1UhoiO3ZJLF60h2hW+WxmIu5DGJ0B4=",
         "user_id": 2,
         "created": "2023-05-05T23:56:04"
     }]});
@@ -352,7 +352,7 @@ async fn test_device_permissions() {
 
     let device = json!({
         "name": "device_5",
-        "wireguard_pubkey": "LQKsT6/3HWKuJmMulH63R8iK+5sI8FyYEL6WDIi6lQU=",
+        "wireguard_pubkey": "qhLnyggsD1nVOcLdTk0q43kOZHHknPQgftBY+ZLy40Q=",
     });
     let response = client
         .post("/api/v1/device/hpotter")
@@ -363,7 +363,7 @@ async fn test_device_permissions() {
     let device = json!({"devices": [{
         "name": "device_6",
         "wireguard_ip": "10.0.0.7",
-        "wireguard_pubkey": "LQKsT6/3HWKuJmMulH63R8iK+5sI8FyYEL6WDIi6lQU=",
+        "wireguard_pubkey": "xGLqgxVAnmk9+tsj5X/wzwouwx3bF1b3W+VWAb4NLjM=",
         "user_id": 2,
         "created": "2023-05-05T23:56:04"
     }]});
@@ -376,7 +376,7 @@ async fn test_device_permissions() {
 
     let device = json!({
         "name": "device_7",
-        "wireguard_pubkey": "LQKsT6/3HWKuJmMulH63R8iK+5sI8FyYEL6WDIi6lQU=",
+        "wireguard_pubkey": "J4p/w6R0xt4c2dIBDJ73BmzGJeF0QLW/iihPnISJMkg=",
     });
     let response = client
         .post("/api/v1/device/admin")
@@ -387,7 +387,7 @@ async fn test_device_permissions() {
     let device = json!({"devices": [{
         "name": "device_8",
         "wireguard_ip": "10.0.0.9",
-        "wireguard_pubkey": "LQKsT6/3HWKuJmMulH63R8iK+5sI8FyYEL6WDIi6lQU=",
+        "wireguard_pubkey": "A2cg4qMe+s0MSFlV6xyhz7XY6PrET6mli9GVSUshXAk=",
         "user_id": 1,
         "created": "2023-05-05T23:56:04"
     }]});
@@ -504,7 +504,7 @@ async fn test_device_pubkey() {
     let devices = json!({"devices": [{
         "name": "device_2",
         "wireguard_ip": "10.0.0.9",
-        "wireguard_pubkey": "LQKsT6/3HWKuJmMulH63R8iK+5sI8FyYEL6WDIi6lQU=",
+        "wireguard_pubkey": "o/8q3kmv5nnbrcb/7aceQWGE44a0yI707wObXRyyWGU=",
         "user_id": 1,
         "created": "2023-05-05T23:56:04"
     },