markbase/markbase-core/src/ssh_server/auth.rs

// SSH认证协议实现（Phase 5）
// 参考OpenSSH auth.c, auth-passwd.c

use crate::ssh_server::packet::{SshPacket, PacketType};
use std::io::{Read, Write};  // 导入Write trait（OpenSSH标准）
use anyhow::{Result, anyhow};
use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};
use log::{info, warn, debug};
use rusqlite::{Connection, params};
use bcrypt::{verify, DEFAULT_COST};

/// SSH认证处理器（参考OpenSSH auth2.c）
pub struct AuthHandler {
    db_path: String,  // SQLite数据库路径
}

impl AuthHandler {
    /// 创建认证处理器
    pub fn new() -> Result<Self> {
        let db_path = "data/auth.sqlite".to_string();
        
        // 验证数据库是否存在
        let conn = Connection::open(&db_path)?;
        drop(conn);  // rusqlite会自动关闭
        
        info!("AuthHandler initialized with database: {}", db_path);
        Ok(Self { db_path })
    }
    
    /// 处理SSH_MSG_USERAUTH_REQUEST（参考OpenSSH auth2.c: userauth_request()）
    pub fn handle_userauth_request(&mut self, packet: &SshPacket) -> Result<AuthResult> {
        info!("Processing SSH_MSG_USERAUTH_REQUEST");
        
        let mut cursor = std::io::Cursor::new(packet.payload.as_slice());
        
        // Packet type
        let packet_type = cursor.read_u8()?;
        if packet_type != PacketType::SSH_MSG_USERAUTH_REQUEST as u8 {
            return Err(anyhow!("Invalid packet type for USERAUTH_REQUEST"));
        }
        
        // 读取用户名（SSH string）
        let user = read_ssh_string(&mut cursor)?;
        
        // 读取服务名称（SSH string）
        let service = read_ssh_string(&mut cursor)?;
        
        // 读取认证方法名称（SSH string）
        let method = read_ssh_string(&mut cursor)?;
        
        info!("Auth request: user={}, service={}, method={}", user, service, method);
        
        // 检查服务名称（OpenSSH要求：ssh-connection）
        if service != "ssh-connection" {
            warn!("Unsupported service: {}", service);
            return Ok(AuthResult::Failure("Unsupported service".to_string()));
        }
        
        // 根据认证方法处理（参考OpenSSH auth2.c）
        if method == "password" {
            self.handle_password_auth(&mut cursor, &user)
        } else if method == "publickey" {
            // Phase 5仅实现password认证，publickey留待Phase 9优化
            warn!("Public key auth not implemented in Phase 5");
            // SSH_MSG_USERAUTH_FAILURE必须返回可继续使用的认证方法列表
            Ok(AuthResult::Failure("password".to_string()))
        } else if method == "none" {
            // OpenSSH：none认证总是失败（用于查询支持的认证方法）
            // 返回支持的认证方法列表：password, publickey
            warn!("None auth request - returning supported methods");
            Ok(AuthResult::Failure("password,publickey".to_string()))
        } else {
            warn!("Unsupported auth method: {}", method);
            Ok(AuthResult::Failure("Unsupported auth method".to_string()))
        }
    }
    
    /// 处理password认证（参考OpenSSH auth-passwd.c）
    fn handle_password_auth(&mut self, cursor: &mut std::io::Cursor<&[u8]>, user: &str) -> Result<AuthResult> {
        info!("Handling password auth for user: {}", user);
        
        // 读取是否修改密码标志（boolean，OpenSSH password认证格式）
        let change_password = cursor.read_u8()? != 0;
        
        if change_password {
            warn!("Password change not supported");
            return Ok(AuthResult::Failure("Password change not supported".to_string()));
        }
        
        // 读取密码（SSH string）
        let password = read_ssh_string(cursor)?;
        
        debug!("Password auth attempt: user={}, password length={}", user, password.len());
        
        // 查询数据库获取password_hash
        let conn = Connection::open(&self.db_path)?;
        
        let password_hash_result = conn.query_row(
            "SELECT password_hash FROM sftpgo_users WHERE username = ?1 AND status = 1",
            params![user],
            |row| row.get::<_, String>(0)
        );
        
        // 关闭连接（rusqlite会自动关闭）
        drop(conn);
        
        // 验证用户是否存在
        let password_hash = match password_hash_result {
            Ok(hash) => Some(hash),
            Err(rusqlite::Error::QueryReturnedNoRows) => None,
            Err(e) => return Err(anyhow!("Database query error: {}", e)),
        };
        
        if password_hash.is_none() {
            warn!("User not found or disabled: {}", user);
            // SSH_MSG_USERAUTH_FAILURE必须返回可继续使用的认证方法列表（RFC 4253）
            return Ok(AuthResult::Failure("password,publickey".to_string()));
        }
        
        // 使用bcrypt验证密码
        let stored_hash = password_hash.unwrap();
        let valid = verify(&password, &stored_hash)?;
        
        if valid {
            info!("Password auth successful for user: {}", user);
            Ok(AuthResult::Success)
        } else {
            warn!("Password auth failed for user: {}", user);
            // SSH_MSG_USERAUTH_FAILURE必须返回可继续使用的认证方法列表（RFC 4253）
            Ok(AuthResult::Failure("password,publickey".to_string()))
        }
    }
    
    /// 构建SSH_MSG_USERAUTH_SUCCESS packet（参考OpenSSH auth2.c）
    pub fn build_userauth_success() -> Result<SshPacket> {
        let payload = vec![PacketType::SSH_MSG_USERAUTH_SUCCESS as u8];
        Ok(SshPacket::new(payload))
    }
    
    /// 构建SSH_MSG_USERAUTH_FAILURE packet（参考OpenSSH auth2.c）
    pub fn build_userauth_failure(methods: &[String], partial_success: bool) -> Result<SshPacket> {
        let mut payload = Vec::new();
        
        // Packet type
        payload.write_u8(PacketType::SSH_MSG_USERAUTH_FAILURE as u8)?;
        
        // 认证方法列表（SSH string，逗号分隔）
        let methods_str = methods.join(",");
        payload.write_u32::<BigEndian>(methods_str.len() as u32)?;
        payload.write_all(methods_str.as_bytes())?;
        
        // partial_success标志（boolean）
        payload.write_u8(if partial_success { 1 } else { 0 })?;
        
        Ok(SshPacket::new(payload))
    }
    
    /// 构建SSH_MSG_USERAUTH_BANNER packet（可选，参考OpenSSH auth2.c）
    pub fn build_userauth_banner(message: &str, language: &str) -> Result<SshPacket> {
        let mut payload = Vec::new();
        
        // Packet type
        payload.write_u8(PacketType::SSH_MSG_USERAUTH_BANNER as u8)?;
        
        // Banner message（SSH string）
        payload.write_u32::<BigEndian>(message.len() as u32)?;
        payload.write_all(message.as_bytes())?;
        
        // Language tag（SSH string）
        payload.write_u32::<BigEndian>(language.len() as u32)?;
        payload.write_all(language.as_bytes())?;
        
        Ok(SshPacket::new(payload))
    }
}

/// SSH认证结果（参考OpenSSH auth2.c）
pub enum AuthResult {
    Success,
    Failure(String),  // 失败原因
    PartialSuccess,   // 部分成功（多步骤认证）
}

/// SSH string读取辅助函数
fn read_ssh_string<R: std::io::Read>(reader: &mut R) -> Result<String> {
    let length = reader.read_u32::<BigEndian>()?;
    let mut buffer = vec![0u8; length as usize];
    reader.read_exact(&mut buffer)?;
    Ok(String::from_utf8(buffer)?)
}

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_userauth_success_packet() {
        let packet = AuthHandler::build_userauth_success().unwrap();
        assert_eq!(packet.payload[0], PacketType::SSH_MSG_USERAUTH_SUCCESS as u8);
    }
    
    #[test]
    fn test_userauth_failure_packet() {
        let methods = vec!["password".to_string(), "publickey".to_string()];
        let packet = AuthHandler::build_userauth_failure(&methods, false).unwrap();
        
        assert_eq!(packet.payload[0], PacketType::SSH_MSG_USERAUTH_FAILURE as u8);
    }
}