diff --git a/signal-gateway-bin/src/main.rs b/signal-gateway-bin/src/main.rs index ca66332..f658f2a 100644 --- a/signal-gateway-bin/src/main.rs +++ b/signal-gateway-bin/src/main.rs @@ -244,7 +244,7 @@ limits = [ Duration::from_secs(10) ); assert_eq!(config.gateway.admin_signal_uuids.len(), 2); - assert!(config.gateway.admin_signal_uuids.contains_key("abc-123-uuid")); + assert!(config.gateway.admin_signal_uuids.contains("abc-123-uuid")); let syslog = config.syslog.expect("syslog should be present"); assert_eq!(syslog.listen_addr, "0.0.0.0:1514".parse().unwrap()); diff --git a/signal-gateway/src/gateway/admin_uuids.rs b/signal-gateway/src/gateway/admin_uuids.rs new file mode 100644 index 0000000..ee0f36a --- /dev/null +++ b/signal-gateway/src/gateway/admin_uuids.rs @@ -0,0 +1,169 @@ +//! Admin Signal UUIDs container with flexible deserialization. +//! +//! Supports two formats: +//! - Map: `{"uuid1": ["safety1", "safety2"], "uuid2": []}` +//! - Sequence: `["uuid1", "uuid2"]` (treated as UUIDs with no safety numbers) + +use serde::de::{MapAccess, SeqAccess, Visitor}; +use serde::{Deserialize, Deserializer}; +use std::collections::HashMap; +use std::fmt; + +/// Container for admin Signal UUIDs mapped to their optional safety numbers. +/// +/// Can be deserialized from either: +/// - A map of UUID -> safety numbers: `{"uuid1": ["12345..."], "uuid2": []}` +/// - A sequence of UUIDs (no safety numbers): `["uuid1", "uuid2"]` +#[derive(Clone, Debug, Default)] +pub struct AdminSignalUuids { + map: HashMap>, +} + +impl AdminSignalUuids { + /// Create an empty container. + pub fn new() -> Self { + Self::default() + } + + /// Check if a UUID is a registered admin. + pub fn contains(&self, uuid: &str) -> bool { + self.map.contains_key(uuid) + } + + /// Get all admin UUIDs. + pub fn uuids(&self) -> impl Iterator { + self.map.keys() + } + + /// Get the number of admin UUIDs. + pub fn len(&self) -> usize { + self.map.len() + } + + /// Check if empty. + pub fn is_empty(&self) -> bool { + self.map.is_empty() + } + + /// Iterate over UUID and safety number pairs. + pub fn iter(&self) -> impl Iterator)> { + self.map.iter() + } + + /// Get safety numbers for a specific UUID. + pub fn get(&self, uuid: &str) -> Option<&Vec> { + self.map.get(uuid) + } +} + +impl<'de> Deserialize<'de> for AdminSignalUuids { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + deserializer.deserialize_any(AdminSignalUuidsVisitor) + } +} + +struct AdminSignalUuidsVisitor; + +impl<'de> Visitor<'de> for AdminSignalUuidsVisitor { + type Value = AdminSignalUuids; + + fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result { + formatter.write_str("a map of UUIDs to safety numbers, or a sequence of UUIDs") + } + + fn visit_map(self, mut access: M) -> Result + where + M: MapAccess<'de>, + { + let mut map = HashMap::with_capacity(access.size_hint().unwrap_or(0)); + while let Some((key, value)) = access.next_entry::>()? { + map.insert(key, value); + } + Ok(AdminSignalUuids { map }) + } + + fn visit_seq(self, mut access: S) -> Result + where + S: SeqAccess<'de>, + { + let mut map = HashMap::with_capacity(access.size_hint().unwrap_or(0)); + while let Some(uuid) = access.next_element::()? { + map.insert(uuid, Vec::new()); + } + Ok(AdminSignalUuids { map }) + } +} + +impl FromIterator for AdminSignalUuids { + fn from_iter>(iter: I) -> Self { + Self { + map: iter.into_iter().map(|uuid| (uuid, Vec::new())).collect(), + } + } +} + +impl FromIterator<(String, Vec)> for AdminSignalUuids { + fn from_iter)>>(iter: I) -> Self { + Self { + map: iter.into_iter().collect(), + } + } +} + +impl<'a> IntoIterator for &'a AdminSignalUuids { + type Item = (&'a String, &'a Vec); + type IntoIter = std::collections::hash_map::Iter<'a, String, Vec>; + + fn into_iter(self) -> Self::IntoIter { + self.map.iter() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_deserialize_map() { + let json = r#"{"uuid1": ["safety1", "safety2"], "uuid2": []}"#; + let uuids: AdminSignalUuids = serde_json::from_str(json).unwrap(); + + assert_eq!(uuids.len(), 2); + assert!(uuids.contains("uuid1")); + assert!(uuids.contains("uuid2")); + assert_eq!(uuids.get("uuid1").unwrap(), &vec!["safety1", "safety2"]); + assert_eq!(uuids.get("uuid2").unwrap(), &Vec::::new()); + } + + #[test] + fn test_deserialize_seq() { + let json = r#"["uuid1", "uuid2", "uuid3"]"#; + let uuids: AdminSignalUuids = serde_json::from_str(json).unwrap(); + + assert_eq!(uuids.len(), 3); + assert!(uuids.contains("uuid1")); + assert!(uuids.contains("uuid2")); + assert!(uuids.contains("uuid3")); + // All should have empty safety numbers + assert!(uuids.get("uuid1").unwrap().is_empty()); + assert!(uuids.get("uuid2").unwrap().is_empty()); + assert!(uuids.get("uuid3").unwrap().is_empty()); + } + + #[test] + fn test_empty_map() { + let json = r#"{}"#; + let uuids: AdminSignalUuids = serde_json::from_str(json).unwrap(); + assert!(uuids.is_empty()); + } + + #[test] + fn test_empty_seq() { + let json = r#"[]"#; + let uuids: AdminSignalUuids = serde_json::from_str(json).unwrap(); + assert!(uuids.is_empty()); + } +} diff --git a/signal-gateway/src/gateway/mod.rs b/signal-gateway/src/gateway/mod.rs index 470670b..d58ce2a 100644 --- a/signal-gateway/src/gateway/mod.rs +++ b/signal-gateway/src/gateway/mod.rs @@ -20,14 +20,7 @@ use http::{Method, Request, Response, StatusCode}; use http_body::Body; use http_body_util::BodyExt; use prometheus_http_client::{AlertStatus, ExtractLabels}; -use std::{ - collections::HashMap, - fmt::Write, - net::SocketAddr, - path::PathBuf, - sync::Mutex, - time::Duration, -}; +use std::{fmt::Write, net::SocketAddr, path::PathBuf, sync::Mutex, time::Duration}; use tokio::{ join, sync::mpsc::{UnboundedReceiver, UnboundedSender, unbounded_channel}, @@ -35,6 +28,9 @@ use tokio::{ use tokio_util::{bytes::Buf, sync::CancellationToken}; use tracing::{debug, error, info, warn}; +mod admin_uuids; +pub use admin_uuids::AdminSignalUuids; + mod log_buffer; mod log_handler; use log_handler::{LogHandler, LogHandlerConfig}; @@ -67,9 +63,9 @@ pub struct GatewayConfig { #[conf(long, env, default_value = "5s", value_parser = conf_extra::parse_duration, serde(use_value_parser))] pub signal_cli_retry_delay: Duration, /// Admin UUIDs mapped to their safety numbers (can be empty). - /// Example: `{"uuid1": ["12345...", "67890..."], "uuid2": []}` + /// Accepts either a map `{"uuid1": ["12345..."], "uuid2": []}` or a list `["uuid1", "uuid2"]`. #[conf(long, env, value_parser = serde_json::from_str)] - pub admin_signal_uuids: HashMap>, + pub admin_signal_uuids: AdminSignalUuids, /// If set, alerts are sent to this group instead of individual admins. #[conf(long, env)] pub alert_group_id: Option, @@ -84,12 +80,12 @@ pub struct GatewayConfig { impl GatewayConfig { /// Check if a UUID is a registered admin pub fn is_admin(&self, uuid: &str) -> bool { - self.admin_signal_uuids.contains_key(uuid) + self.admin_signal_uuids.contains(uuid) } /// Get all admin UUIDs pub fn admin_uuids(&self) -> Vec { - self.admin_signal_uuids.keys().cloned().collect() + self.admin_signal_uuids.uuids().cloned().collect() } }