From a071b468b60cbf0a20e620deafca8fab6121ab8a Mon Sep 17 00:00:00 2001 From: Lol3rrr Date: Wed, 13 Mar 2024 18:58:34 +0100 Subject: [PATCH] Added Prometheus metrics and tokens for routes Added first basic prometheus metrics. Now generates a token for newly created routes that can later be used for removing specific routes or otherwise keeping track of them better --- CHANGELOG.md | 4 ++ Cargo.lock | 140 ++++++++++++++++++++++++++++++++++++++++++- Cargo.toml | 4 +- src/backend.rs | 145 ++++++++++++++++++++++++++++++++++++++++++++ src/lib.rs | 159 ++++++++++++------------------------------------- src/main.rs | 111 ++++++++++++++++++++++++++++++---- 6 files changed, 427 insertions(+), 136 deletions(-) create mode 100644 src/backend.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 315451c..349b287 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,4 +1,8 @@ # Changelog +## v0.2.0 +* Added basic prometheus metrics and corresponding endpoint +* Creating a route will now return a string token, that will be needed to remove the route normally in the future + ## v0.1.4 * Another fix for the routes with different protocols * Added more tests to detect more possible regressions diff --git a/Cargo.lock b/Cargo.lock index a04fddc..c15215d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -306,6 +306,17 @@ dependencies = [ "pin-utils", ] +[[package]] +name = "getrandom" +version = "0.2.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be4136b2a15dd319360be1c07d9933517ccf0be8f16bf62a3bee4f0d618df427" +dependencies = [ + "cfg-if", + "libc", + "wasi", +] + [[package]] name = "gimli" version = "0.27.3" @@ -383,10 +394,12 @@ dependencies = [ [[package]] name = "iptables-proxy" -version = "0.1.4" +version = "0.2.0" dependencies = [ "axum", "clap", + "prometheus", + "rand", "serde", "tokio", "tracing", @@ -428,6 +441,16 @@ version = "0.4.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "09fc20d2ca12cb9f044c93e3bd6d32d523e6e2ec3db4f7b2939cd99026ecd3f0" +[[package]] +name = "lock_api" +version = "0.4.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3c168f8615b12bc01f9c17e2eb0cc07dcae1940121185446edc3744920e8ef45" +dependencies = [ + "autocfg", + "scopeguard", +] + [[package]] name = "log" version = "0.4.19" @@ -503,6 +526,29 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b15813163c1d831bf4a13c3610c05c0d03b39feb07f7e09fa234dac9b15aaf39" +[[package]] +name = "parking_lot" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3742b2c103b9f06bc9fff0a37ff4912935851bee6d36f3c02bcc755bcfec228f" +dependencies = [ + "lock_api", + "parking_lot_core", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4c42a9226546d68acdd9c0a280d17ce19bfe27a46bf68784e4066115788d008e" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-targets", +] + [[package]] name = "percent-encoding" version = "2.3.0" @@ -541,6 +587,12 @@ version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" +[[package]] +name = "ppv-lite86" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b40af805b3121feab8a3c29f04d8ad262fa8e0561883e7653e024ae4479e6de" + [[package]] name = "proc-macro2" version = "1.0.64" @@ -550,6 +602,27 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "prometheus" +version = "0.13.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "449811d15fbdf5ceb5c1144416066429cf82316e2ec8ce0c1f6f8a02e7bbcf8c" +dependencies = [ + "cfg-if", + "fnv", + "lazy_static", + "memchr", + "parking_lot", + "protobuf", + "thiserror", +] + +[[package]] +name = "protobuf" +version = "2.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "106dd99e98437432fed6519dedecfade6a06a73bb7b2a1e019fdd2bee5778d94" + [[package]] name = "quote" version = "1.0.29" @@ -559,6 +632,45 @@ dependencies = [ "proc-macro2", ] +[[package]] +name = "rand" +version = "0.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404" +dependencies = [ + "libc", + "rand_chacha", + "rand_core", +] + +[[package]] +name = "rand_chacha" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" +dependencies = [ + "ppv-lite86", + "rand_core", +] + +[[package]] +name = "rand_core" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" +dependencies = [ + "getrandom", +] + +[[package]] +name = "redox_syscall" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4722d768eff46b75989dd134e5c353f0d6296e5aaa3132e776cbdb56be7731aa" +dependencies = [ + "bitflags 1.3.2", +] + [[package]] name = "rustc-demangle" version = "0.1.23" @@ -590,6 +702,12 @@ version = "1.0.14" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fe232bdf6be8c8de797b22184ee71118d63780ea42ac85b61d1baa6d3b782ae9" +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + [[package]] name = "serde" version = "1.0.171" @@ -700,6 +818,26 @@ version = "0.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2047c6ded9c721764247e62cd3b03c09ffc529b2ba5b10ec482ae507a4a70160" +[[package]] +name = "thiserror" +version = "1.0.55" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e3de26b0965292219b4287ff031fcba86837900fe9cd2b34ea8ad893c0953d2" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.55" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "268026685b2be38d7103e9e507c938a1fcb3d7e6eb15e87870b617bf37b6d581" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "thread_local" version = "1.1.7" diff --git a/Cargo.toml b/Cargo.toml index 90ad01b..194b247 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "iptables-proxy" -version = "0.1.4" +version = "0.2.0" edition = "2021" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html @@ -12,3 +12,5 @@ tracing = "0.1.37" tracing-subscriber = "0.3.17" serde = { version = "1.0.171", features = ["derive"] } clap = { version = "4.3.11", features = ["derive"] } +prometheus = "0.13.3" +rand = { version = "0.8.5", features = ["small_rng"] } diff --git a/src/backend.rs b/src/backend.rs new file mode 100644 index 0000000..3bda289 --- /dev/null +++ b/src/backend.rs @@ -0,0 +1,145 @@ +use std::fmt::Debug; + +use crate::ForwardingRoute; + +pub trait Command { + fn execute(self) -> impl core::future::Future>; +} + +pub trait Backend { + type Cmd: Debug + Command; + + fn register_cmds(&self, route: &ForwardingRoute) -> Result, ()>; + + fn deregister_cmds(&self, route: &ForwardingRoute) -> Result, ()>; +} + +impl Command for tokio::process::Command { + async fn execute(mut self) -> Result<(), ()> { + self.output().await.map(|o| ()).map_err(|e| ()) + } +} + +pub mod iptables { + use std::borrow::Cow; + + use crate::ForwardingRoute; + + use super::Backend; + + #[derive(Debug)] + pub struct IPTablesBackend {} + + impl IPTablesBackend { + pub fn new() -> Self { + Self {} + } + + fn register_args( + &self, + route: &ForwardingRoute, + ) -> impl Iterator>> { + [ + vec![ + "-I".into(), + "FORWARD".into(), + "-d".into(), + format!("{}", route.dest_ip).into(), + "-m".into(), + "comment".into(), + "--comment".into(), + "[iptables-proxy] SD - Accept to forward traffic".into(), + "-m".into(), + route.protocol.as_str().into(), + "-p".into(), + route.protocol.as_str().into(), + "--dport".into(), + format!("{}", route.pub_port).into(), + "-j".into(), + "ACCEPT".into(), + ], + vec![ + "-I".into(), + "FORWARD".into(), + "-m".into(), + "comment".into(), + "--comment".into(), + "[iptables-proxy] DS - Accept to forward return traffic".into(), + "-s".into(), + format!("{}", route.dest_ip).into(), + "-m".into(), + route.protocol.as_str().into(), + "-p".into(), + route.protocol.as_str().into(), + "--sport".into(), + format!("{}", route.dest_port).into(), + "-j".into(), + "ACCEPT".into(), + ], + vec![ + "-t".into(), + "nat".into(), + "-I".into(), + "PREROUTING".into(), + "-m".into(), + route.protocol.as_str().into(), + "-p".into(), + route.protocol.as_str().into(), + "--dport".into(), + format!("{}", route.pub_port).into(), + "-m".into(), + "comment".into(), + "--comment".into(), + "[iptables-proxy] redirect pkts to homeserver".into(), + "-j".into(), + "DNAT".into(), + "--to-destination".into(), + format!("{}:{}", route.dest_ip, route.dest_port).into(), + ], + ] + .into_iter() + } + fn deregister_args( + &self, + route: &ForwardingRoute, + ) -> impl Iterator>> { + self.register_args(route).map(|mut args| { + for arg in args.iter_mut() { + if arg == "-I" { + *arg = "-D".into(); + } + } + args + }) + } + } + + impl Backend for IPTablesBackend { + type Cmd = tokio::process::Command; + + fn register_cmds(&self, route: &crate::ForwardingRoute) -> Result, ()> { + let cmds = self + .register_args(route) + .map(|args| { + let mut cmd = tokio::process::Command::new("iptables"); + cmd.args(args.into_iter().map(|c| c.to_string())); + cmd + }) + .collect(); + + Ok(cmds) + } + + fn deregister_cmds(&self, route: &crate::ForwardingRoute) -> Result, ()> { + let cmds = self + .deregister_args(route) + .map(|args| { + let mut cmd = tokio::process::Command::new("iptables"); + cmd.args(args.into_iter().map(|c| c.to_string())); + cmd + }) + .collect(); + Ok(cmds) + } + } +} diff --git a/src/lib.rs b/src/lib.rs index 18fe6ca..8a51f20 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,33 +1,55 @@ +use std::collections::HashMap; + +use rand::{Rng, SeedableRng}; use serde::Deserialize; -use std::borrow::Cow; mod rule_parser; +pub mod backend; + pub struct Routes { - routes: Vec, + routes: HashMap, + rng: rand::rngs::SmallRng, } impl Routes { pub fn new() -> Self { - Self { routes: Vec::new() } + Self { + routes: HashMap::new(), + rng: rand::rngs::SmallRng::from_entropy(), + } } - pub fn add(&mut self, route: ForwardingRoute) -> Option { - if let Some((idx, _)) = self.routes.iter().enumerate().find(|(_, r)| { + pub fn add(&mut self, route: ForwardingRoute) -> (String, Option) { + let token = loop { + let value: String = (&mut self.rng) + .sample_iter(&rand::distributions::Alphanumeric) + .take(10) + .map(char::from) + .collect(); + + if !self.routes.contains_key(&value) { + break value; + } + }; + + let previous_route = if let Some((idx, _)) = self.routes.iter().find(|(_, r)| { r.public_ip() == route.public_ip() && r.public_port() == route.public_port() && r.protocol() == route.protocol() }) { - let previous = self.routes.swap_remove(idx); + let previous = self.routes.remove(&idx.clone()); - self.routes.push(route); + self.routes.insert(token.clone(), route); - Some(previous) + previous } else { - self.routes.push(route); + self.routes.insert(token.clone(), route); None - } + }; + + (token, previous_route) } pub fn remove( @@ -36,10 +58,10 @@ impl Routes { public_port: u16, protocol: &Protocol, ) -> Option { - match self.routes.iter().enumerate().find(|(_, r)| { + match self.routes.iter().find(|(_, r)| { r.public_ip() == public_ip && r.public_port() == public_port && r.protocol() == protocol }) { - Some((i, _)) => Some(self.routes.swap_remove(i)), + Some((i, _)) => self.routes.remove(&i.clone()), None => None, } } @@ -93,113 +115,6 @@ impl ForwardingRoute { pub fn protocol(&self) -> &Protocol { &self.protocol } - - fn register_args(&self) -> impl Iterator>> { - [ - vec![ - "-I".into(), - "FORWARD".into(), - "-d".into(), - format!("{}", self.dest_ip).into(), - "-m".into(), - "comment".into(), - "--comment".into(), - "[iptables-proxy] SD - Accept to forward traffic".into(), - "-m".into(), - self.protocol.as_str().into(), - "-p".into(), - self.protocol.as_str().into(), - "--dport".into(), - format!("{}", self.pub_port).into(), - "-j".into(), - "ACCEPT".into(), - ], - vec![ - "-I".into(), - "FORWARD".into(), - "-m".into(), - "comment".into(), - "--comment".into(), - "[iptables-proxy] DS - Accept to forward return traffic".into(), - "-s".into(), - format!("{}", self.dest_ip).into(), - "-m".into(), - self.protocol.as_str().into(), - "-p".into(), - self.protocol.as_str().into(), - "--sport".into(), - format!("{}", self.dest_port).into(), - "-j".into(), - "ACCEPT".into(), - ], - vec![ - "-t".into(), - "nat".into(), - "-I".into(), - "PREROUTING".into(), - "-m".into(), - self.protocol.as_str().into(), - "-p".into(), - self.protocol.as_str().into(), - "--dport".into(), - format!("{}", self.pub_port).into(), - "-m".into(), - "comment".into(), - "--comment".into(), - "[iptables-proxy] redirect pkts to homeserver".into(), - "-j".into(), - "DNAT".into(), - "--to-destination".into(), - format!("{}:{}", self.dest_ip, self.dest_port).into(), - ], - ] - .into_iter() - } - fn deregister_args(&self) -> impl Iterator>> { - self.register_args().map(|mut args| { - for arg in args.iter_mut() { - if arg == "-I" { - *arg = "-D".into(); - } - } - args - }) - } - - pub fn dry_register(&self) -> impl Iterator { - self.register_args() - .map(|args| ("iptables".to_string(), args.join(" "))) - } - - pub fn register(&self) -> impl Iterator { - self.register_args().map(|args| { - let mut cmd = tokio::process::Command::new("iptables"); - cmd.args(args.into_iter().map(|c| c.to_string())); - cmd - }) - } - - pub fn dry_deregister(&self) -> impl Iterator { - self.deregister_args() - .map(|args| ("iptables".to_string(), args.join(" "))) - } - pub fn deregister(&self) -> impl Iterator { - self.deregister_args().map(|args| { - let mut cmd = tokio::process::Command::new("iptables"); - cmd.args(args.into_iter().map(|c| c.to_string())); - cmd - }) - } -} - -pub struct Config { - static_routes: Vec, -} - -impl Config { - pub fn load() -> Self { - todo!("Work on loading the default standard Config") - } } #[cfg(test)] @@ -217,7 +132,7 @@ mod tests { dest_port: 1234, protocol: Protocol::Tcp, }); - assert_eq!(None, prev); + assert_eq!(None, prev.1); } #[test] @@ -231,7 +146,7 @@ mod tests { dest_port: 1234, protocol: Protocol::Tcp, }); - assert_eq!(None, prev); + assert_eq!(None, prev.1); let prev = routes.add(ForwardingRoute { pub_ip: "".into(), @@ -240,6 +155,6 @@ mod tests { dest_port: 1234, protocol: Protocol::Udp, }); - assert_eq!(None, prev); + assert_eq!(None, prev.1); } } diff --git a/src/main.rs b/src/main.rs index fc9c6e3..19d7f65 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,13 +1,22 @@ use clap::Parser; -use iptables_proxy::{ForwardingRoute, Protocol}; +use iptables_proxy::{ + backend::{Backend, Command}, + ForwardingRoute, Protocol, +}; use tracing_subscriber::{layer::Filter, prelude::__tracing_subscriber_SubscriberExt, Layer}; -use axum::{extract::State, http::StatusCode, routing::post, Json, Router}; +use axum::{ + extract::State, + http::StatusCode, + routing::{get, post}, + Json, Router, +}; use std::{ net::{IpAddr, SocketAddr}, sync::{Arc, Mutex}, }; +use prometheus::Encoder; use serde::Deserialize; #[derive(Debug, Parser)] @@ -39,10 +48,12 @@ struct AppState { routes: Mutex, dry_run: bool, public_ip: String, + metrics: prometheus::Registry, + requests: prometheus::CounterVec, } impl AppState { - pub fn add(&self, route: ForwardingRoute) -> Option { + pub fn add(&self, route: ForwardingRoute) -> (String, Option) { let mut existing_routes = self.routes.lock().unwrap(); existing_routes.add(route) @@ -72,13 +83,27 @@ async fn main() { tracing::info!("Running in dry run mode"); } + let registry = prometheus::Registry::new_custom(Some("iptables_proxy".into()), None).unwrap(); + + let request_counter = prometheus::CounterVec::new( + prometheus::Opts::new("requests", "The Number of requests received by the proxy"), + &["endpoint"], + ) + .unwrap(); + registry + .register(Box::new(request_counter.clone())) + .unwrap(); + let app_state = Arc::new(AppState { routes: Mutex::new(iptables_proxy::Routes::new()), dry_run: args.dry_run, public_ip: format!("{}", args.public_ip), + metrics: registry, + requests: request_counter, }); let app = Router::new() + .route("/metrics", get(metrics)) .route("/create", post(create)) .route("/remove", post(remove)) .with_state(app_state); @@ -103,41 +128,78 @@ struct CreateRequest { async fn create( State(state): State>, Json(payload): Json, -) -> StatusCode { +) -> axum::response::Response { tracing::debug!("Received request to create route: {:?}", payload); + state + .requests + .get_metric_with_label_values(&["create"]) + .unwrap() + .inc(); + let route = ForwardingRoute::new( (state.public_ip.clone(), payload.public_port), (payload.inner_ip.clone(), payload.inner_port), payload.protocol.clone(), ); - if let Some(old_route) = state.add(route.clone()) { + let backend = iptables_proxy::backend::iptables::IPTablesBackend::new(); + + let (token, prev_route) = state.add(route.clone()); + if let Some(old_route) = prev_route { tracing::error!("Route already exists, replacing existing route"); // Deregister the old route from iptables - for mut cmd in old_route.deregister() { + let deregister_cmds = match backend.deregister_cmds(&old_route) { + Ok(r) => r, + Err(e) => { + tracing::error!("Getting Deregister Commands: {:?}", e); + return axum::response::Response::builder() + .status(StatusCode::INTERNAL_SERVER_ERROR) + .body("".into()) + .unwrap(); + } + }; + for cmd in deregister_cmds { tracing::debug!("Running Command: {:?}", cmd); if !state.dry_run { - if let Err(e) = cmd.status().await { + if let Err(e) = cmd.execute().await { tracing::error!("Executing Command: {:?}", e); + return axum::response::Response::builder() + .status(StatusCode::INTERNAL_SERVER_ERROR) + .body("Removing previous entry".into()) + .unwrap(); } } } } - for mut cmd in route.register() { + let register_cmds = match backend.register_cmds(&route) { + Ok(c) => c, + Err(e) => { + tracing::error!("Getting Register Commands: {:?}", e); + return axum::response::Response::builder() + .status(StatusCode::INTERNAL_SERVER_ERROR) + .body("".into()) + .unwrap(); + } + }; + for cmd in register_cmds { tracing::debug!("Running Command: {:?}", cmd); if !state.dry_run { - if let Err(e) = cmd.status().await { + if let Err(e) = cmd.execute().await { tracing::error!("Executing Command: {:?}", e); + return axum::response::Response::builder() + .status(StatusCode::INTERNAL_SERVER_ERROR) + .body("Creating new entry".into()) + .unwrap(); } } } tracing::debug!("Created route"); - StatusCode::OK + axum::response::Response::builder().body(token).unwrap() } #[derive(Debug, Deserialize)] @@ -153,6 +215,12 @@ async fn remove( ) -> StatusCode { tracing::debug!("Received request to remove route: {:?}", payload); + state + .requests + .get_metric_with_label_values(&["remove"]) + .unwrap() + .inc(); + let route = state.remove(&payload); let route = match route { @@ -163,10 +231,19 @@ async fn remove( } }; - for mut cmd in route.deregister() { + let backend = iptables_proxy::backend::iptables::IPTablesBackend::new(); + + let deregister_cmds = match backend.deregister_cmds(&route) { + Ok(r) => r, + Err(e) => { + tracing::error!("Getting Deregister Commands: {:?}", e); + return StatusCode::INTERNAL_SERVER_ERROR; + } + }; + for cmd in deregister_cmds { tracing::debug!("Running Command: {:?}", cmd); if !state.dry_run { - if let Err(e) = cmd.status().await { + if let Err(e) = cmd.execute().await { tracing::error!("Executing Command: {:?}", e); } } @@ -176,3 +253,13 @@ async fn remove( StatusCode::OK } + +#[tracing::instrument(skip(state))] +async fn metrics(State(state): State>) -> impl axum::response::IntoResponse { + let mut buffer = Vec::new(); + let encoder = prometheus::TextEncoder::new(); + let metric_families = state.metrics.gather(); + encoder.encode(&metric_families, &mut buffer).unwrap(); + + String::from_utf8(buffer).unwrap() +}