256 lines
9.5 KiB
Rust
256 lines
9.5 KiB
Rust
//! SSH 连接池（P2 连接复用）。
|
||
//!
|
||
//! # 语义
|
||
//!
|
||
//! 同一「身份」（用户名 + 主机 + 端口 + 认证材料指纹）的多个会话
|
||
//! **共享同一条 SSH 连接**：第二个会话跳过 TCP / 握手 / 认证，直接在
|
||
//! 既有连接上开新的会话通道。与 OpenSSH ControlMaster 的行为一致：
|
||
//! - 打开：首个会话建立连接；
|
||
//! - 共享：后续会话引用计数 +1；
|
||
//! - 关闭：会话关闭只减引用并关闭**自己的 shell 通道**；
|
||
//! **最后一个引用释放时**才断开底层连接（含跳板机链）。
|
||
//!
|
||
//! # 为什么 Handle 必须经由池共享
|
||
//!
|
||
//! `russh::client::Handle` 不实现 `Clone`（内含 session actor 的接收端），
|
||
//! 此前每个 `SshSession` 独占一个 Handle，无从复用。池条目持有
|
||
//! `Arc<tokio::sync::Mutex<Option<Handle>>>` 槽位，会话间共享同一 Arc；
|
||
//! tokio Mutex（而非 std）是为了允许**连接建立期间跨 `.await` 持锁**——
|
||
//! 它天然串行化了「双击两个标签同时连同一主机」的竞态：后到者等待，
|
||
//! 先到者成功后直接复用。
|
||
//!
|
||
//! # 跳板机链的归属
|
||
//!
|
||
//! 经跳板链建立的连接，其跳板 Handle 挂在**池条目**上而不是首建会话上：
|
||
//! 否则首建会话关闭时会连带剪断仍在被其他会话使用的隧道。
|
||
//! 最后一个引用释放时，跳板与目标连接一起断开。
|
||
//!
|
||
//! # 已知取舍
|
||
//!
|
||
//! - 共享连接的 keepalive / 加密参数取自**首个**建立它的会话；
|
||
//! - 共享连接断开（网络故障）时，挂在上面的所有会话一起进入 Closed——
|
||
//! 这与「它们本来就在同一条 TCP 上」的物理事实一致；
|
||
//! - 远程转发（-R）的入站路由按端口做**全局**匹配（见 forward 模块），
|
||
//! 因为入站通道总是从属连接级 Handler，而 Handler 的 session_id 属于首建会话。
|
||
use std::collections::HashMap;
|
||
use std::sync::{Arc, Mutex};
|
||
|
||
use russh::client::Handle;
|
||
use russh::Disconnect;
|
||
|
||
use super::SshHandler;
|
||
|
||
/// 可克隆的池句柄（`TerminalManager` 持有一份，每个会话 clone 一份）。
|
||
#[derive(Clone, Default)]
|
||
pub struct ConnectionPool {
|
||
conns: Arc<Mutex<HashMap<String, PooledEntry>>>,
|
||
}
|
||
|
||
struct PooledEntry {
|
||
/// 共享槽位：`None` = 连接建立中（或失败）；`Some` = 已就绪。
|
||
slot: Arc<tokio::sync::Mutex<Option<Handle<SshHandler>>>>,
|
||
/// 仍在使用此连接的会话数
|
||
refcount: usize,
|
||
/// 连接是否已就绪（同步可查；`has_handle` 用）
|
||
ready: bool,
|
||
/// 跳板机链的连接（经跳板建立时非空；随条目共享，最后释放时断开）
|
||
hops: Vec<Handle<SshHandler>>,
|
||
/// 日志用描述
|
||
label: String,
|
||
}
|
||
|
||
impl ConnectionPool {
|
||
/// 取（或创建）某身份的连接槽位，引用计数 +1。
|
||
///
|
||
/// 返回的 Arc 就是池条目里的槽位本身：会话把它存进 `SshInner.handle`，
|
||
/// 连接建立后写 `Some(handle)`，同键的其他会话即刻可见。
|
||
pub fn slot(&self, key: &str) -> Arc<tokio::sync::Mutex<Option<Handle<SshHandler>>>> {
|
||
let mut conns = self.conns.lock().unwrap_or_else(|e| e.into_inner());
|
||
let entry = conns.entry(key.to_string()).or_insert_with(|| PooledEntry {
|
||
slot: Arc::new(tokio::sync::Mutex::new(None)),
|
||
refcount: 0,
|
||
ready: false,
|
||
hops: Vec::new(),
|
||
label: String::new(),
|
||
});
|
||
entry.refcount += 1;
|
||
entry.slot.clone()
|
||
}
|
||
|
||
/// 设置日志用描述（连接建立成功后调用）。
|
||
pub fn set_label(&self, key: &str, label: &str) {
|
||
if let Some(e) = self
|
||
.conns
|
||
.lock()
|
||
.unwrap_or_else(|x| x.into_inner())
|
||
.get_mut(key)
|
||
{
|
||
e.label = label.to_string();
|
||
}
|
||
}
|
||
|
||
/// 标记连接已就绪（do_connect 写入 Handle 之后）。
|
||
pub fn mark_ready(&self, key: &str) {
|
||
if let Some(e) = self
|
||
.conns
|
||
.lock()
|
||
.unwrap_or_else(|x| x.into_inner())
|
||
.get_mut(key)
|
||
{
|
||
e.ready = true;
|
||
}
|
||
}
|
||
|
||
/// 把跳板机连接挂到池条目上（fresh 连接路径、有跳板时调用一次）。
|
||
pub fn attach_hops(&self, key: &str, hops: Vec<Handle<SshHandler>>) {
|
||
if let Some(e) = self
|
||
.conns
|
||
.lock()
|
||
.unwrap_or_else(|x| x.into_inner())
|
||
.get_mut(key)
|
||
{
|
||
e.hops = hops;
|
||
}
|
||
}
|
||
|
||
/// 连接是否已就绪（同步可查；替代原 `SshSession::has_handle` 的语义）。
|
||
pub fn is_ready(&self, key: &str) -> bool {
|
||
self.conns
|
||
.lock()
|
||
.unwrap_or_else(|e| e.into_inner())
|
||
.get(key)
|
||
.is_some_and(|e| e.ready)
|
||
}
|
||
|
||
/// 释放一个会话的引用。
|
||
///
|
||
/// 返回 `Some((槽位, 跳板连接))` 表示这是**最后一个**引用——调用方负责
|
||
/// 断开底层连接与跳板（异步任务里做，见 `kill`）。非最后引用返回 `None`，
|
||
/// 调用方只需关闭自己的 shell 通道。
|
||
pub fn release(
|
||
&self,
|
||
key: &str,
|
||
) -> Option<(
|
||
Arc<tokio::sync::Mutex<Option<Handle<SshHandler>>>>,
|
||
Vec<Handle<SshHandler>>,
|
||
)> {
|
||
let mut conns = self.conns.lock().unwrap_or_else(|e| e.into_inner());
|
||
let Some(entry) = conns.get_mut(key) else {
|
||
return None;
|
||
};
|
||
entry.refcount = entry.refcount.saturating_sub(1);
|
||
if entry.refcount > 0 {
|
||
return None;
|
||
}
|
||
// 最后一个引用：移除条目并交出断开责任
|
||
let entry = conns.remove(key)?;
|
||
Some((entry.slot, entry.hops))
|
||
}
|
||
}
|
||
|
||
/// 连接池身份键：用户名 + 主机 + 端口 + 认证方式 + 认证材料指纹。
|
||
///
|
||
/// 认证材料（密码或私钥文本）取短哈希入键——同一主机配置两份不同密钥/密码时
|
||
/// 不应共享连接（那等于用 A 的身份看了 B 的会话）。
|
||
pub fn pool_key_of(
|
||
username: &str,
|
||
host: &str,
|
||
port: u16,
|
||
auth_method: &str,
|
||
auth_material: Option<&str>,
|
||
) -> String {
|
||
let marker = match auth_material {
|
||
Some(m) => short_hash(m),
|
||
None => "none".to_string(),
|
||
};
|
||
format!("{username}|{host}:{port}|{auth_method}|{marker}")
|
||
}
|
||
|
||
/// 材料指纹：SHA-256 前 8 字节的十六进制（16 字符）。
|
||
fn short_hash(material: &str) -> String {
|
||
use sha2::{Digest, Sha256};
|
||
let digest = Sha256::digest(material.as_bytes());
|
||
digest[..8].iter().map(|b| format!("{b:02x}")).collect()
|
||
}
|
||
|
||
/// 断开一个 Handle（kill 与最后引用释放共用的收尾动作）。
|
||
pub async fn disconnect(handle: Handle<SshHandler>) {
|
||
let _ = handle
|
||
.disconnect(Disconnect::ByApplication, "closed by user", "")
|
||
.await;
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
/// 条目级语义测试：不涉及真实 Handle（槽位保持 None 即可）。
|
||
#[test]
|
||
fn acquire_increments_and_release_removes_on_last() {
|
||
let pool = ConnectionPool::default();
|
||
let s1 = pool.slot("k");
|
||
let s2 = pool.slot("k");
|
||
// 同键两次 acquire 返回同一个 Arc（这才是「共享」）
|
||
assert!(Arc::ptr_eq(&s1, &s2));
|
||
assert!(pool.release("k").is_none(), "还有 1 个引用，不应触发拆除");
|
||
let (slot, hops) = pool.release("k").expect("最后一个引用应触发拆除");
|
||
assert!(Arc::ptr_eq(&slot, &s1));
|
||
assert!(hops.is_empty());
|
||
// 移除后再次 acquire 得到全新条目
|
||
let s3 = pool.slot("k");
|
||
assert!(!Arc::ptr_eq(&s3, &s1));
|
||
pool.release("k");
|
||
}
|
||
|
||
#[test]
|
||
fn independent_keys_are_independent() {
|
||
let pool = ConnectionPool::default();
|
||
let a = pool.slot("a");
|
||
let b = pool.slot("b");
|
||
assert!(!Arc::ptr_eq(&a, &b));
|
||
assert!(pool.release("a").is_some());
|
||
assert!(pool.release("b").is_some());
|
||
}
|
||
|
||
#[test]
|
||
fn ready_flag_and_hops_follow_entry_lifecycle() {
|
||
let pool = ConnectionPool::default();
|
||
assert!(!pool.is_ready("k"));
|
||
pool.slot("k");
|
||
pool.slot("k"); // 两个会话共享
|
||
pool.mark_ready("k");
|
||
assert!(pool.is_ready("k"));
|
||
pool.attach_hops("k", Vec::new());
|
||
// 释放一个引用后条目仍在（另一个会话还在用），ready 保持
|
||
assert!(pool.release("k").is_none());
|
||
assert!(pool.is_ready("k"));
|
||
// 最后一个引用释放后条目消失
|
||
assert!(pool.release("k").is_some());
|
||
assert!(!pool.is_ready("k"));
|
||
}
|
||
|
||
#[test]
|
||
fn over_release_is_safe() {
|
||
let pool = ConnectionPool::default();
|
||
assert!(pool.release("ghost").is_none());
|
||
pool.slot("k");
|
||
pool.release("k");
|
||
// 多余的 release 不应 panic（saturating 语义）
|
||
let _ = pool.release("k");
|
||
}
|
||
|
||
#[test]
|
||
fn pool_key_distinguishes_identity() {
|
||
let k1 = pool_key_of("ops", "srv", 22, "key", Some("keytext"));
|
||
let k2 = pool_key_of("ops", "srv", 22, "key", Some("other-key"));
|
||
let k3 = pool_key_of("ops", "srv", 22, "key", Some("keytext"));
|
||
let k4 = pool_key_of("root", "srv", 22, "key", Some("keytext"));
|
||
assert_ne!(k1, k2, "不同认证材料不应共享连接");
|
||
assert_eq!(k1, k3, "相同身份应命中同一池条目");
|
||
assert_ne!(k1, k4, "不同用户不应共享连接");
|
||
// 无认证材料（理论上不出现）也不与他人混淆
|
||
let k5 = pool_key_of("ops", "srv", 22, "password", None);
|
||
assert_ne!(k1, k5);
|
||
}
|
||
}
|