技術

UNSW CSE の autotest を手元で使う小さな Rust ツール

2024年2月15日木曜日 / , ,
English日本語
要約：
UNSW CSE の autotest と give は大学のサーバーでしか使えません。そこで、手元のディレクトリを転送し、コマンドを実行して結果を持ち帰る小さなツールを Rust で作りました。

UNSW CSE の課題では、ほぼ必ず二つのコマンドを使います。テストを実行する autotest と、課題を提出する give です。ただし、どちらも大学のサーバーにしかありません。どこでコードを書いても、最後は CSE のサーバーにログインして実行する必要があります。

今学期は COMP6991 で Rust を学んでいます。練習用に何か小さなものを作りたいと思っていたので、この用途はちょうどよさそうでした。そうして作ったのが 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 なら、この三つを一つの enum で表せます。それぞれの認証に必要なデータも、対応するバリアントに持たせられます。

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 オプションも付けました。ここでは少し問題がありました。CSE のサーバーでは libssh2 の setenv が使えるとは限りません。そこで、失敗したときはコマンドの前に 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...