Fix e2e lockfile and tests (#945)

* revert e2e lockfile to old version

* remove global state from tests

* thread local

* generalize global values

* sort imports, cleanup

src/db/models/settings.rs

@@ -1,40 +1,23 @@
-use std::{
-    collections::HashMap,
-    str::FromStr,
-    sync::{RwLock, RwLockReadGuard},
-};
+use std::{collections::HashMap, str::FromStr};
 
 use sqlx::{query, query_as, PgExecutor, PgPool, Type};
 use struct_patch::Patch;
 use thiserror::Error;
 
-use crate::secret::SecretString;
+use crate::{global_value, secret::SecretString};
 
-// wrap in `Option` since a static cannot be initialized with a non-const function
-static SETTINGS: RwLock<Option<Settings>> = RwLock::new(None);
-
-pub(crate) fn set_settings(new_settings: Settings) {
-    *SETTINGS
-        .write()
-        .expect("Failed to acquire lock on current settings.") = Some(new_settings);
-}
-
-pub(crate) fn get_settings() -> RwLockReadGuard<'static, Option<Settings>> {
-    SETTINGS
-        .read()
-        .expect("Failed to acquire lock on current settings.")
-}
+global_value!(SETTINGS, Option<Settings>, None, set_settings, get_settings);
 
 /// Initializes global `SETTINGS` struct at program startup
 pub async fn initialize_current_settings(pool: &PgPool) -> Result<(), sqlx::Error> {
     debug!("Initializing global settings strut");
     match Settings::get(pool).await? {
         Some(settings) => {
-            set_settings(settings);
+            set_settings(Some(settings));
         }
         None => {
             debug!("Settings not found in DB. Using default values to initialize global settings struct");
-            set_settings(Settings::default());
+            set_settings(Some(Settings::default()));
         }
     }
     Ok(())
@@ -47,7 +30,7 @@ pub async fn update_current_settings(
 ) -> Result<(), sqlx::Error> {
     debug!("Updating current settings to: {new_settings:?}");
     new_settings.save(pool).await?;
-    set_settings(new_settings);
+    set_settings(Some(new_settings));
     Ok(())
 }
 

src/enterprise/directory_sync/google.rs

@@ -1,11 +1,12 @@
 use std::{str::FromStr, time::Duration};
 
-use super::{DirectoryGroup, DirectorySync, DirectorySyncError, DirectoryUser};
 use chrono::{DateTime, TimeDelta, Utc};
 #[cfg(not(test))]
 use jsonwebtoken::{encode, Algorithm, EncodingKey, Header};
 use reqwest::{header::AUTHORIZATION, Url};
 
+use super::{DirectoryGroup, DirectorySync, DirectorySyncError, DirectoryUser};
+
 #[cfg(not(test))]
 const SCOPES: &str = "openid email profile https://www.googleapis.com/auth/admin.directory.customer.readonly https://www.googleapis.com/auth/admin.directory.group.readonly https://www.googleapis.com/auth/admin.directory.user.readonly";
 const ACCESS_TOKEN_URL: &str = "https://oauth2.googleapis.com/token";

src/enterprise/directory_sync/mod.rs

@@ -1,17 +1,16 @@
 use std::collections::{HashMap, HashSet};
 
-use sqlx::PgPool;
-
-use crate::{
-    db::{Group, Id, User},
-    enterprise::db::models::openid_provider::DirectorySyncUserBehavior,
-};
 use sqlx::error::Error as SqlxError;
+use sqlx::PgPool;
 use thiserror::Error;
 
 use super::db::models::openid_provider::{DirectorySyncTarget, OpenIdProvider};
 #[cfg(not(test))]
 use super::is_enterprise_enabled;
+use crate::{
+    db::{Group, Id, User},
+    enterprise::db::models::openid_provider::DirectorySyncUserBehavior,
+};
 
 #[derive(Debug, Error)]
 pub enum DirectorySyncError {
@@ -600,12 +599,13 @@ pub(crate) async fn do_directory_sync(pool: &PgPool) -> Result<(), DirectorySync
 
 #[cfg(test)]
 mod test {
+    use secrecy::ExposeSecret;
+
     use super::*;
     use crate::{
         config::DefGuardConfig, enterprise::db::models::openid_provider::DirectorySyncTarget,
         SERVER_CONFIG,
     };
-    use secrecy::ExposeSecret;
 
     async fn make_test_provider(
         pool: &PgPool,

src/enterprise/license.rs

@@ -1,7 +1,4 @@
-use std::{
-    sync::{RwLock, RwLockReadGuard},
-    time::Duration,
-};
+use std::time::Duration;
 
 use anyhow::Result;
 use base64::prelude::*;
@@ -13,28 +10,21 @@ use sqlx::{error::Error as SqlxError, PgPool};
 use thiserror::Error;
 use tokio::time::sleep;
 
+use super::limits::Counts;
 use crate::{
     db::{models::settings::update_current_settings, Settings},
-    server_config, VERSION,
+    global_value, server_config, VERSION,
 };
 
-use super::limits::Counts;
-
 const LICENSE_SERVER_URL: &str = "https://pkgs.defguard.net/api/license/renew";
 
-static LICENSE: RwLock<Option<License>> = RwLock::new(None);
-
-pub fn set_cached_license(license: Option<License>) {
-    *LICENSE
-        .write()
-        .expect("Failed to acquire lock on the license mutex.") = license;
-}
-
-pub fn get_cached_license() -> RwLockReadGuard<'static, Option<License>> {
-    LICENSE
-        .read()
-        .expect("Failed to acquire lock on the license mutex.")
-}
+global_value!(
+    LICENSE,
+    Option<License>,
+    None,
+    set_cached_license,
+    get_cached_license
+);
 
 tonic::include_proto!("license");
 
@@ -584,7 +574,7 @@ pub async fn run_periodic_license_check(pool: &PgPool) -> Result<(), LicenseErro
             let license = get_cached_license();
             debug!("Checking if the license {license:?} requires a renewal...");
 
-            if let Some(license) = &*license {
+            if let Some(license) = license.as_ref() {
                 if license.requires_renewal() {
                     // check if we are pass the maximum expiration date, after which we don't
                     // want to try to renew the license anymore

src/enterprise/limits.rs

@@ -1,39 +1,24 @@
 use sqlx::{error::Error as SqlxError, query, PgPool};
-use std::sync::{RwLock, RwLockReadGuard};
 
 #[cfg(test)]
 use super::license::get_cached_license;
 use super::license::License;
+use crate::global_value;
 
 // Limits for free users
 pub const DEFAULT_USERS_LIMIT: u32 = 5;
 pub const DEFAULT_DEVICES_LIMIT: u32 = 10;
 pub const DEFAULT_LOCATIONS_LIMIT: u32 = 1;
 
-#[derive(Debug, Default)]
-pub(crate) struct Counts {
+#[derive(Debug)]
+#[cfg_attr(test, derive(Clone))]
+pub struct Counts {
     user: u32,
     device: u32,
     wireguard_network: u32,
 }
 
-static COUNTS: RwLock<Counts> = RwLock::new(Counts {
-    user: 0,
-    device: 0,
-    wireguard_network: 0,
-});
-
-fn set_counts(new_counts: Counts) {
-    *COUNTS
-        .write()
-        .expect("Failed to acquire lock on the enterprise limit counts.") = new_counts;
-}
-
-pub(crate) fn get_counts() -> RwLockReadGuard<'static, Counts> {
-    COUNTS
-        .read()
-        .expect("Failed to acquire lock on the enterprise limit counts.")
-}
+global_value!(COUNTS, Counts, Counts::default(), set_counts, get_counts);
 
 /// Update the counts of users, devices, and wireguard networks stored in the memory.
 // TODO: Use it with database triggers when they are implemented
@@ -80,6 +65,14 @@ pub async fn do_count_update(pool: &PgPool) -> Result<(), SqlxError> {
 }
 
 impl Counts {
+    pub(crate) const fn default() -> Self {
+        Self {
+            user: 0,
+            device: 0,
+            wireguard_network: 0,
+        }
+    }
+
     #[cfg(test)]
     pub(crate) fn new(user: u32, device: u32, wireguard_network: u32) -> Self {
         Self {
@@ -97,7 +90,7 @@ impl Counts {
         let maybe_license = get_cached_license();
 
         // validate limits against license if available, use defaults otherwise
-        match &*maybe_license {
+        match &maybe_license {
             Some(license) => {
                 debug!("Cached license found. Validating license limits...");
                 self.is_over_license_limits(license)
@@ -177,9 +170,8 @@ impl LimitsExceeded {
 mod test {
     use chrono::{TimeDelta, Utc};
 
-    use crate::enterprise::license::{set_cached_license, License, LicenseLimits};
-
     use super::*;
+    use crate::enterprise::license::{set_cached_license, License, LicenseLimits};
 
     #[test]
     fn test_counts() {

src/globals.rs

@@ -0,0 +1,54 @@
+#[macro_export]
+/// Specify a global value that can be accessed from anywhere in the application.
+/// Positional arguments:
+/// - `$name`: The name of the global value. This will be the name of the variable that holds the value.
+/// - `$type`: The type of the global value.
+/// - `$init`: The initial value of the global value.
+/// - `$set_fn`: The name of the function that will be used to set the global value.
+/// - `$get_fn`: The name of the function that will be used to get the global value.
+///
+/// The macro will also automatically generate boilerplate code for unit tests to work correctly.
+macro_rules! global_value {
+    ($name:ident, $type:ty, $init:expr, $set_fn:ident, $get_fn:ident) => {
+        use std::sync::RwLock;
+        #[cfg(not(test))]
+        use std::sync::RwLockReadGuard;
+
+        #[cfg(test)]
+        thread_local! {
+            static $name: RwLock<$type> = const { RwLock::new($init) };
+        }
+
+        #[cfg(not(test))]
+        static $name: RwLock<$type> = RwLock::new($init);
+
+        #[cfg(not(test))]
+        pub fn $set_fn(value: $type) {
+            *$name.write().expect("Failed to acquire lock on the mutex.") = value;
+        }
+
+        #[cfg(not(test))]
+        pub fn $get_fn() -> RwLockReadGuard<'static, $type> {
+            $name.read().expect("Failed to acquire lock on the mutex.")
+        }
+
+        #[cfg(test)]
+        pub fn $set_fn(new_value: $type) {
+            $name.with(|value| {
+                *value.write().expect("Failed to acquire lock on the mutex.") = new_value;
+            });
+        }
+
+        // This is not really a 1:1 replacement for the non-test RwLockReadGuard<'static, $type> as the RwLock may be tried to be
+        // dereferenced
+        #[cfg(test)]
+        pub fn $get_fn() -> $type {
+            $name.with(|value| {
+                value
+                    .read()
+                    .expect("Failed to acquire lock on the mutex.")
+                    .clone()
+            })
+        }
+    };
+}

src/lib.rs

@@ -130,6 +130,7 @@ pub mod config;
 pub mod db;
 pub mod enterprise;
 mod error;
+pub mod globals;
 pub mod grpc;
 pub mod handlers;
 pub mod headers;

src/templates.rs

@@ -263,7 +263,6 @@ pub fn gateway_reconnected_mail(
     context.insert("gateway_ip", gateway_ip);
     context.insert("network_name", network_name);
     tera.add_raw_template("mail_gateway_reconnected", MAIL_GATEWAY_RECONNECTED)?;
-    println!("dupa");
     Ok(tera.render("mail_gateway_reconnected", &context)?)
 }
 

src/updates.rs

@@ -1,16 +1,16 @@
-use std::{
-    env,
-    sync::{RwLock, RwLockReadGuard},
-};
+use std::env;
 
 use chrono::NaiveDate;
 use semver::Version;
 
+use crate::global_value;
+
 const PRODUCT_NAME: &str = "Defguard";
 const UPDATES_URL: &str = "https://update-service-dev.defguard.net/api/update/check";
 const VERSION: &str = env!("CARGO_PKG_VERSION");
 
 #[derive(Deserialize, Debug, Serialize)]
+#[cfg_attr(test, derive(Clone))]
 pub struct Update {
     version: String,
     release_date: NaiveDate,
@@ -20,19 +20,7 @@ pub struct Update {
     notes: String,
 }
 
-static NEW_UPDATE: RwLock<Option<Update>> = RwLock::new(None);
-
-fn set_update(update: Update) {
-    *NEW_UPDATE
-        .write()
-        .expect("Failed to acquire lock on the update.") = Some(update);
-}
-
-pub fn get_update() -> RwLockReadGuard<'static, Option<Update>> {
-    NEW_UPDATE
-        .read()
-        .expect("Failed to acquire lock on the update.")
-}
+global_value!(NEW_UPDATE, Option<Update>, None, set_update, get_update);
 
 async fn fetch_update() -> Result<Update, anyhow::Error> {
     let body = serde_json::json!({
@@ -63,7 +51,7 @@ pub(crate) async fn do_new_version_check() -> Result<(), anyhow::Error> {
                 update.version, update.release_date
             );
         }
-        set_update(update);
+        set_update(Some(update));
     } else {
         debug!("New version check done. You are using the latest version of Defguard.");
     }