use anyhow::Result;
use axum::{Router, routing::get};
use prometheus::{Encoder, IntCounterVec, IntGaugeVec, Opts, Registry, TextEncoder};
use crate::{
app::VERSION,
config::{MetricsConfig, MetricsLabel},
};
#[derive(Debug, Clone)]
pub struct TenXProgrammer {
pub registry: Registry,
pub counters: TenXProgrammerCounters,
}
#[derive(Debug, Clone)]
pub struct TenXProgrammerCounters {
pub request_counter: IntCounterVec,
pub garbage_served_counter: IntCounterVec,
pub maze_depth: IntCounterVec,
pub challenge_counter: IntCounterVec,
}
impl TenXProgrammer {
fn build_label_names(config: &MetricsConfig) -> Vec<&'static str> {
let mut labels = Vec::new();
if config.labels.contains(&MetricsLabel::Host) {
labels.push("host");
}
if config.labels.contains(&MetricsLabel::UserAgent) {
labels.push("user_agent");
}
if config.labels.contains(&MetricsLabel::UserAgentGroup) {
labels.push("user_agent_group");
}
if config.labels.contains(&MetricsLabel::Verdict) {
labels.push("verdict");
}
labels
}
#[must_use]
pub fn build_label_values<'a>(
config: &'a MetricsConfig,
headers: &'a axum::http::HeaderMap,
verdict: &'a str,
) -> Vec<String> {
let mut values = Vec::new();
if config.labels.contains(&MetricsLabel::Host) {
if let Some(host) = headers.get("host") {
let host = String::from_utf8_lossy(host.as_bytes()).to_string();
values.push(host);
} else {
values.push("<unknown>".into());
}
}
if config.labels.contains(&MetricsLabel::UserAgent)
|| config.labels.contains(&MetricsLabel::UserAgentGroup)
{
let ua = headers.get("user-agent").map_or_else(
|| "<unknown>".into(),
|agent| String::from_utf8_lossy(agent.as_bytes()).to_string(),
);
if config.labels.contains(&MetricsLabel::UserAgent) {
values.push(ua.clone());
}
if config.labels.contains(&MetricsLabel::UserAgentGroup) {
let group = config
.agent_group
.iter()
.find(|agent_group_config| agent_group_config.agent.is_match(&ua));
let group = group.map_or("<unknown>", |cfg| &cfg.group);
values.push(group.to_owned());
}
}
if config.labels.contains(&MetricsLabel::Verdict) {
values.push(verdict.to_owned());
}
values
}
pub fn new(config: &MetricsConfig) -> Result<Option<Self>> {
if !config.enable {
return Ok(None);
}
let labels = Self::build_label_names(config);
let registry = Registry::new();
let request_counter_opts =
Opts::new("iocaine_requests_total", "Total number of requests served");
let request_counter = IntCounterVec::new(request_counter_opts, &labels)?;
registry.register(Box::new(request_counter.clone()))?;
let garbage_served_counter_opts = Opts::new(
"iocaine_garbage_served",
"Total amount of garbage served (in bytes)",
);
let garbage_served_counter = IntCounterVec::new(garbage_served_counter_opts, &labels)?;
registry.register(Box::new(garbage_served_counter.clone()))?;
let maze_depth_opts = Opts::new(
"iocaine_maze_depth",
"Maximum explored depth of the maze (in path parts)",
);
let maze_depth = IntCounterVec::new(maze_depth_opts, &labels)?;
registry.register(Box::new(maze_depth.clone()))?;
let challenge_counter_opts =
Opts::new("iocaine_challenges", "Number of challenges presented");
let challenge_counter = IntCounterVec::new(challenge_counter_opts, &labels)?;
registry.register(Box::new(challenge_counter.clone()))?;
let version_opts = Opts::new(
"iocaine_version",
"Version of the running iocaine (in the 'version' label)",
);
let version = IntGaugeVec::new(version_opts, &["version"])?;
registry.register(Box::new(version.clone()))?;
version.with_label_values(&[VERSION]).set(1);
Ok(Some(Self {
registry,
counters: TenXProgrammerCounters {
request_counter,
garbage_served_counter,
maze_depth,
challenge_counter,
},
}))
}
pub fn app(self) -> Router {
Router::new().route(
"/metrics",
get(|| async move {
let encoder = TextEncoder::new();
let mut buffer = Vec::<u8>::new();
let metrics = self.registry.gather();
encoder.encode(&metrics, &mut buffer).unwrap();
let metrics = prometheus::gather();
encoder.encode(&metrics, &mut buffer).unwrap();
String::from_utf8(buffer).unwrap()
}),
)
}
}