harmony/opnsense-config/src/modules/load_balancer.rs

use std::{collections::HashSet, sync::Arc};

use log::warn;
use opnsense_config_xml::{
    Frontend, HAProxy, HAProxyBackend, HAProxyHealthCheck, HAProxyServer, OPNsense,
};

use crate::{config::OPNsenseShell, Error};

pub struct LoadBalancerConfig<'a> {
    opnsense: &'a mut OPNsense,
    opnsense_shell: Arc<dyn OPNsenseShell>,
}

impl<'a> LoadBalancerConfig<'a> {
    pub fn new(opnsense: &'a mut OPNsense, opnsense_shell: Arc<dyn OPNsenseShell>) -> Self {
        Self {
            opnsense,
            opnsense_shell,
        }
    }

    pub fn get_full_config(&self) -> &Option<HAProxy> {
        &self.opnsense.opnsense.haproxy
    }

    fn with_haproxy<F, R>(&mut self, f: F) -> R
    where
        F: FnOnce(&mut HAProxy) -> R,
    {
        match &mut self.opnsense.opnsense.haproxy.as_mut() {
            Some(haproxy) => f(haproxy),
            None => unimplemented!(
                "Adding a backend is not supported when haproxy config does not exist yet"
            ),
        }
    }

    pub fn enable(&mut self, enabled: bool) {
        self.with_haproxy(|haproxy| haproxy.general.enabled = enabled as i32);
    }

    /// Configures a service by removing any existing service on the same port
    /// and then adding the new definition. This ensures idempotency.
    pub fn configure_service(
        &mut self,
        frontend: Frontend,
        backend: HAProxyBackend,
        servers: Vec<HAProxyServer>,
        healthcheck: Option<HAProxyHealthCheck>,
    ) {
        self.remove_service_by_bind_address(&frontend.bind);
        self.add_new_service(frontend, backend, servers, healthcheck);
    }

    /// Removes a service and its dependent components based on the frontend's bind address.
    /// This performs a cascading delete of the frontend, backend, servers, and health check.
    fn remove_service_by_bind_address(&mut self, bind_address: &str) {
        self.with_haproxy(|haproxy| {
            let Some(old_frontend) = remove_frontend_by_bind_address(haproxy, bind_address) else {
                return;
            };

            let Some(old_backend) = remove_backend(haproxy, old_frontend) else {
                return;
            };

            remove_healthcheck(haproxy, &old_backend);
            remove_servers(haproxy, &old_backend);
        });
    }

    /// Adds the components of a new service to the HAProxy configuration.
    fn add_new_service(
        &mut self,
        frontend: Frontend,
        backend: HAProxyBackend,
        servers: Vec<HAProxyServer>,
        healthcheck: Option<HAProxyHealthCheck>,
    ) {
        self.with_haproxy(|haproxy| {
            if let Some(check) = healthcheck {
                haproxy.healthchecks.healthchecks.push(check);
            }
            haproxy.servers.servers.extend(servers);
            haproxy.backends.backends.push(backend);
            haproxy.frontends.frontend.push(frontend);
        });
    }

    pub async fn reload_restart(&self) -> Result<(), Error> {
        self.opnsense_shell.exec("configctl haproxy stop").await?;
        self.opnsense_shell
            .exec("configctl template reload OPNsense/HAProxy")
            .await?;
        self.opnsense_shell
            .exec("configctl template reload OPNsense/Syslog")
            .await?;
        self.opnsense_shell
            .exec("/usr/local/sbin/haproxy -c -f /usr/local/etc/haproxy.conf.staging")
            .await?;

        // This script copies the staging config to production config. I am not 100% sure it is
        // required in the context
        self.opnsense_shell
            .exec("/usr/local/opnsense/scripts/OPNsense/HAProxy/setup.sh deploy")
            .await?;

        self.opnsense_shell
            .exec("configctl haproxy configtest")
            .await?;
        self.opnsense_shell.exec("configctl haproxy start").await?;
        Ok(())
    }
}

fn remove_frontend_by_bind_address(haproxy: &mut HAProxy, bind_address: &str) -> Option<Frontend> {
    let pos = haproxy
        .frontends
        .frontend
        .iter()
        .position(|f| f.bind == bind_address);

    match pos {
        Some(pos) => Some(haproxy.frontends.frontend.remove(pos)),
        None => None,
    }
}

fn remove_backend(haproxy: &mut HAProxy, old_frontend: Frontend) -> Option<HAProxyBackend> {
    let pos = haproxy
        .backends
        .backends
        .iter()
        .position(|b| b.uuid == old_frontend.default_backend);

    match pos {
        Some(pos) => Some(haproxy.backends.backends.remove(pos)),
        None => None, // orphaned frontend, shouldn't happen
    }
}

fn remove_healthcheck(haproxy: &mut HAProxy, backend: &HAProxyBackend) {
    if let Some(uuid) = &backend.health_check.content {
        haproxy
            .healthchecks
            .healthchecks
            .retain(|h| h.uuid != *uuid);
    }
}

/// Remove the backend's servers. This assumes servers are not shared between services.
fn remove_servers(haproxy: &mut HAProxy, backend: &HAProxyBackend) {
    if let Some(server_uuids_str) = &backend.linked_servers.content {
        let server_uuids_to_remove: HashSet<_> = server_uuids_str.split(',').collect();
        haproxy
            .servers
            .servers
            .retain(|s| !server_uuids_to_remove.contains(s.uuid.as_str()));
    }
}