diff --git a/markbase-core/src/sftp/server.rs b/markbase-core/src/sftp/server.rs index c57ac49..99cc446 100644 --- a/markbase-core/src/sftp/server.rs +++ b/markbase-core/src/sftp/server.rs @@ -1,19 +1,28 @@ +use crate::provider::DataProvider; use crate::sftp::audit::AuditLog; use crate::sftp::config::SftpConfig; -use crate::sftp::pty::PtySession; +use crate::sftp::handler::SftpHandler; use crate::sftp::shell::ShellHandler; use russh::server::{Auth, Msg, Server, Session}; -use russh::{keys, Channel, ChannelId, MethodSet}; +use russh::{Channel, ChannelId}; use russh_keys::PrivateKey; +use std::collections::HashMap; use std::path::Path; -use std::sync::{Arc, Mutex}; +use std::sync::Arc; use std::time::Duration; use tokio::sync::Mutex; +use anyhow::Result; pub struct MarkBaseSftpServer { - user_id: String, config: Arc, + provider: Arc, +} + +impl MarkBaseSftpServer { + pub fn new(config: Arc, provider: Arc) -> Self { + Self { config, provider } + } } impl Server for MarkBaseSftpServer { @@ -24,94 +33,70 @@ impl Server for MarkBaseSftpServer { .unwrap_or_else(|_| AuditLog::new("/tmp/sftp_audit.log").unwrap()); SshSession { - user_id: self.user_id.clone(), + user_id: None, config: self.config.clone(), + provider: self.provider.clone(), clients: Arc::new(Mutex::new(HashMap::new())), audit, - pty_sessions: Arc::new(Mutex::new(HashMap::new())), } - -async fn channel_open_session( - &mut self, - mut channel: Channel, - session: &mut Session, - ) -> Result { - log::info!("SSH channel open session: channel_id={}", channel.id()); - self.clients.lock().unwrap().insert(channel.id(), channel.clone()); - Ok(true) - } - - async fn subsystem_request( - &mut self, - channel: ChannelId, - name: &str, - session: &mut Session, - ) -> Result<(), Self::Error> { - log::info!("SSH subsystem request: channel={}, name={}", channel, name); - - if name == "sftp" { - log::info!("Starting SFTP subsystem"); - let sftp_handler = crate::sftp::handler::SftpHandler::new_with_config( - self.user_id.clone(), - self.config.clone(), - ); - let channel_stream = self.get_channel(channel).await.unwrap(); - russh_sftp::server::run(channel_stream.into_stream(), sftp_handler).await; - } else if name == "shell" { - log::info!("Starting shell subsystem"); - let shell_handler = ShellHandler::new(self.config.clone()); - let channel_stream = self.get_channel(channel).await.unwrap(); - log::warn!("Shell subsystem not fully implemented"); - } else { - log::warn!("Unknown subsystem: {}", name); - } - - Ok(()) - } - - async fn exec_request( - &mut self, - channel: ChannelId, - data: &[u8], - session: &mut Session, - ) -> Result<(), Self::Error> { - let command = String::from_utf8_lossy(data); - log::info!("SSH exec request: channel={}, command={}", channel, command); - - let command_str = command.to_string(); - - if command_str.starts_with("rsync --server") { - log::info!("Handling rsync command"); - let channel_obj = self.get_channel(channel).await; - if let Some(ch) = channel_obj { - self.handle_rsync_command(ch, &command_str).await?; - } - } else if command_str.starts_with("scp") { - log::warn!("SCP command received but not implemented: {}", command_str); - self.handle_exec_placeholder(channel, &command_str).await?; - } else { - log::warn!("Unsupported exec command: {}", command_str); - self.handle_exec_placeholder(channel, &command_str).await?; - } - - Ok(()) - } - - async fn shell_request( - &mut self, - channel: ChannelId, - session: &mut Session, - ) -> Result<(), Self::Error> { - log::info!("SSH shell request: channel={}", channel); - let shell_handler = ShellHandler::new(self.config.clone()); - let channel_obj = self.get_channel(channel).await; - if let Some(ch) = channel_obj { - log::warn!("Shell request not fully implemented"); - } - Ok(()) } } +pub struct SshSession { + user_id: Option, + config: Arc, + provider: Arc, + clients: Arc>>>, + audit: AuditLog, +} + +impl russh::server::Handler for SshSession { + type Error = anyhow::Error; + + async fn auth_request( + &mut self, + user: &str, + method: russh::server::Method, + ) -> Result { + log::info!("Auth request for user: {}, method: {:?}", user, method); + + match method { + russh::server::Method::Password { password } => { + let password_str = std::str::from_utf8(password) + .map_err(|_| anyhow::anyhow!("Invalid password encoding"))?; + + if self.provider.check_password(user, password_str)? { + log::info!("Password authentication successful for user: {}", user); + self.user_id = Some(user.to_string()); + Ok(Auth::Accept) + } else { + log::warn!("Password authentication failed for user: {}", user); + Ok(Auth::Reject { proceed_with_methods: false }) + } + } + russh::server::Method::PublicKey { key } => { + log::info!("Public key authentication for user: {}", user); + + let pubkey_bytes = key.public_key_bytes(); + let pubkey_str = base64::encode(pubkey_bytes); + + let keys = self.provider.get_public_keys(user)?; + if keys.iter().any(|k| k.contains(&pubkey_str) || k == &pubkey_str) { + log::info!("Public key authentication successful for user: {}", user); + self.user_id = Some(user.to_string()); + Ok(Auth::Accept) + } else { + log::warn!("Public key not found for user: {}", user); + Ok(Auth::Reject { proceed_with_methods: false }) + } + } + _ => { + log::warn!("Unsupported authentication method for user: {}", user); + Ok(Auth::Reject { proceed_with_methods: false }) + } + } + } + async fn channel_open_session( &mut self, channel: Channel, @@ -133,227 +118,96 @@ async fn channel_open_session( ) -> Result<(), Self::Error> { log::info!("Subsystem request: {}", name); + let user_id = self.user_id.clone().unwrap_or_else(|| "unknown".to_string()); + if name == "sftp" { - let channel = self.get_channel(channel_id).await; - let sftp_handler = crate::sftp::handler::SftpHandler::new_with_config( - &self.user_id, - self.config.clone(), - )?; + let channel = { + let clients = self.clients.lock().await; + clients.get(&channel_id).cloned() + }; - session.channel_success(channel_id)?; - - log::info!("Starting SFTP subsystem for user: {}", self.user_id); - russh_sftp::server::run(channel.into_stream(), sftp_handler).await; + if let Some(channel) = channel { + let sftp_handler = SftpHandler::new_with_config(&user_id, self.config.clone())?; + session.channel_success(channel_id)?; + log::info!("Starting SFTP subsystem for user: {}", user_id); + russh_sftp::server::run(channel.into_stream(), sftp_handler).await; + } else { + session.channel_failure(channel_id)?; + } } else if name == "shell" { - let channel = self.get_channel(channel_id).await; - - // 检查shell是否启用 if !self.config.shell.enabled { - log::warn!("Shell disabled for user {}", self.user_id); + log::warn!("Shell disabled for user {}", user_id); session.channel_failure(channel_id)?; return Ok(()); } - session.channel_success(channel_id)?; - - log::info!("Starting Shell subsystem for user: {}", self.user_id); - - // 启动shell处理（简化实现） - let shell_handler = - ShellHandler::new(&self.user_id, self.config.clone(), self.audit.clone()); - self.handle_shell_subsystem(channel, shell_handler).await?; + log::info!("Shell subsystem request for user: {}", user_id); } else { session.channel_failure(channel_id)?; } Ok(()) } -} -impl SshSession { - async fn handle_rsync_command( + async fn exec_request( &mut self, - mut channel: Channel, - command_str: &str, - ) -> Result<()> { - log::info!("Handling rsync command for user {}", self.user_id); + channel_id: ChannelId, + data: &[u8], + session: &mut Session, + ) -> Result<(), Self::Error> { + let command = String::from_utf8_lossy(data); + log::info!("SSH exec request: channel={}, command={}", channel_id, command); - // 创建rsync handler - let rsync_config = crate::rsync::RsyncConfig::default(); - let rsync_handler = crate::rsync::RsyncHandler::new( - &self.user_id, - std::sync::Arc::new(rsync_config), - &self.config.sftp.base_path, - ); + let user_id = self.user_id.clone().unwrap_or_else(|| "unknown".to_string()); - // 解析rsync命令 - let rsync_cmd = rsync_handler.parse_command(command_str)?; + session.channel_success(channel_id)?; - log::info!( - "Rsync mode: sender={}, path={}", - rsync_cmd.is_sender_mode(), - rsync_cmd.path - ); - - // 获取文件路径 - let file_path = rsync_handler.get_file_path(&rsync_cmd.path)?; - -// 简化实现：sender模式发送文件数据 -if rsync_cmd.is_sender_mode() { -log::info!("Rsync sender mode: sending file {}", file_path); - -// Step 1: 创建握手并生成checksum seed -let mut handshake = crate::rsync::protocol::RsyncHandshake::new(); -handshake.perform_sender_handshake()?; -let checksum_seed = handshake.get_checksum_seed(); - -log::info!("Checksum seed generated: {}", checksum_seed); - -// Step 2: 读取文件 -let data = tokio::fs::read(&file_path).await?; -log::info!("File read: {} bytes", data.len()); - -// Step 3: 计算block checksums（用于delta传输） -let config = rsync_handler.get_config(); -let block_checksums = if config.delta_enabled { -crate::rsync::checksum::compute_block_checksums_with_seed( -&data, -config.block_size, -checksum_seed -) -} else { -vec![] -}; - -log::info!("Block checksums computed: {} blocks", block_checksums.len()); - -// Step 4: 压缩数据 -let send_data = if config.compression { -crate::rsync::compress::compress_data(&data, config.compression_level)? -} else { -data.clone() -}; - -log::info!("Sending {} bytes (compressed)", send_data.len()); - -// Step 5: 发送数据到channel -channel.data(&send_data[..]).await?; - -// Step 6: 发送退出状态 -channel.exit_status(0).await?; - -log::info!("Rsync sender completed successfully: seed={}, blocks={}", -checksum_seed, block_checksums.len()); -} else { -log::info!("Rsync receiver mode: receiving file {}", file_path); - -// Receiver模式：不实现（技术障碍） -log::warn!("Rsync receiver mode not supported (requires channel.read())"); - -// 发送失败状态（暂时不支持） -channel.exit_status(1).await?; -} + if command.starts_with("rsync --server") { + log::info!("Rsync command for user {}", user_id); + } else if command.starts_with("scp") { + log::info!("SCP command for user {}", user_id); + } else { + log::info!("Generic command: {}", command); + } Ok(()) } - async fn handle_shell_subsystem( + async fn shell_request( &mut self, - _channel: Channel, - shell_handler: ShellHandler, - ) -> Result<()> { - log::info!("Shell subsystem started for user {}", self.user_id); + channel_id: ChannelId, + session: &mut Session, + ) -> Result<(), Self::Error> { + log::info!("Shell request: channel={}", channel_id); + + let user_id = self.user_id.clone().unwrap_or_else(|| "unknown".to_string()); - // 检查shell是否启用 if !self.config.shell.enabled { - log::warn!("Shell disabled for user {}", self.user_id); + session.channel_failure(channel_id)?; return Ok(()); } - // 创建PTY session - let mut pty_session = PtySession::new("xterm", 80, 24, shell_handler.get_shell_path())?; - - // 启动shell进程 - pty_session.start_shell().await?; - - log::info!("Shell process started for user {}", self.user_id); - - // 简化实现：等待shell进程退出 - // 完整交互需要channel.read()支持（russh API限制） - - // 清理shell进程 - pty_session.kill()?; - - Ok(()) - } - - async fn execute_command( - &mut self, - mut channel: Channel, - command: &str, - shell_handler: ShellHandler, - ) -> Result<()> { - log::info!("Executing command '{}' for user {}", command, self.user_id); - - // 执行命令 - let result = shell_handler.execute_command(command).await; - - match result { - Ok(output) => { - log::info!("Command '{}' succeeded: {} bytes", command, output.len()); - - // 发送stdout到channel - if !output.is_empty() { - channel.data(&output.as_bytes()[..]).await?; - } - - // 发送退出状态 - channel.exit_status(0).await?; - } - Err(e) => { - log::error!("Command '{}' failed: {}", command, e); - - // 发送stderr到channel - let error_msg = format!("Error: {}\r\n", e); - channel.data(&error_msg.as_bytes()[..]).await?; - - // 发送退出状态（非0表示失败） - channel.exit_status(1).await?; - } - } + session.channel_success(channel_id)?; + log::info!("Shell started for user: {}", user_id); Ok(()) } } -pub async fn run_server(config: SftpConfig, user_id: &str) -> Result<()> { +pub async fn run_server(config: SftpConfig, provider: Arc) -> Result<()> { if !config.sftp.enabled { log::warn!("SFTP server disabled in config"); return Ok(()); } let port = config.sftp.port; - let log_level = match config.logging.level.as_str() { - "debug" => log::LevelFilter::Debug, - "info" => log::LevelFilter::Info, - "warn" => log::LevelFilter::Warn, - "error" => log::LevelFilter::Error, - _ => log::LevelFilter::Info, - }; - - env_logger::builder().filter_level(log_level).init(); - let addr = format!("127.0.0.1:{}", port); log::info!("SFTP server starting on {}", addr); - log::info!("User: {}", user_id); - log::info!("Config loaded: base_path={}", config.sftp.base_path); + log::info!("Config: base_path={}", config.sftp.base_path); - println!("=== MarkBase SFTP Server ==="); + println!("=== MarkBase SFTP Server (russh) ==="); println!("Listening on {}", addr); - println!("User: {}", user_id); - println!("Config: {}", config.sftp.base_path); - println!(""); println!("Press Ctrl+C to stop"); let russh_config = russh::server::Config { @@ -361,7 +215,7 @@ pub async fn run_server(config: SftpConfig, user_id: &str) -> Result<()> { auth_rejection_time_initial: Some(Duration::from_secs(0)), keys: { let host_key_path = "config/ssh_host_ed25519_key"; - + if Path::new(host_key_path).exists() { log::info!("Loading existing SSH host key from {}", host_key_path); vec![PrivateKey::load(host_key_path).unwrap_or_else(|_| { @@ -369,25 +223,19 @@ pub async fn run_server(config: SftpConfig, user_id: &str) -> Result<()> { PrivateKey::random(&mut rand::rng(), ssh_key::Algorithm::Ed25519).unwrap() })] } else { - log::info!("Generating new SSH host key and saving to {}", host_key_path); - let key = PrivateKey::random(&mut rand::rng(), ssh_key::Algorithm::Ed25519).unwrap(); - key.save(host_key_path).unwrap_or_else(|e| { - log::warn!("Failed to save host key: {}", e); - }); - vec![key] + log::info!("Generating new SSH host key"); + vec![PrivateKey::random(&mut rand::rng(), ssh_key::Algorithm::Ed25519).unwrap()] } }, ..Default::default() }; - let mut server = MarkBaseSftpServer { - user_id: user_id.to_string(), - config: Arc::new(config), - }; + let config_arc = Arc::new(config); + let server = MarkBaseSftpServer::new(config_arc, provider); server .run_on_address(Arc::new(russh_config), ("127.0.0.1", port)) .await?; Ok(()) -} +} \ No newline at end of file