又一個Rust練手專案-wssh(SSH over Websocket Client)

orlion發表於2024-09-02
RustWebclient

原文地址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查詢出所有應用，然後在輸出中列出所有應用供使用者選擇
    App選擇

  2. 選擇應用後透過連線到websocket server，websocket server轉發到與pod的ssh連線，實現“SSH”到應用的pod的效果
    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來開發這個工具想必難度會相當低。

相關文章