rust-work/ssh-proxy/src/logger.rs

use serde::{Deserialize, Serialize};
use std::fs::{File, OpenOptions};
use std::io::Write;
use std::path::PathBuf;
use std::sync::Mutex;

/// 日志记录器
pub struct RequestLogger {
    log_file: Mutex<Option<File>>,
    enabled: bool,
}

impl RequestLogger {
    pub fn new(log_path: Option<&str>) -> Self {
        let (log_file, enabled) = if let Some(path) = log_path {
            let path = expand_path(path);
            if let Some(parent) = path.parent() {
                let _ = std::fs::create_dir_all(parent);
            }
            match OpenOptions::new()
                .create(true)
                .append(true)
                .open(&path)
            {
                Ok(file) => (Some(file), true),
                Err(e) => {
                    eprintln!("[Logger] Failed to open log file {}: {}", path.display(), e);
                    (None, false)
                }
            }
        } else {
            (None, false)
        };

        Self {
            log_file: Mutex::new(log_file),
            enabled,
        }
    }

    pub fn is_enabled(&self) -> bool {
        self.enabled
    }

    /// 记录请求日志
    pub fn log(&self, entry: &LogEntry) {
        if !self.enabled {
            return;
        }

        let json = match serde_json::to_string(entry) {
            Ok(j) => j,
            Err(e) => {
                eprintln!("[Logger] Failed to serialize log: {}", e);
                return;
            }
        };

        if let Ok(mut file) = self.log_file.lock() {
            if let Some(ref mut f) = *file {
                let _ = writeln!(f, "{}", json);
            }
        }
    }
}

/// 日志条目
#[derive(Debug, Serialize, Deserialize)]
pub struct LogEntry {
    pub timestamp: String,
    pub level: String,
    #[serde(rename = "type")]
    pub log_type: String,
    pub client: String,
    pub endpoint: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub conn: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub server: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub sql: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub command: Option<String>,
    #[serde(rename = "durationMs")]
    pub duration_ms: u64,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub rows: Option<usize>,
    #[serde(skip_serializing_if = "Option::is_none")]
    #[serde(rename = "exitCode")]
    pub exit_code: Option<i32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub error: Option<String>,
}

impl LogEntry {
    pub fn new(endpoint: &str, client: &str) -> Self {
        Self {
            timestamp: current_timestamp(),
            level: "INFO".to_string(),
            log_type: "request".to_string(),
            client: client.to_string(),
            endpoint: endpoint.to_string(),
            conn: None,
            server: None,
            sql: None,
            command: None,
            duration_ms: 0,
            rows: None,
            exit_code: None,
            error: None,
        }
    }

    pub fn with_conn(mut self, conn: &str) -> Self {
        self.conn = Some(conn.to_string());
        self
    }

    pub fn with_server(mut self, server: &str) -> Self {
        self.server = Some(server.to_string());
        self
    }

    pub fn with_sql(mut self, sql: &str) -> Self {
        // 截断过长的 SQL
        self.sql = Some(truncate_string(sql, 1000));
        self
    }

    pub fn with_command(mut self, command: &str) -> Self {
        // 截断过长的命令
        self.command = Some(truncate_string(command, 500));
        self
    }

    pub fn with_duration(mut self, ms: u64) -> Self {
        self.duration_ms = ms;
        // 慢请求告警: >3s 为 WARN, >10s 为 ERROR
        // 注意: 后台命令(nohup)可能耗时很长但 exitCode=0，不标记错误
        if ms > 10000 && self.level == "INFO" {
            self.level = "ERROR".to_string();
        } else if ms > 3000 && self.level == "INFO" {
            self.level = "WARN".to_string();
        }
        self
    }

    pub fn with_rows(mut self, rows: usize) -> Self {
        self.rows = Some(rows);
        self
    }

    pub fn with_exit_code(mut self, code: i32) -> Self {
        self.exit_code = Some(code);
        // 只标记真正的错误 (>=2 或 127)
        // exitCode: -1 (pkill未找到), 0 (成功), 1 (grep未匹配) 视为正常
        if code >= 2 || code == 127 {
            if self.level == "INFO" {
                self.level = "ERROR".to_string();
            }
        }
        self
    }

    pub fn with_error(mut self, error: &str) -> Self {
        self.level = "ERROR".to_string();
        self.error = Some(error.to_string());
        self
    }
}

fn current_timestamp() -> String {
    time::OffsetDateTime::now_utc()
        .format(&time::macros::format_description!("[year]-[month]-[day]T[hour]:[minute]:[second]"))
        .unwrap_or_default()
}

fn truncate_string(s: &str, max_len: usize) -> String {
    if s.chars().count() <= max_len {
        s.to_string()
    } else {
        format!("{}... (truncated)", s.chars().take(max_len).collect::<String>())
    }
}

fn expand_path(path: &str) -> PathBuf {
    if path.starts_with('~') {
        let home = std::env::var("HOME")
            .or_else(|_| std::env::var("USERPROFILE"))
            .unwrap_or_default();
        PathBuf::from(path.replacen('~', &home, 1))
    } else {
        PathBuf::from(path)
    }
}