技術

用 Rust 寫了個小工具，在本機跑 UNSW CSE 的 autotest

2024年2月15日星期四 / , ,
English正體中文
摘要：
UNSW CSE 的 autotest 和 give 只能在學校的伺服器上執行，所以我用 Rust 寫了個小工具，把本機目錄同步過去，執行指令，再把輸出帶回來。

UNSW CSE 的作業，基本上都會用到兩個指令：autotest 用來跑測試，give 用來繳交作業。但這些指令只在學校的伺服器上有，所以不管作業在哪裡寫，最後都得登入 CSE 的伺服器去執行。

這學期正在學 Rust（COMP6991），我一直想找個小東西練練手，這個需求正好合適，於是就有了 CSERun。

CSERun：在本機終端機執行 CSE 伺服器上的命令。
CSERun：在本機終端機執行 CSE 伺服器上的命令。

本機開發的麻煩

如果直接在伺服器上寫程式，其實沒什麼問題，用 vim 或 VS Code 的遠端功能都可以。但包括我在內，很多人更習慣在本機寫，編輯器和各種設定都已經準備好了，只有最後跑測試時才需要用到伺服器。

這樣一來，每改一次程式，就得繁瑣地重複以下步驟：

  1. 在本機改好程式。
  2. 把檔案傳到伺服器上，可以用 scp、rsync，或是 git push 之後再到伺服器上 pull。
  3. 用 ssh 登入，再 cd 到對應的目錄。
  4. 執行 autotest，查看結果。
  5. 回到本機，繼續修改。

想要的樣子

所以我想要的很簡單，在本機的作業目錄裡，一個指令就能搞定。

Shell
cd ~/COMP6991/lab01/exercise01
cserun "6991 autotest"

它應該自己連上伺服器、登入、把目前的目錄傳過去、在那邊執行指令，最後把輸出原樣顯示在本機終端機裡。

Mermaid diagram

每次執行都用一個新目錄

為了避免舊檔案干擾，我讓每次執行都在伺服器上使用一個新的目錄。目錄名稱是當下的時間，建在 ~/.cserun/temp/ 底下，上傳的檔案放在裡面的 container。另外也會把執行的指令記錄到 command.txt，之後想查看當時跑了什麼也方便。

Rust
// 用当前时间做目录名，比如 .cserun/temp/2024-02-14-01-10-40-224/
let temp_dir_name = chrono::Local::now()
    .format("%Y-%m-%d-%H-%M-%S-%3f")
    .to_string();
let remote_dir = format!(".cserun/temp/{}", temp_dir_name);

指令是在 container 裡面執行的，所以對它來說，就像是在一個乾淨的作業目錄裡執行一樣。

Rust
pre_exec_command.push_str(&format!("cd {}/container && ", remote_dir));
let command = format!("{}{}", pre_exec_command, conf.command);
channel.exec(&command)?;

同步檔案

連線和上傳用的是 ssh2 這個函式庫，它封裝了 libssh2，可以處理登入、SFTP 和執行指令。

至於要傳哪些檔案，我用了 ignore 這個函式庫來走訪本機目錄，它也是 ripgrep 使用的那一個。好處是它會自動讀取 .gitignore 和 .ignore，所以 target 和 node_modules 這類目錄，只要原本就列在 .gitignore 裡，就不會被傳上去，我也不用再寫一份排除規則。

Rust
let walker = WalkBuilder::new(local_path)
    .ignore(true)
    .git_ignore(true)
    .build();

for entry in walker.flatten() {
    let path = entry.path();
    if let Ok(strip_path) = path.strip_prefix(local_path) {
        let remote_path = remote_base_path.join(strip_path);
        if path.is_dir() {
            sftp.mkdir(&remote_path, 0o755)?;
        } else {
            upload_file(sftp, path, &remote_path)?;
        }
    }
}

不過有些指令不需要本機檔案，例如 6991 classrun -sturec，所以我又加了一個 --no-sync 選項。使用時會跳過上傳，只建立一個空目錄。

登入

登入資訊放在 ~/.config/cserun/config.toml 設定檔裡。第一次執行時，如果檔案不存在，程式會自動產生一份範本，然後提示你填寫。

驗證方式一共支援三種：密碼、金鑰檔案和 ssh-agent。在 Rust 裡，這三種正好可以用一個列舉來表示，各自需要的資料直接放在對應的變體裡。

Rust
pub enum Auth {
    Password(String),
    AuthKey(AuthKey),
    Agent,
}

再用 match 分別處理。這樣一來，新增或漏掉某種情況，編譯器都會提醒我。

Rust
match conf.auth {
    Auth::Password(p) => {
        sess.userauth_password(conf.username.as_str(), p.as_str())?;
    }
    Auth::AuthKey(auth_key) => {
        sess.userauth_pubkey_file(/* ... */)?;
    }
    Auth::Agent => {
        sess.userauth_agent(conf.username.as_str())?;
    }
}

如果選了密碼驗證，但設定檔裡沒有填密碼，程式就會在執行時用 rpassword 提示輸入，輸入的內容不會顯示在畫面上。

輸出

指令開始執行後，我希望輸出能像在本機執行一樣，邊跑邊顯示，而不是等全部結束才一次印出。所以我把連線設成非阻塞，然後在迴圈裡持續讀取標準輸出和標準錯誤，讀到就立刻印出。

Rust
sess.set_blocking(false);

loop {
    if channel.eof() {
        break;
    }

    let mut is_data_available = false;

    match channel.read(&mut buffer) {
        Ok(size) if size > 0 => {
            print!("{}", String::from_utf8_lossy(&buffer[..size]));
            is_data_available = true;
        }
        Ok(_) => {}
        Err(e) if e.kind() == io::ErrorKind::WouldBlock => {}
        Err(e) => return Err(e.into()),
    }

    // 标准错误也是同样的处理
    // ...

    if !is_data_available {
        // 没有数据的时候睡一会儿，不然会把 CPU 占满
        std::thread::sleep(Duration::from_millis(100));
    }
}

非阻塞時，沒有資料會回傳 WouldBlock 錯誤。這其實只是暫時沒有資料，所以要另外比對出來，直接跳過。

另外，我也加了 --env 選項，可以替指令設定環境變數。這裡遇到一個小問題：libssh2 的 setenv 在 CSE 伺服器上不一定有用，所以失敗時就改成在指令前面接上一個 export。

Rust
channel.setenv(env[0], env[1]).unwrap_or_else(|_| {
    pre_exec_command.push_str(&format!("export {}={} && ", env[0], env[1]));
});

命令列參數

命令列參數是用 clap 的 derive 寫的，只要在結構上加幾個屬性，就能產生說明文字並解析參數。

Rust
#[derive(Parser, Debug)]
#[clap(version, about, long_about = None)]
struct Args {
    /// Do not sync files before running the command
    #[clap(long, requires = "command")]
    no_sync: bool,

    /// Set environment variables, in the format of KEY:VALUE
    #[clap(long, value_name = "KEY:VALUE", value_parser = parse_env, requires = "command")]
    env: Vec<String>,

    /// The command to run on the cse server
    command: Option<String>,

    /// Show the path of config file
    #[clap(long, conflicts_with_all = &["no_sync", "env", "command"])]
    config: bool,
}

參數之間的關係也能直接在屬性裡宣告。例如 --no-sync 和 --env 必須搭配指令使用，而 --config 只能單獨使用，就不用我自己再寫判斷。

用起來

先安裝好 Rust，接著只要一個指令就能安裝。

Shell
cargo install cserun --git https://github.com/xxxbrian/cserun

裝好之後，在作業目錄裡執行就可以了。下面是不同步檔案、同時設定環境變數的例子。

Shell
cserun --no-sync --env uni:unsw --env faculty:cse "echo \$uni; echo \$faculty"
[1/5] 🌐   Connecting to cse.unsw.edu.au:22
[2/5] 🔗   Handshake successful
[3/5] 🔒   Authenticated as z5339247
[4/5] 🚫   Skipped syncing local files
[5/5] ✨   Environment variables set
[5/5] 🚀   Command sented: echo $uni; echo $faculty
=============== Output ===============
unsw
cse
======================================
Exit status: Success

平常用的話，就是 cserun "6991 autotest" 這樣一個指令。在本機改完程式，直接執行，測試結果就會出現在本機終端機裡。

  • Loading...
  • Loading...
  • Loading...
  • Loading...
  • Loading...