又一个Rust练手项目-wssh(SSH over Websocket Client)

原文地址https://blog.fanscore.cn/a/61/

1. wssh

1.1 开发背景

公司内部的发布系统提供一个连接到k8s pod的web终端，可以在网页中连接到k8s pod内。实现原理大概为通过websocket协议代理了k8s pod ssh，然后在前端通过xterm.js+websocket实现了web终端的效果。

但是每次需要进pod内调试点东西都需要打开浏览器进到发布系统里一通点点点才能进入，而发布系统页面加载的又非常慢，所以效率非常低。

因此使用Rust实现了一个命令行工具，可以在本机终端中通过命令连接到k8s pod，实现了类似于ssh client的效果。这样一来不仅简化了我登陆pod的过程，又熟悉了Rust，还输出了篇博客。

项目地址:github.com/Orlion/wssh

1.2 效果

1. 通过-e test指定为测试环境，执行后会先调用发布系统的应用列表api查询出所有应用，然后在输出中列出所有应用供用户选择
   
   

2. 选择应用后通过连接到websocket server，websocket server转发到与pod的ssh连接，实现“SSH”到应用的pod的效果
   
   

2. 原理

公司发布系统的现状:

首先我们的发布系统提供了一个Websocket Server，这个server实际代理了到k8s pod ssh连接。然后在前端通过xterm.js模拟了一个终端，通过websocket连接到server。

wssh替换了前端：

3. 实现细节

3.1 命令行参数解析

wssh命令行参数解析使用了clap这个库

let clap_command = clap::Command::new("wssh")
    .version("0.1.0") // 指定版本号
    .author("Orlion") // 作者
    .about("SSH over Websocket 客户端")
    .arg(  // 添加命令行参数
        clap::Arg::new("env")
            .long("env")
            .short('e')
            .help("环境 test/preview")
            .value_name("ENV")
            .required(true),
    );
let matches = clap_command.get_matches();
// 获取--env参数值
let env = matches.get_one::<String>("env").expect("请输入--env参数");

3.2 发布系统登录

如1.1节所述，wssh会调用发布系统的api，发布系统需要先登录才能调用，但是调用登录api比较麻烦，还需要用户输入账号密码，因此wssh使用了github.com/thewh1teagle/rookie 库直接读取发布系统域名下的cookie，免去了输入账号密码的麻烦，非常的简单。

let domains = vec!["jumpserver.domain.com".into()];
let cookies = rookie::chrome(Some(domains)).map_err(|e| { // 使用rookie从chrome获取jumpserver的cookie
    error::from_string(format!("获取jumpserver cookie失败: {}", e.to_string()))
})?;

let mut cookie_map: HashMap<String, Cookie> = HashMap::new();
for cookie in cookies {
    if cookie.name == "sessionid" || cookie.name == "JUMPSERVER_SESS_ID" {
        cookie_map.insert(cookie.name.clone(), cookie);
    }
}

let cookies = cookie_map
    .values()
    .map(|cookie| format!("{}={}", cookie.name, cookie.value))
    .collect::<Vec<String>>()
    .join("; ");
}

3.3 命令行中输出应用列表

在命令行中输出列表供用户选择如果手动输出的话出来的效果是比较差的，因此找到了dialoguer这个库，这个库提供了一个模糊搜索的组件FuzzySelect

let app_index =
    dialoguer::FuzzySelect::with_theme(&dialoguer::theme::ColorfulTheme::default())
        .with_prompt("请选择应用") // 提示信息
        .item("0. 退出") // 为用户提供退出的选项
        .items(&app_selections) // 输出应用列表
        .default(0) // 默认选择退出
        .interact()
        .map_err(|e| error::from_string(format!("选择应用失败: {}", e.to_string())))?;

3.4 通过websocket登陆到pod

首先使用tokio_tungstenite库建立websocket连接。

let uri = format!(
    "wss://jumpserver.domain.com/ssh?ssh_token={}",
    urlencoding::encode(ssh_token),
);
let (socket, response) = tokio_tungstenite::connect_async(uri)
    .await
    .map_err(|e| error::from_string(format!("websocket连接失败: {}", e.to_string())))?;

开发这部分连接功能时踩了个“坑”，原因是刚开始开发时对Rust的异步特性不熟悉，所以想使用同步多线程的方案，所以开始使用了tungstenite::connect()创建了同步连接，后来在进行两个线程并行读写时遇到了问题，原因是connect返回的对象的read()方法和write()方法接收的是&mut self，因为Rust不允许同时存在两个可变引用，所以并发读写是不可能的。

所以后来换成了tokio_tungstenite::connect_async()函数，这个函数返回的对象提供了split()方法可以将一个连接切分成一个读句柄和一个写句柄，这样就可以并行读写了。

另外查阅文档的过程中也得知了TCP连接可拆分而TLS连接是不可拆分的，所以如果你的websocket server可以通过ws而没有强制wss的话可以使用rs-websocket这个古老的库，这个库的同步连接方法返回的TCP连接是可以拆分的。

3.5 标准输出的调整

要在本地输出远程ssh server输出的内容之前还需要做以下三个调整。

1. 发送window-change请求
   
   本地终端窗口大小初始化和发生变更时都需要同步ssh server的，以便获得一致的显示效果，如果不发送可能会导致显示内容被截断或者格式不正确，并且vim等命令依赖于准确的终端尺寸来显示界面。
2. 将标准输出设置为raw模式。在raw模式下，标准输出表现为
   • 没有行缓存，会逐字节输出
   • 不会回显输入，必须由程序写入
   • 输出未规范化（例如，\n 表示“向下一行”，而不是“换行符”）

let mut stdout = std::io::stdout().into_raw_mode()

4. 总结

通过这个项目又加深了对Rust的理解，过程中还首次用到了反人类的生命周期标注‍♀️（虽然后面简化掉了)，收获很大，Rust远比看上去简单。

同时越发感慨Go的简易性，Go的协程结合channel、select等组件无疑极大降低了并发编程的难度，如果使用Go来开发这个工具想必难度会相当低。