use std::net::IpAddr;
use tracing::{span, Level, Span};
use uuid::Uuid;
#[derive(Debug, Clone)]
pub struct SessionContext {
pub session_id: String,
pub client_ip: IpAddr,
pub protocol: String,
pub username: Option<String>,
}
impl SessionContext {
pub fn new(client_ip: IpAddr, protocol: impl Into<String>) -> Self {
Self {
session_id: Uuid::new_v4().to_string(),
client_ip,
protocol: protocol.into(),
username: None,
}
}
pub fn with_id(
session_id: impl Into<String>,
client_ip: IpAddr,
protocol: impl Into<String>,
) -> Self {
Self {
session_id: session_id.into(),
client_ip,
protocol: protocol.into(),
username: None,
}
}
pub fn set_username(&mut self, username: impl Into<String>) {
self.username = Some(username.into());
}
pub fn session_id(&self) -> &str {
&self.session_id
}
pub fn client_ip(&self) -> IpAddr {
self.client_ip
}
pub fn protocol(&self) -> &str {
&self.protocol
}
pub fn username(&self) -> Option<&str> {
self.username.as_deref()
}
}
#[derive(Debug)]
pub struct SessionLogger {
context: SessionContext,
span: Span,
}
impl SessionLogger {
pub fn new(context: SessionContext) -> Self {
let span = if let Some(ref username) = context.username {
span!(
Level::INFO,
"session",
session_id = %context.session_id,
client_ip = %context.client_ip,
protocol = %context.protocol,
username = %username,
)
} else {
span!(
Level::INFO,
"session",
session_id = %context.session_id,
client_ip = %context.client_ip,
protocol = %context.protocol,
)
};
Self { context, span }
}
pub fn enter(&self) -> tracing::span::Entered<'_> {
self.span.enter()
}
pub fn context(&self) -> &SessionContext {
&self.context
}
pub fn set_username(&mut self, username: impl Into<String>) {
self.context.set_username(username);
let username_str = self.context.username.as_deref().unwrap_or("<unknown>");
self.span = span!(
Level::INFO,
"session",
session_id = %self.context.session_id,
client_ip = %self.context.client_ip,
protocol = %self.context.protocol,
username = %username_str,
);
}
pub fn span(&self) -> &Span {
&self.span
}
}
#[macro_export]
macro_rules! session_info {
($logger:expr, $($arg:tt)*) => {
{
let _guard = $logger.enter();
tracing::info!($($arg)*);
}
};
}
#[macro_export]
macro_rules! session_debug {
($logger:expr, $($arg:tt)*) => {
{
let _guard = $logger.enter();
tracing::debug!($($arg)*);
}
};
}
#[macro_export]
macro_rules! session_warn {
($logger:expr, $($arg:tt)*) => {
{
let _guard = $logger.enter();
tracing::warn!($($arg)*);
}
};
}
#[macro_export]
macro_rules! session_error {
($logger:expr, $($arg:tt)*) => {
{
let _guard = $logger.enter();
tracing::error!($($arg)*);
}
};
}
#[macro_export]
macro_rules! session_trace {
($logger:expr, $($arg:tt)*) => {
{
let _guard = $logger.enter();
tracing::trace!($($arg)*);
}
};
}
pub fn format_session_header(context: &SessionContext) -> String {
context.session_id.clone()
}
pub fn format_session_json(context: &SessionContext) -> String {
let username_str = context.username.as_deref().unwrap_or("");
format!(
r#"{{"session_id":"{}","client_ip":"{}","protocol":"{}","username":"{}"}}"#,
context.session_id, context.client_ip, context.protocol, username_str
)
}
#[cfg(test)]
mod tests {
use super::*;
use std::net::IpAddr;
#[test]
fn test_session_context_creation() {
let ip = IpAddr::from([192, 168, 1, 100]);
let session = SessionContext::new(ip, "SMTP");
assert!(!session.session_id.is_empty());
assert_eq!(session.client_ip, ip);
assert_eq!(session.protocol, "SMTP");
assert!(session.username.is_none());
}
#[test]
fn test_session_context_with_id() {
let ip = IpAddr::from([10, 0, 0, 1]);
let custom_id = "test-session-123";
let session = SessionContext::with_id(custom_id, ip, "IMAP");
assert_eq!(session.session_id, custom_id);
assert_eq!(session.client_ip, ip);
assert_eq!(session.protocol, "IMAP");
}
#[test]
fn test_set_username() {
let ip = IpAddr::from([127, 0, 0, 1]);
let mut session = SessionContext::new(ip, "POP3");
assert!(session.username.is_none());
session.set_username("alice");
assert_eq!(session.username.as_deref(), Some("alice"));
}
#[test]
fn test_session_logger_creation() {
let ip = IpAddr::from([192, 168, 1, 1]);
let context = SessionContext::new(ip, "JMAP");
let logger = SessionLogger::new(context);
assert_eq!(logger.context().protocol, "JMAP");
assert_eq!(logger.context().client_ip, ip);
}
#[test]
fn test_session_logger_set_username() {
let ip = IpAddr::from([172, 16, 0, 1]);
let context = SessionContext::new(ip, "SMTP");
let mut logger = SessionLogger::new(context);
assert!(logger.context().username.is_none());
logger.set_username("bob");
assert_eq!(logger.context().username.as_deref(), Some("bob"));
}
#[test]
fn test_format_session_header() {
let ip = IpAddr::from([127, 0, 0, 1]);
let session = SessionContext::with_id("abc-123", ip, "JMAP");
let header = format_session_header(&session);
assert_eq!(header, "abc-123");
}
#[test]
fn test_format_session_json() {
let ip = IpAddr::from([192, 168, 1, 50]);
let mut session = SessionContext::with_id("test-id", ip, "IMAP");
session.set_username("alice");
let json = format_session_json(&session);
assert!(json.contains(r#""session_id":"test-id""#));
assert!(json.contains(r#""client_ip":"192.168.1.50""#));
assert!(json.contains(r#""protocol":"IMAP""#));
assert!(json.contains(r#""username":"alice""#));
}
#[test]
fn test_format_session_json_no_username() {
let ip = IpAddr::from([10, 0, 0, 1]);
let session = SessionContext::with_id("session-123", ip, "SMTP");
let json = format_session_json(&session);
assert!(json.contains(r#""username":"""#));
}
#[test]
fn test_session_id_is_uuid() {
let ip = IpAddr::from([127, 0, 0, 1]);
let session = SessionContext::new(ip, "POP3");
let parsed = Uuid::parse_str(&session.session_id);
assert!(parsed.is_ok(), "Session ID should be a valid UUID");
}
#[test]
fn test_unique_session_ids() {
let ip = IpAddr::from([127, 0, 0, 1]);
let session1 = SessionContext::new(ip, "SMTP");
let session2 = SessionContext::new(ip, "SMTP");
assert_ne!(
session1.session_id, session2.session_id,
"Session IDs should be unique"
);
}
}