Complete Phase 15: Window Control + sshbuf zero-copy + SCP support
- Fix Window Control: decrease local_window on SSH_MSG_CHANNEL_DATA (critical fix) - Implement SSH_MSG_CHANNEL_WINDOW_ADJUST (OpenSSH channels.c style) - Add sshbuf.rs: zero-copy buffer management (339 lines, OpenSSH sshbuf.c reference) - Add SCP command detection (scp -t/-f support for legacy protocol) - Add handle_scp_exec() and handle_interactive_exec() for SCP/rsync - Verify: rsync 100MB transfer successful, SCP legacy protocol working - Security: All crypto using RustCrypto authoritative libraries (x25519-dalek, ed25519-dalek, aes, hmac)
This commit is contained in:
340
markbase-core/src/ssh_server/sshbuf.rs
Normal file
340
markbase-core/src/ssh_server/sshbuf.rs
Normal file
@@ -0,0 +1,340 @@
|
||||
// SSH Buffer 零拷贝实现(参考 OpenSSH sshbuf.c)
|
||||
// 提供高效的 buffer 管理,消除临时 buffer
|
||||
|
||||
use anyhow::{Result, anyhow};
|
||||
use std::io::{Read, Write};
|
||||
|
||||
/// SSH Buffer(参考 OpenSSH struct sshbuf)
|
||||
///
|
||||
/// OpenSSH 实现:
|
||||
/// ```c
|
||||
/// struct sshbuf {
|
||||
/// u_char *d; // Data (可变数据指针)
|
||||
/// size_t off; // First available byte is buf->d + buf->off
|
||||
/// size_t size; // Last byte is buf->d + buf->size - 1
|
||||
/// size_t alloc; // Total bytes allocated to buf->d
|
||||
/// };
|
||||
/// ```
|
||||
pub struct SshBuf {
|
||||
data: Vec<u8>, // Data buffer (对应 OpenSSH buf->d)
|
||||
off: usize, // Offset (对应 OpenSSH buf->off)
|
||||
size: usize, // Size (对应 OpenSSH buf->size)
|
||||
max_size: usize, // Maximum size (对应 OpenSSH buf->max_size)
|
||||
}
|
||||
|
||||
impl SshBuf {
|
||||
/// 创建新的 SSH Buffer
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
data: Vec::new(),
|
||||
off: 0,
|
||||
size: 0,
|
||||
max_size: 128 * 1024 * 1024, // 128MB (OpenSSH SSHBUF_SIZE_MAX)
|
||||
}
|
||||
}
|
||||
|
||||
/// 创建指定大小的 SSH Buffer
|
||||
pub fn with_capacity(capacity: usize) -> Self {
|
||||
Self {
|
||||
data: Vec::with_capacity(capacity),
|
||||
off: 0,
|
||||
size: 0,
|
||||
max_size: 128 * 1024 * 1024,
|
||||
}
|
||||
}
|
||||
|
||||
/// 设置最大大小
|
||||
pub fn set_max_size(&mut self, max_size: usize) -> Result<()> {
|
||||
if max_size > 128 * 1024 * 1024 {
|
||||
return Err(anyhow!("max_size too large (max 128MB)"));
|
||||
}
|
||||
self.max_size = max_size;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 获取 buffer 长度(对应 OpenSSH sshbuf_len)
|
||||
///
|
||||
/// OpenSSH: `sshbuf_len = buf->size - buf->off`
|
||||
pub fn len(&self) -> usize {
|
||||
self.size - self.off
|
||||
}
|
||||
|
||||
/// 检查 buffer 是否为空
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.len() == 0
|
||||
}
|
||||
|
||||
/// 获取可用空间(对应 OpenSSH sshbuf_avail)
|
||||
///
|
||||
/// OpenSSH: `sshbuf_avail = buf->max_size - buf->size`
|
||||
pub fn avail(&self) -> usize {
|
||||
self.max_size - self.size
|
||||
}
|
||||
|
||||
/// 获取可变指针(对应 OpenSSH sshbuf_mutable_ptr)
|
||||
///
|
||||
/// OpenSSH 实现:
|
||||
/// ```c
|
||||
/// u_char *sshbuf_mutable_ptr(const struct sshbuf *buf) {
|
||||
/// return buf->d + buf->off;
|
||||
/// }
|
||||
/// ```
|
||||
///
|
||||
/// Rust 实现:返回 `&mut [u8]` slice(零拷贝)
|
||||
pub fn mutable_ptr(&mut self) -> &mut [u8] {
|
||||
&mut self.data[self.off..self.size]
|
||||
}
|
||||
|
||||
/// 获取不可变指针(对应 OpenSSH sshbuf_ptr)
|
||||
pub fn ptr(&self) -> &[u8] {
|
||||
&self.data[self.off..self.size]
|
||||
}
|
||||
|
||||
/// 预分配空间(对应 OpenSSH sshbuf_reserve)
|
||||
///
|
||||
/// OpenSSH 实现:
|
||||
/// ```c
|
||||
/// int sshbuf_reserve(struct sshbuf *buf, size_t len, u_char **dpp) {
|
||||
/// if ((r = sshbuf_allocate(buf, len)) != 0)
|
||||
/// return r;
|
||||
///
|
||||
/// dp = buf->d + buf->size;
|
||||
/// buf->size += len;
|
||||
/// *dpp = dp;
|
||||
/// return 0;
|
||||
/// }
|
||||
/// ```
|
||||
///
|
||||
/// Rust 实现:返回 `&mut [u8]` slice(零拷贝,可直接 write)
|
||||
pub fn reserve(&mut self, len: usize) -> Result<&mut [u8]> {
|
||||
if len > self.avail() {
|
||||
return Err(anyhow!("no buffer space (avail={})", self.avail()));
|
||||
}
|
||||
|
||||
// 预分配空间
|
||||
let current_size = self.size;
|
||||
let new_size = current_size + len;
|
||||
|
||||
// 确保 Vec 有足够容量
|
||||
if new_size > self.data.len() {
|
||||
self.data.resize(new_size, 0);
|
||||
}
|
||||
|
||||
// 更新 size
|
||||
self.size = new_size;
|
||||
|
||||
// 返回新空间的 slice(零拷贝)
|
||||
Ok(&mut self.data[current_size..new_size])
|
||||
}
|
||||
|
||||
/// 消费数据(对应 OpenSSH sshbuf_consume)
|
||||
///
|
||||
/// OpenSSH 实现:
|
||||
/// ```c
|
||||
/// int sshbuf_consume(struct sshbuf *buf, size_t len) {
|
||||
/// buf->off += len;
|
||||
///
|
||||
/// if (buf->off == buf->size)
|
||||
/// buf->off = buf->size = 0;
|
||||
///
|
||||
/// return 0;
|
||||
/// }
|
||||
/// ```
|
||||
///
|
||||
/// Rust 实现:移动偏移量(零拷贝,不实际删除数据)
|
||||
pub fn consume(&mut self, len: usize) -> Result<()> {
|
||||
if len > self.len() {
|
||||
return Err(anyhow!("message incomplete (len={}, consume={})", self.len(), len));
|
||||
}
|
||||
|
||||
self.off += len;
|
||||
|
||||
// 如果 buffer 空,重置
|
||||
if self.off == self.size {
|
||||
self.off = 0;
|
||||
self.size = 0;
|
||||
|
||||
// OpenSSH: pack buffer(移除已消费的数据)
|
||||
// Rust: 我们保留 Vec,但重置指针
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 从末尾消费数据(对应 OpenSSH sshbuf_consume_end)
|
||||
///
|
||||
/// OpenSSH 实现:
|
||||
/// ```c
|
||||
/// int sshbuf_consume_end(struct sshbuf *buf, size_t len) {
|
||||
/// buf->size -= len;
|
||||
/// return 0;
|
||||
/// }
|
||||
/// ```
|
||||
pub fn consume_end(&mut self, len: usize) -> Result<()> {
|
||||
if len > self.len() {
|
||||
return Err(anyhow!("message incomplete"));
|
||||
}
|
||||
|
||||
self.size -= len;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 直接从 fd read 到 buffer(对应 OpenSSH sshbuf_read)
|
||||
///
|
||||
/// OpenSSH 实现:
|
||||
/// ```c
|
||||
/// int sshbuf_read(int fd, struct sshbuf *buf, size_t maxlen, size_t *rlen) {
|
||||
/// if ((r = sshbuf_reserve(buf, maxlen, &d)) != 0)
|
||||
/// return r;
|
||||
///
|
||||
/// rr = read(fd, d, maxlen); // 直接 read 到 buffer
|
||||
///
|
||||
/// if ((adjust = maxlen - rr) != 0)
|
||||
/// sshbuf_consume_end(buf, adjust); // 调整大小
|
||||
///
|
||||
/// return 0;
|
||||
/// }
|
||||
/// ```
|
||||
///
|
||||
/// Rust 实现:零拷贝,直接 read 到 buffer
|
||||
pub fn read_from<R: Read>(&mut self, reader: &mut R, maxlen: usize) -> Result<usize> {
|
||||
// 1. reserve 空间
|
||||
let space = self.reserve(maxlen)?;
|
||||
|
||||
// 2. 直接 read 到 buffer(零拷贝)
|
||||
let n = reader.read(space)?;
|
||||
|
||||
// 3. 调整大小(移除未使用的空间)
|
||||
if maxlen > n {
|
||||
self.consume_end(maxlen - n)?;
|
||||
}
|
||||
|
||||
Ok(n)
|
||||
}
|
||||
|
||||
/// 直接从 buffer write 到 fd(对应 OpenSSH channel_handle_wfd)
|
||||
///
|
||||
/// OpenSSH 实现:
|
||||
/// ```c
|
||||
/// buf = sshbuf_mutable_ptr(c->output); // 获取指针
|
||||
/// len = write(c->wfd, buf, dlen); // 直接 write
|
||||
/// sshbuf_consume(c->output, len); // 消费已写入的数据
|
||||
/// ```
|
||||
///
|
||||
/// Rust 实现:零拷贝,直接 write 从 buffer
|
||||
pub fn write_to<W: Write>(&mut self, writer: &mut W) -> Result<usize> {
|
||||
if self.is_empty() {
|
||||
return Ok(0);
|
||||
}
|
||||
|
||||
// 1. 获取数据指针(零拷贝)
|
||||
let data = self.ptr();
|
||||
|
||||
// 2. 直接 write(零拷贝)
|
||||
let n = writer.write(data)?;
|
||||
|
||||
// 3. 消费已写入的数据(零拷贝,只移动偏移)
|
||||
self.consume(n)?;
|
||||
|
||||
Ok(n)
|
||||
}
|
||||
|
||||
/// 添加数据(对应 OpenSSH sshbuf_put)
|
||||
///
|
||||
/// 用于不需要零拷贝的场景
|
||||
pub fn put(&mut self, data: &[u8]) -> Result<()> {
|
||||
let space = self.reserve(data.len())?;
|
||||
space.copy_from_slice(data);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 清空 buffer
|
||||
pub fn reset(&mut self) {
|
||||
self.off = 0;
|
||||
self.size = 0;
|
||||
// OpenSSH: 保留 Vec,只重置指针
|
||||
}
|
||||
|
||||
/// Debug: 打印 buffer 状态
|
||||
pub fn debug_info(&self) -> String {
|
||||
format!(
|
||||
"SshBuf: off={}, size={}, len={}, alloc={}, max_size={}",
|
||||
self.off, self.size, self.len(), self.data.len(), self.max_size
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for SshBuf {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::io::Cursor;
|
||||
|
||||
#[test]
|
||||
fn test_sshbuf_basic() {
|
||||
let mut buf = SshBuf::new();
|
||||
|
||||
// Test reserve
|
||||
let space = buf.reserve(10).unwrap();
|
||||
assert_eq!(space.len(), 10);
|
||||
assert_eq!(buf.len(), 10);
|
||||
|
||||
// Test mutable_ptr
|
||||
space[0] = 1;
|
||||
space[1] = 2;
|
||||
let ptr = buf.mutable_ptr();
|
||||
assert_eq!(ptr[0], 1);
|
||||
assert_eq!(ptr[1], 2);
|
||||
|
||||
// Test consume
|
||||
buf.consume(2).unwrap();
|
||||
assert_eq!(buf.len(), 8);
|
||||
assert_eq!(buf.off, 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sshbuf_zero_copy_read() {
|
||||
let mut buf = SshBuf::with_capacity(100);
|
||||
let mut reader = Cursor::new("hello world");
|
||||
|
||||
// 零拷贝 read
|
||||
let n = buf.read_from(&mut reader, 20).unwrap();
|
||||
assert_eq!(n, 11); // "hello world" length
|
||||
assert_eq!(buf.len(), 11);
|
||||
|
||||
// 检查数据
|
||||
let data = buf.ptr();
|
||||
assert_eq!(data, "hello world".as_bytes());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sshbuf_zero_copy_write() {
|
||||
let mut buf = SshBuf::new();
|
||||
buf.put("hello world".as_bytes()).unwrap();
|
||||
|
||||
let mut writer = Vec::new();
|
||||
|
||||
// 零拷贝 write
|
||||
let n = buf.write_to(&mut writer).unwrap();
|
||||
assert_eq!(n, 11);
|
||||
assert_eq!(buf.len(), 0); // 已消费
|
||||
|
||||
// 检查数据
|
||||
assert_eq!(writer, "hello world".as_bytes());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sshbuf_max_size() {
|
||||
let mut buf = SshBuf::new();
|
||||
buf.set_max_size(1000).unwrap();
|
||||
|
||||
// 尝试 reserve 超过 max_size
|
||||
let result = buf.reserve(2000);
|
||||
assert!(result.is_err());
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user