diff --git a/crates/defguard/src/main.rs b/crates/defguard/src/main.rs index 0971412a5e..678beec7a6 100644 --- a/crates/defguard/src/main.rs +++ b/crates/defguard/src/main.rs @@ -9,7 +9,10 @@ use defguard_common::{ config::{Command, DefGuardConfig, SERVER_CONFIG}, db::{ init_db, - models::{ActiveWizard, Settings, Wizard, settings::initialize_current_settings}, + models::{ + ActiveWizard, Settings, Wizard, gateway::Gateway, proxy::Proxy, + settings::initialize_current_settings, + }, }, messages::peer_stats_update::PeerStatsUpdate, types::proxy::ProxyControlMessage, @@ -228,6 +231,14 @@ async fn main() -> Result<(), anyhow::Error> { GatewayTxSet::new(gateway_tx.clone(), peer_stats_tx), ); + debug!("Resetting proxy connection state on startup"); + Proxy::mark_all_disconnected(&pool).await?; + debug!("Proxy connection states reset"); + + debug!("Resetting gateway connection state on startup"); + Gateway::mark_all_disconnected(&pool).await?; + debug!("Gateway connection states reset"); + // run services tokio::select! { res = proxy_manager.run() => error!("ProxyManager returned early: {res:?}"), diff --git a/crates/defguard_common/src/db/models/gateway.rs b/crates/defguard_common/src/db/models/gateway.rs index c57c006cb7..e76b1f443b 100644 --- a/crates/defguard_common/src/db/models/gateway.rs +++ b/crates/defguard_common/src/db/models/gateway.rs @@ -72,6 +72,23 @@ impl Gateway { } impl Gateway { + /// Mark all gateways currently considered connected as disconnected. + pub async fn mark_all_disconnected<'e, E>(executor: E) -> sqlx::Result<()> + where + E: PgExecutor<'e>, + { + query( + "UPDATE gateway \ + SET disconnected_at = NOW() \ + WHERE connected_at IS NOT NULL \ + AND (disconnected_at IS NULL OR disconnected_at <= connected_at)", + ) + .execute(executor) + .await?; + + Ok(()) + } + pub async fn find_by_location_id<'e, E>( executor: E, location_id: Id, diff --git a/crates/defguard_common/src/db/models/proxy.rs b/crates/defguard_common/src/db/models/proxy.rs index 011c63efc8..91efead60f 100644 --- a/crates/defguard_common/src/db/models/proxy.rs +++ b/crates/defguard_common/src/db/models/proxy.rs @@ -67,6 +67,23 @@ impl Proxy { } impl Proxy { + /// Mark all proxies currently considered connected as disconnected. + pub async fn mark_all_disconnected<'e, E>(executor: E) -> sqlx::Result<()> + where + E: sqlx::PgExecutor<'e>, + { + sqlx::query( + "UPDATE proxy \ + SET disconnected_at = NOW() \ + WHERE connected_at IS NOT NULL \ + AND (disconnected_at IS NULL OR disconnected_at < connected_at)", + ) + .execute(executor) + .await?; + + Ok(()) + } + /// Fetch all enabled Proxies. pub async fn all_enabled<'e, E>(executor: E) -> sqlx::Result> where diff --git a/crates/defguard_common/src/db/models/wizard.rs b/crates/defguard_common/src/db/models/wizard.rs index cd29d241dd..004be1866e 100644 --- a/crates/defguard_common/src/db/models/wizard.rs +++ b/crates/defguard_common/src/db/models/wizard.rs @@ -4,12 +4,11 @@ use serde::{Deserialize, Serialize}; use sqlx::{FromRow, PgExecutor, Type}; use tracing::info; +use super::setup_auto_adoption::AutoAdoptionWizardStep; use crate::db::models::{ InitialSetupState, InitialSetupStep, setup_auto_adoption::AutoAdoptionWizardState, }; -use super::setup_auto_adoption::AutoAdoptionWizardStep; - /// Which wizard is currently active. Stored as a PostgreSQL enum column. #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Type)] #[sqlx(type_name = "active_wizard", rename_all = "snake_case")] diff --git a/crates/defguard_core/src/enterprise/ldap/tests.rs b/crates/defguard_core/src/enterprise/ldap/tests.rs index d49fae12a7..ccc7d3b1a9 100644 --- a/crates/defguard_core/src/enterprise/ldap/tests.rs +++ b/crates/defguard_core/src/enterprise/ldap/tests.rs @@ -4,7 +4,6 @@ use defguard_common::db::{models::settings::initialize_current_settings, setup_p use ldap3::SearchEntry; use sqlx::postgres::{PgConnectOptions, PgPoolOptions}; -use super::*; use super::{ model::{extract_rdn_value, get_users_without_ldap_path, user_from_searchentry}, sync::{ @@ -12,6 +11,7 @@ use super::{ extract_intersecting_users, }, test_client::{LdapEvent, group_to_test_attrs, user_to_test_attrs}, + *, }; use crate::{ enterprise::{ diff --git a/crates/defguard_setup/src/migration.rs b/crates/defguard_setup/src/migration.rs index 4b452f6df0..90a1ad00ba 100644 --- a/crates/defguard_setup/src/migration.rs +++ b/crates/defguard_setup/src/migration.rs @@ -12,7 +12,12 @@ use axum::{ use axum_extra::extract::cookie::Key; use defguard_common::{VERSION, db::models::Settings, types::proxy::ProxyControlMessage}; use defguard_core::{ + appstate::AppState, auth::failed_login::FailedLoginMap, + db::AppEvent, + enterprise::handlers::openid_login::{auth_callback, get_auth_info}, + events::ApiEvent, + grpc::GatewayEvent, handle_404, handlers::{ auth::{ @@ -37,18 +42,9 @@ use tokio::{ }; use tracing::{info, instrument}; -use defguard_core::{ - appstate::AppState, - db::AppEvent, - enterprise::handlers::openid_login::{auth_callback, get_auth_info}, - events::ApiEvent, - grpc::GatewayEvent, -}; - -use crate::handlers::migration::{get_migration_state, set_general_config, update_migration_state}; use crate::handlers::{ initial_wizard::{create_ca, get_ca, upload_ca}, - migration::finish_setup, + migration::{finish_setup, get_migration_state, set_general_config, update_migration_state}, }; /// FIXME: This is a workaround which enables us to reuse the same API handlers