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, } impl<'a> LoadBalancerConfig<'a> { pub fn new(opnsense: &'a mut OPNsense, opnsense_shell: Arc) -> Self { Self { opnsense, opnsense_shell, } } pub fn get_full_config(&self) -> &Option { &self.opnsense.opnsense.haproxy } fn with_haproxy(&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, healthcheck: Option, ) { 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, healthcheck: Option, ) { 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 { 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 { 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())); } }