diff --git a/Cargo.lock b/Cargo.lock index 771edb5..cccce54 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1673,6 +1673,7 @@ name = "signal-gateway-bin" version = "0.1.0" dependencies = [ "conf", + "conf-extra", "dotenvy", "hyper", "hyper-util", diff --git a/prometheus-http-client/src/plot.rs b/prometheus-http-client/src/plot.rs index 05dc5a0..84af566 100644 --- a/prometheus-http-client/src/plot.rs +++ b/prometheus-http-client/src/plot.rs @@ -192,7 +192,7 @@ impl PlotStyle { path: impl AsRef, mts: &[MetricTimeseries], plot_threshold: Option, - ) -> Result<(), Box> + ) -> Result<(), Box> where KV: Clone + Debug, K: Display, diff --git a/signal-gateway-bin/Cargo.toml b/signal-gateway-bin/Cargo.toml index 67b96f6..ed9aac0 100644 --- a/signal-gateway-bin/Cargo.toml +++ b/signal-gateway-bin/Cargo.toml @@ -19,6 +19,7 @@ rustls-tls = ["signal-gateway/rustls-tls"] signal-gateway = { path = "../signal-gateway" } conf = { workspace = true } +conf-extra = { workspace = true } dotenvy = { workspace = true } hyper = { workspace = true } hyper-util = { workspace = true } diff --git a/signal-gateway-bin/src/admin_netcat/mod.rs b/signal-gateway-bin/src/admin_netcat/mod.rs new file mode 100644 index 0000000..808c026 --- /dev/null +++ b/signal-gateway-bin/src/admin_netcat/mod.rs @@ -0,0 +1,78 @@ +//! Admin netcat TCP client for forwarding messages to a TCP server +//! +//! This module handles admin messages not handled by the gateway by opening a TCP connection, +//! writing the message terminated with CRLF, and reading the response until CRLF. + +use conf::Conf; +use signal_gateway::MessageHandlerResult; +use std::time::Duration; +use tokio::{ + io::{AsyncBufReadExt, AsyncWriteExt, BufReader}, + net::TcpStream, + time::timeout, +}; + +/// Configuration for the admin netcat TCP client +#[derive(Clone, Conf, Debug)] +pub struct AdminNetcatConfig { + /// TCP address to forward admin commands to + #[conf(long, env)] + pub tcp_addr: String, + /// Timeout for connecting, writing, and reading + #[conf(long, env, default_value = "5s", value_parser = conf_extra::parse_duration)] + pub timeout: Duration, +} + +impl AdminNetcatConfig { + /// Create a message handler function from this config. + /// + /// The returned handler opens a TCP connection to the configured address, + /// writes the message terminated with CRLF, and reads the response until CRLF. + pub fn into_handler( + self, + ) -> Box< + dyn Fn( + String, + ) + -> std::pin::Pin + Send>> + + Send + + Sync, + > { + Box::new(move |message: String| { + let config = self.clone(); + Box::pin(async move { handle_message(&config, message).await }) + }) + } +} + +/// Handle a message by forwarding it to the configured TCP server +async fn handle_message(config: &AdminNetcatConfig, message: String) -> MessageHandlerResult { + // Connect to server + let mut stream = timeout(config.timeout, TcpStream::connect(&config.tcp_addr)) + .await + .map_err(|_| (504u16, "connecting: timeout".into()))? + .map_err(|err| (502u16, format!("connecting: {err}").into()))?; + + // Write message with CRLF terminator + let message = format!("{message}\r\n"); + timeout(config.timeout, stream.write_all(message.as_bytes())) + .await + .map_err(|_| (504u16, "writing: timeout".into()))? + .map_err(|err| (502u16, format!("writing: {err}").into()))?; + + // Read response until CR + let mut reader = BufReader::new(stream); + let mut buf = Vec::new(); + timeout(config.timeout, reader.read_until(b'\r', &mut buf)) + .await + .map_err(|_| (504u16, "reading: timeout".into()))? + .map_err(|err| (502u16, format!("reading: {err}").into()))?; + + // Convert to string and trim the trailing CR + let text = std::str::from_utf8(&buf) + .map_err(|err| (502u16, format!("utf8: {err}").into()))? + .trim_end_matches(['\r', '\n']) + .to_owned(); + + Ok((text, vec![])) +} diff --git a/signal-gateway-bin/src/main.rs b/signal-gateway-bin/src/main.rs index dfe0910..2308b3b 100644 --- a/signal-gateway-bin/src/main.rs +++ b/signal-gateway-bin/src/main.rs @@ -9,6 +9,9 @@ use tokio_util::sync::CancellationToken; use tracing::{error, info, warn}; use tracing_subscriber::EnvFilter; +mod admin_netcat; +use admin_netcat::AdminNetcatConfig; + mod syslog; use syslog::SyslogUdpConfig; @@ -22,6 +25,8 @@ struct Config { http_listen_addr: SocketAddr, #[conf(flatten, prefix)] syslog_udp: Option, + #[conf(flatten, prefix)] + admin_netcat: Option, #[conf(flatten)] gateway: GatewayConfig, } @@ -68,7 +73,8 @@ async fn main() { let token = CancellationToken::new(); - let gateway = Arc::new(Gateway::new(config.gateway, token.clone()).await); + let message_handler = config.admin_netcat.map(|c| c.into_handler()); + let gateway = Arc::new(Gateway::new(config.gateway, token.clone(), message_handler).await); let listener = TcpListener::bind(config.http_listen_addr).await.unwrap(); info!("Listening for http on {}", config.http_listen_addr); diff --git a/signal-gateway/src/gateway/mod.rs b/signal-gateway/src/gateway/mod.rs index c60f0b1..61e38a0 100644 --- a/signal-gateway/src/gateway/mod.rs +++ b/signal-gateway/src/gateway/mod.rs @@ -12,22 +12,30 @@ use http_body::Body; use http_body_util::BodyExt; use prometheus_http_client::{AlertStatus, ExtractLabels}; use std::{ - collections::HashMap, error::Error, fmt::Write, net::SocketAddr, path::PathBuf, time::Duration, + collections::HashMap, error::Error, fmt::Write, future::Future, net::SocketAddr, + path::PathBuf, pin::Pin, }; use tokio::{ - io::{AsyncBufReadExt, AsyncWriteExt, BufReader}, join, - net::TcpStream, sync::{ Mutex, RwLock, mpsc::{UnboundedReceiver, UnboundedSender, unbounded_channel}, }, - time::timeout, }; +use std::time::Duration; use tokio_util::bytes::Buf; use tokio_util::sync::CancellationToken; use tracing::{debug, error, info, warn}; +/// Response from a message handler: either success with text and optional attachments, +/// or an error with status code and message. +pub type MessageHandlerResult = Result<(String, Vec), (u16, Box)>; + +/// Handler function for admin messages that don't start with `/`. +/// Takes the message text and returns a response. +pub type MessageHandler = + Box Pin + Send>> + Send + Sync>; + mod circular_buffer; mod log_handler; use log_handler::{LogHandler, LogHandlerConfig}; @@ -41,10 +49,6 @@ pub struct GatewayConfig { pub signal_cli_tcp_addr: SocketAddr, #[conf(long, env)] pub signal_account: String, - #[conf(long, env)] - pub cbmm_tcp_addr: String, - #[conf(long, env, default_value = "5s", value_parser = conf_extra::parse_duration)] - pub cbmm_timeout: Duration, #[conf(repeat, long, env)] pub admin_uuid: Vec, #[conf(flatten)] @@ -142,8 +146,8 @@ struct AdminMessage { /// The gateway manages sending messages to signal-cli and receiving messages from signal-cli. /// It maintains a queue of messages to be sent to all admins, generated by alerts etc. -/// It also subscribes to messages received from signal and processes them one-by-one, possibly -/// making TCP request to cbmm to handle them. +/// It also subscribes to messages received from signal and processes them one-by-one, +/// calling a user-provided handler for non-command messages. /// /// This is the only task that communicates directly with signal-cli, and by design it linearizes /// all interaction, which prevents any possible races. @@ -159,10 +163,16 @@ pub struct Gateway { prometheus: Option, /// Log handlers keyed by origin (app + host). Lazily created when first message from an origin arrives. log_handlers: RwLock>, + /// Handler for admin messages that don't start with `/` + message_handler: Option, } impl Gateway { - pub async fn new(config: GatewayConfig, token: CancellationToken) -> Self { + pub async fn new( + config: GatewayConfig, + token: CancellationToken, + message_handler: Option, + ) -> Self { let (admin_mq_tx, admin_mq_rx) = unbounded_channel(); let prometheus = config @@ -179,6 +189,7 @@ impl Gateway { token, prometheus, log_handlers: RwLock::new(HashMap::new()), + message_handler, } } @@ -283,7 +294,7 @@ impl Gateway { let (message, attachments) = resp.unwrap_or_else( |(code, msg)| { let text = format!("{code}: {msg}"); - error!("(cbmm) {text}"); + error!("Message handler error: {text}"); (text, vec![]) } ); @@ -303,52 +314,25 @@ impl Gateway { } } - // Returns Err in case of a timeout + // Returns Err in case of a timeout or handler error // Returns Ok when success or error text is generated async fn handle_signal_admin_message( &self, msg: &Envelope, - ) -> Result<(String, Vec), (u16, Box)> { + ) -> Result<(String, Vec), (u16, Box)> { let data = msg.data_message.as_ref().unwrap(); // Admin messages starting with / are handled by gateway - // Other messages are forwarded to cbmm + // Other messages are passed to the configured message handler if data.message.starts_with("/") { // Parse the command using conf - let cmd = parse_gateway_command(&data.message).map_err(|err| (400, err.into()))?; + let cmd = parse_gateway_command(&data.message).map_err(|err| (400u16, err.into()))?; self.handle_gateway_command(cmd).await + } else if let Some(handler) = &self.message_handler { + handler(data.message.clone()).await } else { - // Connect to cbmm. Note that we could use a keep-alive strategy here maybe... - let mut cbmm_stream = timeout( - self.config.cbmm_timeout, - TcpStream::connect(&self.config.cbmm_tcp_addr), - ) - .await - .map_err(format_err("connecting", 504))? - .map_err(format_err("connecting", 502))?; - - timeout( - self.config.cbmm_timeout, - cbmm_stream.write_all(data.message.as_bytes()), - ) - .await - .map_err(format_err("writing", 504))? - .map_err(format_err("writing", 502))?; - - let _ = cbmm_stream.shutdown().await; - - // Wrap as BufReader so that we can use "read_until" which simplifies things - let mut reader = BufReader::new(cbmm_stream); - let mut buf = vec![]; - timeout(self.config.cbmm_timeout, reader.read_until(b'\r', &mut buf)) - .await - .map_err(format_err("reading", 504))? - .map_err(format_err("reading", 502))?; - - let s = str::from_utf8(&buf).map_err(format_err("utf8", 502))?; - let text = s.trim().to_owned(); - Ok((text, vec![])) + Err((501u16, "No message handler configured".into())) } } @@ -417,7 +401,7 @@ impl Gateway { async fn handle_gateway_command( &self, cmd: GatewayCommand, - ) -> Result<(String, Vec), (u16, Box)> { + ) -> Result<(String, Vec), (u16, Box)> { match cmd { GatewayCommand::Log { filter } => { let handlers = self.log_handlers.read().await; @@ -445,7 +429,7 @@ impl Gateway { let prometheus = self .prometheus .as_ref() - .ok_or_else(|| (501, "prometheus was not configured".into()))?; + .ok_or_else(|| (501u16, "prometheus was not configured".into()))?; match prometheus.oneoff_query(query).await { Ok(( @@ -481,7 +465,7 @@ impl Gateway { let prometheus = self .prometheus .as_ref() - .ok_or_else(|| (501, "prometheus was not configured".into()))?; + .ok_or_else(|| (501u16, "prometheus was not configured".into()))?; prometheus.purge_old_plots(); match prometheus.create_oneoff_plot(query.clone(), duration).await { @@ -497,7 +481,7 @@ impl Gateway { let prometheus = self .prometheus .as_ref() - .ok_or_else(|| (501, "prometheus was not configured".into()))?; + .ok_or_else(|| (501u16, "prometheus was not configured".into()))?; let matcher_refs: Vec<&str> = matchers.iter().map(|s| s.as_str()).collect(); match prometheus.series(&matcher_refs).await { @@ -518,7 +502,7 @@ impl Gateway { let prometheus = self .prometheus .as_ref() - .ok_or_else(|| (501, "prometheus was not configured".into()))?; + .ok_or_else(|| (501u16, "prometheus was not configured".into()))?; let matcher_refs: Vec<&str> = matchers.iter().map(|s| s.as_str()).collect(); match prometheus.labels(&matcher_refs).await { @@ -539,7 +523,7 @@ impl Gateway { let prometheus = self .prometheus .as_ref() - .ok_or_else(|| (501, "prometheus was not configured".into()))?; + .ok_or_else(|| (501u16, "prometheus was not configured".into()))?; match prometheus.alerts().await { Ok(data) => { @@ -714,14 +698,6 @@ impl Drop for Gateway { } } -// Returns a lambda that expresses an error as a (u16, String) with given context info -fn format_err( - context: &'static str, - code: u16, -) -> impl Fn(E) -> (u16, Box) { - move |err: E| -> (u16, Box) { (code, format!("{context}: {err}").into()) } -} - #[cfg(test)] mod tests { use super::*; diff --git a/signal-gateway/src/lib.rs b/signal-gateway/src/lib.rs index 0f24177..3ca6281 100644 --- a/signal-gateway/src/lib.rs +++ b/signal-gateway/src/lib.rs @@ -7,5 +7,5 @@ pub(crate) mod log_message; pub(crate) mod prometheus; pub(crate) mod transports; -pub use gateway::{Gateway, GatewayConfig}; +pub use gateway::{Gateway, GatewayConfig, MessageHandler, MessageHandlerResult}; pub use log_message::{Level, LogMessage, LogMessageBuilder}; diff --git a/signal-gateway/src/prometheus/mod.rs b/signal-gateway/src/prometheus/mod.rs index e325901..2e184ec 100644 --- a/signal-gateway/src/prometheus/mod.rs +++ b/signal-gateway/src/prometheus/mod.rs @@ -6,6 +6,8 @@ use prometheus_http_client::{ use std::error::Error; use tracing::info; +type BoxError = Box; + #[cfg(feature = "plot")] mod plot; #[cfg(feature = "plot")] @@ -52,7 +54,7 @@ impl Prometheus { pub async fn oneoff_query( &self, query: String, - ) -> Result<(ExtractLabels, Vec>), Box> { + ) -> Result<(ExtractLabels, Vec>), BoxError> { info!("Prom query: {query}"); let vector: Vec = QueryRequest { query, time: None } .send_with_client(&self.reqwest_client, &self.config.prometheus_host) @@ -69,7 +71,7 @@ impl Prometheus { pub async fn series( &self, matches: impl IntoIterator>, - ) -> Result, Box> { + ) -> Result, BoxError> { Ok(SeriesRequest { matches: matches.into_iter().map(|s| s.as_ref().to_owned()).collect(), } @@ -81,7 +83,7 @@ impl Prometheus { pub async fn labels( &self, matches: impl IntoIterator>, - ) -> Result, Box> { + ) -> Result, BoxError> { Ok(LabelsRequest { matches: matches.into_iter().map(|s| s.as_ref().to_owned()).collect(), } @@ -90,7 +92,7 @@ impl Prometheus { } /// Get the list of current alerts - pub async fn alerts(&self) -> Result, Box> { + pub async fn alerts(&self) -> Result, BoxError> { Ok(AlertsRequest {} .send_with_client(&self.reqwest_client, &self.config.prometheus_host) .await? @@ -106,7 +108,7 @@ impl Prometheus { } /// Create a new plot corresponding to a given alert. Returns a pathbuf if it is present - pub async fn create_alert_plot(&self, alert: &Alert) -> Result> { + pub async fn create_alert_plot(&self, alert: &Alert) -> Result { use chrono::{TimeDelta, Utc}; let expr = alert.parse_expr_from_generator_url()?; @@ -151,7 +153,7 @@ impl Prometheus { &self, query: String, since: Duration, - ) -> Result> { + ) -> Result { info!("Prom range query: {query}"); let matrix = QueryRangeRequest::builder(query.to_owned()) .since(since) @@ -188,7 +190,7 @@ fn build_label_selector( /// /// Returns (base_query, threshold) where base_query is the part before the comparator. #[cfg(feature = "plot")] -fn parse_alert_expr(expr: &str) -> Result<(String, PlotThreshold), Box> { +fn parse_alert_expr(expr: &str) -> Result<(String, PlotThreshold), BoxError> { // Look for comparison operators with surrounding spaces for comparator in [" < ", " > "] { if let Some(pos) = expr.find(comparator) { diff --git a/signal-gateway/src/prometheus/plot.rs b/signal-gateway/src/prometheus/plot.rs index 1da04f3..b68733d 100644 --- a/signal-gateway/src/prometheus/plot.rs +++ b/signal-gateway/src/prometheus/plot.rs @@ -1,8 +1,9 @@ +use super::BoxError; use chrono::{FixedOffset, Local, Offset, Utc}; use chrono_tz::Tz; use conf::Conf; use rand::RngCore; -use std::{error::Error, path::PathBuf, time::Duration}; +use std::{path::PathBuf, time::Duration}; use tracing::{error, warn}; use walkdir::WalkDir; @@ -46,7 +47,7 @@ impl PlotConfig { matrix: &[MetricTimeseries], threshold: Option, title: Option<&str>, - ) -> Result> { + ) -> Result { let mut filename = self.dir.clone(); filename.push(format!("plot-{}", rand::rng().next_u64())); filename.set_extension("png");