技术

用 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...