backend: remove the database fields for WgKeys in use

Author: ffreddow

backend/src/lib.rs

@@ -308,6 +308,27 @@ pub(crate) mod tests {
         web::Data::new(state)
     }
 
+    pub async fn mark_keys_as_in_use(state: &web::Data<AppState>, key_ids: Vec<uuid::Uuid>) {
+        let mut keys_in_use = state.keys_in_use.lock().await;
+        for key_id in key_ids {
+            keys_in_use.insert(key_id);
+        }
+    }
+
+    pub async fn get_charger_key_ids(
+        state: &web::Data<AppState>,
+        charger_uuid: uuid::Uuid,
+    ) -> Vec<uuid::Uuid> {
+        use db_connector::schema::wg_keys::dsl::*;
+
+        let mut conn = state.pool.get().unwrap();
+        wg_keys
+            .filter(charger_id.eq(charger_uuid))
+            .select(id)
+            .load::<uuid::Uuid>(&mut conn)
+            .unwrap_or_default()
+    }
+
     pub fn configure(cfg: &mut ServiceConfig) {
         let pool = db_connector::test_connection_pool();
 

backend/src/main.rs

@@ -19,15 +19,20 @@
 
 mod monitoring;
 
-use std::{collections::{HashMap, HashSet}, net::UdpSocket, num::NonZeroUsize, sync::Arc, time::Duration};
+use std::{
+    collections::{HashMap, HashSet},
+    net::UdpSocket,
+    num::NonZeroUsize,
+    sync::Arc,
+    time::Duration,
+};
 
 use actix::Arbiter;
 pub use backend::*;
 use backend::{rate_limit::IPRateLimiter, utils::get_connection};
 
 use actix_web::{middleware::Logger, web, App, HttpServer};
-use db_connector::{get_connection_pool, run_migrations, Pool};
-use diesel::prelude::*;
+use db_connector::{get_connection_pool, run_migrations};
 use futures_util::lock::Mutex;
 use lettre::{transport::smtp::authentication::Credentials, SmtpTransport};
 use lru::LruCache;
@@ -40,16 +45,6 @@ use udp_server::packet::{
     ManagementPacketHeader,
 };
 
-fn reset_wg_keys(pool: &Pool) {
-    use db_connector::schema::wg_keys::dsl::*;
-
-    let mut conn = pool.get().unwrap();
-    diesel::update(wg_keys)
-        .set(in_use.eq(false))
-        .execute(&mut conn)
-        .unwrap();
-}
-
 fn cleanup_thread(state: web::Data<AppState>) {
     loop {
         std::thread::sleep(Duration::from_secs(60));
@@ -129,8 +124,6 @@ async fn main() -> std::io::Result<()> {
     let mut conn = pool.get().expect("Failed to get connection from pool");
     run_migrations(&mut conn).expect("Failed to run migrations");
 
-    reset_wg_keys(&pool);
-
     let mailer = {
         let email = std::env::var("EMAIL_USER").expect("EMAIL_USER must be set");
         let pass = std::env::var("EMAIL_PASS").expect("EMAIL_PASS must be set");

backend/src/routes/charger/add.rs

@@ -396,7 +396,6 @@ async fn add_wg_key(
         id: uuid::Uuid::new_v4(),
         user_id: uid,
         charger_id: cid,
-        in_use: false,
         charger_pub: keys.charger_public,
         web_private: keys.web_private,
         psk: keys.psk,

backend/src/routes/charger/allow_user.rs

@@ -79,7 +79,6 @@ async fn add_keys(
                 id: uuid::Uuid::new_v4(),
                 user_id: uid,
                 charger_id: cid,
-                in_use: false,
                 charger_pub: key.charger_public,
                 web_private: key.web_private,
                 psk: key.psk,

backend/src/routes/charger/get_key.rs

@@ -93,19 +93,21 @@ pub async fn get_key(
         })
         .await?;
 
-        // Count how many are in use
-        charger_key_ids.iter().filter(|key_id| keys_in_use_cache.contains(*key_id)).count()
+        charger_key_ids
+            .into_iter()
+            .filter(|key_id| keys_in_use_cache.contains(key_id))
+            .collect::<Vec<_>>()
     };
 
-    if keys_in_use_count >= 5 {
+    if keys_in_use_count.len() >= 5 {
         return Err(Error::AllKeysInUse.into());
     }
 
     let mut conn = get_connection(&state)?;
     let key: Option<WgKey> = web_block_unpacked(move || {
         match WgKey::belonging_to(&user)
             .filter(charger_id.eq(&cid))
-            .filter(in_use.eq(false))
+            .filter(id.ne_all(keys_in_use_count))
             .select(WgKey::as_select())
             .get_result(&mut conn)
         {
@@ -134,15 +136,16 @@ pub async fn get_key(
 
 #[cfg(test)]
 mod tests {
-    use std::str::FromStr;
-
     use super::*;
     use actix_web::{cookie::Cookie, test, App};
-    use db_connector::test_connection_pool;
     use rand::TryRngCore;
     use rand_core::OsRng;
 
-    use crate::{middleware::jwt::JwtMiddleware, routes::user::tests::TestUser, tests::configure};
+    use crate::{
+        middleware::jwt::JwtMiddleware,
+        routes::user::tests::TestUser,
+        tests::{configure, get_charger_key_ids, mark_keys_as_in_use},
+    };
 
     #[actix_web::test]
     async fn test_get_key() {
@@ -168,26 +171,23 @@ mod tests {
     }
 
     #[actix_web::test]
-    async fn test_get_key_none_left() {
-        use db_connector::schema::wg_keys::dsl::*;
+    async fn test_get_key_all_keys_in_use() {
+        use std::str::FromStr;
 
         let (mut user, _) = TestUser::random().await;
         user.login().await;
 
         let charger_uid = OsRng.try_next_u32().unwrap() as i32;
         let charger = user.add_charger(charger_uid).await;
 
-        let pool = test_connection_pool();
-        let mut conn = pool.get().unwrap();
+        let state = crate::tests::create_test_state(None);
+        let charger_uuid = uuid::Uuid::from_str(&charger.uuid).unwrap();
+        let key_ids = get_charger_key_ids(&state, charger_uuid).await;
 
-        diesel::update(wg_keys)
-            .filter(charger_id.eq(uuid::Uuid::from_str(&charger.uuid).unwrap()))
-            .set(in_use.eq(true))
-            .execute(&mut conn)
-            .unwrap();
+        mark_keys_as_in_use(&state, key_ids).await;
 
         let app = App::new()
-            .configure(configure)
+            .app_data(state)
             .wrap(JwtMiddleware)
             .service(get_key);
         let app = test::init_service(app).await;
@@ -197,6 +197,7 @@ mod tests {
             .cookie(Cookie::new("access_token", user.get_access_token()))
             .to_request();
 
+        // Should fail with 404 since all keys are in use
         let resp = test::call_service(&app, req).await;
         assert!(resp.status().is_client_error());
         assert_eq!(resp.status().as_u16(), 404);

backend/src/ws_udp_bridge.rs

@@ -156,27 +156,13 @@ impl WebClient {
     }
 
     pub async fn stop(self) {
-        use db_connector::schema::wg_keys::dsl::*;
-
         log::debug!("Closed connection to charger '{}'", self.charger_id);
 
         {
             let mut keys_in_use = self.app_state.keys_in_use.lock().await;
             keys_in_use.remove(&self.key_id);
         }
 
-        let mut conn = match self.app_state.pool.get() {
-            Ok(conn) => conn,
-            Err(_err) => {
-                log::error!(
-                    "Failed to release connection '{}' for charger '{}'",
-                    self.key_id,
-                    self.charger_id
-                );
-                return;
-            }
-        };
-
         let meta = RemoteConnMeta {
             charger_id: self.charger_id,
             conn_no: self.conn_no,
@@ -226,21 +212,6 @@ impl WebClient {
             }
         }
 
-        match diesel::update(wg_keys)
-            .filter(id.eq(self.key_id))
-            .set(in_use.eq(false))
-            .execute(&mut conn)
-        {
-            Ok(_) => (),
-            Err(_err) => {
-                log::error!(
-                    "Failed to release connection '{}' for charger '{}'",
-                    self.key_id,
-                    self.charger_id
-                );
-            }
-        }
-
         self.session.close(None).await.ok();
     }
 }
@@ -295,8 +266,7 @@ async fn start_ws(
 ) -> Result<HttpResponse, actix_web::Error> {
     use db_connector::schema::wg_keys::dsl as wg_keys;
 
-    let key_uuid = uuid::Uuid::from_str(&key_id.key_id)
-        .map_err(|_| Error::WgKeysDoNotExist)?;
+    let key_uuid = uuid::Uuid::from_str(&key_id.key_id).map_err(|_| Error::WgKeysDoNotExist)?;
 
     let mut keys_in_use = state.keys_in_use.lock().await;
     if keys_in_use.contains(&key_uuid) {
@@ -347,20 +317,6 @@ async fn start_ws(
 
     if resp.status() == 101 {
         keys_in_use.insert(keys.id);
-
-        let mut conn = get_connection(&state)?;
-        use db_connector::schema::wg_keys::dsl::*;
-        web_block_unpacked(move || {
-            if let Err(_err) = diesel::update(wg_keys)
-                .filter(id.eq(&keys.id))
-                .set(in_use.eq(true))
-                .execute(&mut conn)
-            {
-                return Err(Error::InternalError);
-            }
-            Ok(())
-        })
-        .await?;
     }
     drop(keys_in_use);
 
@@ -408,8 +364,18 @@ async fn start_ws(
 
 #[cfg(test)]
 mod tests {
-    // #[actix_web::test]
-    // async fn test_connecting_ws() {
+    use super::*;
+    use uuid::Uuid;
 
-    // }
+    #[test]
+    fn test_validate_key_id_valid() {
+        let valid_uuid = Uuid::new_v4().to_string();
+        assert!(validate_key_id(&valid_uuid).is_ok());
+    }
+
+    #[test]
+    fn test_validate_key_id_invalid() {
+        let invalid_uuid = "not-a-uuid";
+        assert!(validate_key_id(invalid_uuid).is_err());
+    }
 }

db_connector/migrations/2025-10-20-131714_wg_keys_drop_in_use/down.sql

@@ -0,0 +1,3 @@
+-- This file should undo anything in `up.sql`
+
+ALTER TABLE "wg_keys" ADD COLUMN "in_use" BOOLEAN NOT NULL DEFAULT false;

db_connector/migrations/2025-10-20-131714_wg_keys_drop_in_use/up.sql

@@ -0,0 +1,3 @@
+-- Your SQL goes here
+
+ALTER TABLE "wg_keys" DROP COLUMN "in_use";

db_connector/src/models/wg_keys.rs

@@ -19,7 +19,6 @@ pub struct WgKey {
     pub id: uuid::Uuid,
     pub user_id: uuid::Uuid,
     pub charger_id: uuid::Uuid,
-    pub in_use: bool,
     pub charger_pub: String,
     pub web_private: Vec<u8>,
     pub psk: Vec<u8>,

db_connector/src/schema.rs

@@ -90,7 +90,6 @@ diesel::table! {
     wg_keys (id) {
         id -> Uuid,
         user_id -> Uuid,
-        in_use -> Bool,
         charger_pub -> Varchar,
         web_private -> Bytea,
         psk -> Bytea,