pollux/src/logging.rs

92 lines
No EOL
2.7 KiB
Rust

use tokio::net::TcpStream;
use tokio_rustls::server::TlsStream;
use tracing_subscriber::fmt::time::FormatTime;
struct GeminiTimeFormat;
impl FormatTime for GeminiTimeFormat {
fn format_time(&self, w: &mut tracing_subscriber::fmt::format::Writer<'_>) -> std::fmt::Result {
let now = time::OffsetDateTime::now_utc();
write!(w, "{}-{:02}-{:02}T{:02}:{:02}:{:02}",
now.year(), now.month() as u8, now.day(),
now.hour(), now.minute(), now.second())
}
}
pub struct RequestLogger {
client_ip: String,
request_url: String,
}
impl RequestLogger {
pub fn new(stream: &TlsStream<TcpStream>, request_url: String) -> Self {
let client_ip = extract_client_ip(stream);
Self {
client_ip,
request_url,
}
}
pub fn log_error(self, status_code: u8, error_message: &str) {
let level = match status_code {
41 | 51 => tracing::Level::WARN,
59 => tracing::Level::ERROR,
_ => tracing::Level::ERROR,
};
let request_path = self.request_url.strip_prefix("gemini://localhost").unwrap_or(&self.request_url);
match level {
tracing::Level::WARN => tracing::warn!("{} \"{}\" {} \"{}\"", self.client_ip, request_path, status_code, error_message),
tracing::Level::ERROR => tracing::error!("{} \"{}\" {} \"{}\"", self.client_ip, request_path, status_code, error_message),
_ => {}
}
}
}
fn extract_client_ip(stream: &TlsStream<TcpStream>) -> String {
let (tcp_stream, _) = stream.get_ref();
match tcp_stream.peer_addr() {
Ok(addr) => addr.to_string(),
Err(_) => "unknown".to_string(),
}
}
pub fn init_logging(level: &str) {
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
let level = match level.to_lowercase().as_str() {
"error" => tracing::Level::ERROR,
"warn" => tracing::Level::WARN,
"info" => tracing::Level::INFO,
"debug" => tracing::Level::DEBUG,
"trace" => tracing::Level::TRACE,
_ => {
eprintln!("Warning: Invalid log level '{}', defaulting to 'info'", level);
tracing::Level::INFO
}
};
tracing_subscriber::registry()
.with(tracing_subscriber::fmt::layer()
.compact()
.with_timer(GeminiTimeFormat)
.with_target(false)
.with_thread_ids(false))
.with(tracing_subscriber::filter::LevelFilter::from_level(level))
.init();
}
#[cfg(test)]
mod tests {
#[test]
fn test_basic_functionality() {
// Basic test to ensure logging module compiles
assert!(true);
}
}