Fix WG keys, add DELETE connection method

Author: the2pizza

dev/dev.sql

@@ -96,4 +96,6 @@ ALTER COLUMN password DROP NOT NULL;
 
 ALTER TABLE inbounds ADD COLUMN dns INET[];
 
+ALTER TABLE connections DROP CONSTRAINT connections_user_id_fkey;
+
 

src/bin/api/core/http/handlers/connection.rs

@@ -1,4 +1,5 @@
 use base64::Engine;
+
 use warp::http::Response;
 
 use defguard_wireguard_rs::net::IpAddrMask;
@@ -17,6 +18,8 @@ use pony::state::storage::connection::ApiOp;
 use pony::utils;
 use pony::xray_op::clash::generate_clash_config;
 use pony::xray_op::clash::generate_proxy_config;
+use pony::zmq::message::Action;
+use pony::zmq::message::Message;
 use pony::zmq::publisher::Publisher as ZmqPublisher;
 use pony::Conn as Connection;
 use pony::ConnWithId;
@@ -448,6 +451,123 @@ where
     }
 }
 
+/// Handler deletes connection
+// DELETE /connection?conn_id=
+pub async fn delete_connection_handler<N, C>(
+    conn_param: ConnQueryParam,
+    publisher: ZmqPublisher,
+    state: SyncState<N, C>,
+) -> Result<impl warp::Reply, warp::Rejection>
+where
+    N: NodeStorageOp + Sync + Send + Clone + 'static,
+    C: ConnectionApiOp
+        + ConnectionBaseOp
+        + Sync
+        + Send
+        + Clone
+        + 'static
+        + From<Connection>
+        + PartialEq,
+    Connection: From<C>,
+{
+    let conn_id = conn_param.conn_id;
+    let conn_opt = {
+        let mem = state.memory.lock().await;
+        mem.connections.get(&conn_id).cloned()
+    };
+
+    let Some(conn) = conn_opt else {
+        let response = ResponseMessage::<Option<uuid::Uuid>> {
+            status: StatusCode::NOT_FOUND.as_u16(),
+            message: format!("Connection {} not found", conn_id),
+            response: None,
+        };
+        return Ok(warp::reply::with_status(
+            warp::reply::json(&response),
+            StatusCode::NOT_FOUND,
+        ));
+    };
+
+    if conn.get_deleted() {
+        let response = ResponseMessage::<Option<uuid::Uuid>> {
+            status: StatusCode::NOT_FOUND.as_u16(),
+            message: format!("Connection {} is deleted", conn_id),
+            response: None,
+        };
+        return Ok(warp::reply::with_status(
+            warp::reply::json(&response),
+            StatusCode::NOT_FOUND,
+        ));
+    }
+
+    match SyncOp::delete_connection(&state, &conn_id).await {
+        Ok(StorageOperationStatus::Ok(id)) => {
+            let msg = Message {
+                action: Action::Delete,
+                conn_id: conn_id,
+                password: conn.get_password(),
+                wg: conn.get_wireguard().cloned(),
+                tag: conn.get_proto().proto(),
+            };
+
+            if let Some(node_id) = conn.get_wireguard_node_id() {
+                let _ = publisher.send(&node_id.to_string(), msg).await;
+            } else {
+                let _ = publisher.send(&conn.get_env(), msg).await;
+            }
+
+            let response = ResponseMessage::<Option<uuid::Uuid>> {
+                status: StatusCode::OK.as_u16(),
+                message: format!("Connection {} has been deleted", id),
+                response: Some(id),
+            };
+            Ok(warp::reply::with_status(
+                warp::reply::json(&response),
+                StatusCode::OK,
+            ))
+        }
+
+        Ok(StorageOperationStatus::NotFound(id)) => {
+            let response = ResponseMessage::<Option<uuid::Uuid>> {
+                status: StatusCode::NOT_FOUND.as_u16(),
+                message: format!("Connection {} not found", id),
+                response: None,
+            };
+            Ok(warp::reply::with_status(
+                warp::reply::json(&response),
+                StatusCode::NOT_FOUND,
+            ))
+        }
+
+        Ok(status) => {
+            let response = ResponseMessage::<Option<uuid::Uuid>> {
+                status: StatusCode::BAD_REQUEST.as_u16(),
+                message: format!("Unsupported operation status: {}", status),
+                response: None,
+            };
+            Ok(warp::reply::with_status(
+                warp::reply::json(&response),
+                StatusCode::BAD_REQUEST,
+            ))
+        }
+
+        Err(err) => {
+            let response = ResponseMessage::<Option<uuid::Uuid>> {
+                status: StatusCode::INTERNAL_SERVER_ERROR.as_u16(),
+                message: format!(
+                    "Internal error while deleting connection {}: {}",
+                    conn_id, err
+                ),
+                response: None,
+            };
+            Ok(warp::reply::with_status(
+                warp::reply::json(&response),
+                StatusCode::INTERNAL_SERVER_ERROR,
+            ))
+        }
+    }
+}
+
 /// Handler updates connection
 // PUT /connection
 pub async fn put_connection_handler<N, C>(

src/bin/api/core/http/routes.rs

@@ -164,6 +164,15 @@ where
             .and(with_state(self.state.clone()))
             .and_then(create_connection_handler);
 
+        let delete_connection_route = warp::delete()
+            .and(warp::path("connection"))
+            .and(warp::path::end())
+            .and(auth.clone())
+            .and(warp::query::<ConnQueryParam>())
+            .and(publisher(self.publisher.clone()))
+            .and(with_state(self.state.clone()))
+            .and_then(delete_connection_handler);
+
         let put_connection_route = warp::put()
             .and(warp::path("connection"))
             .and(warp::path::end())
@@ -191,6 +200,7 @@ where
             .or(get_connection_route)
             .or(get_subscription_route)
             .or(post_connection_route)
+            .or(delete_connection_route)
             .or(put_connection_route)
             .recover(rejection);
 

src/bin/api/core/sync/tasks.rs

@@ -45,7 +45,7 @@ where
         conn_id: &uuid::Uuid,
         conn: ConnUpdateRequest,
     ) -> Result<OperationStatus>;
-    async fn delete_connection(&self, conn_id: &uuid::Uuid) -> Result<()>;
+    async fn delete_connection(&self, conn_id: &uuid::Uuid) -> Result<OperationStatus>;
     async fn update_node_status(
         &self,
         uuid: &uuid::Uuid,
@@ -215,17 +215,17 @@ where
             )),
         }
     }
-    async fn delete_connection(&self, conn_id: &uuid::Uuid) -> Result<()> {
+    async fn delete_connection(&self, conn_id: &uuid::Uuid) -> Result<OperationStatus> {
         let mut mem = self.memory.lock().await;
 
         match mem.connections.delete(conn_id) {
             Ok(_) => {
                 self.sync_tx
                     .send(SyncTask::DeleteConn { conn_id: *conn_id })
                     .await?;
-                Ok(())
+                Ok(OperationStatus::Ok(*conn_id))
             }
-            Err(e) => Err(PonyError::Custom(e.to_string())),
+            Err(_) => Ok(OperationStatus::NotFound(*conn_id)),
         }
     }
     async fn update_node_status(

src/config/wireguard.rs

@@ -18,7 +18,7 @@ pub struct WireguardSettings {
 
 impl WireguardSettings {
     pub fn new(config: &WgConfig) -> Self {
-        let (privkey, pubkey) = match (config.pubkey.clone(), config.privkey.clone()) {
+        let (privkey, pubkey) = match (config.privkey.clone(), config.pubkey.clone()) {
             (Some(privkey), Some(pubkey)) => (privkey, pubkey),
             _ => {
                 let keys = WgKeys::default();