Feat/#61 remote exec

judge-control-app/Cargo.toml

@@ -17,3 +17,4 @@ uuid = { version = "1.0", features = ["serde", "v4"] }
 reqwest = "0.12.5"
 anyhow = "1.0"
 thiserror = "1.0.63"
+ssh2 = "0.9.4"

judge-control-app/src/main.rs

@@ -1,5 +1,5 @@
 mod custom_file;
-mod telnet;
+mod remote_exec;
 mod text_resource_repository;
 
 fn main() {

judge-control-app/src/telnet.rs → judge-control-app/src/remote_exec.rs

@@ -1 +1,2 @@
+pub mod ssh;
 pub mod traits;

judge-control-app/src/remote_exec/ssh.rs

@@ -0,0 +1,119 @@
+#[cfg(test)]
+mod tests;
+
+use super::traits::RemoteExec;
+use anyhow::{Context, Result};
+use ssh2::{Channel, Session};
+use std::io::Read;
+use std::time::Duration;
+use thiserror::Error as ThisError;
+use tokio::{
+    net::{TcpStream, ToSocketAddrs},
+    time::timeout,
+};
+
+pub struct SshConnection<AddrType: ToSocketAddrs> {
+    pub addrs: AddrType,
+    pub username: String,
+    pub password: String,
+}
+
+#[derive(ThisError, Debug)]
+pub enum SshExecError {
+    #[error("Execution in remote SSH server failed")]
+    RemoteServerError(#[from] RemoteServerError),
+    #[error("Internal error while SSH execution")]
+    InternalServerError(#[from] InternalServerError),
+}
+
+#[derive(thiserror::Error, Debug)]
+#[error(transparent)]
+pub struct RemoteServerError(anyhow::Error);
+
+#[derive(thiserror::Error, Debug)]
+#[error(transparent)]
+pub struct InternalServerError(anyhow::Error);
+
+impl<AddrType: ToSocketAddrs> RemoteExec<AddrType, SshExecError> for SshConnection<AddrType> {
+    async fn exec(
+        &mut self,
+        cmd: &str,
+        connection_time_limit: Duration,
+        execution_time_limit: Duration,
+    ) -> Result<String, SshExecError> {
+        let channel = self
+            .connect_with_timeout(connection_time_limit)
+            .await
+            .map_err(SshExecError::InternalServerError)?;
+        let output = self
+            .exec_inner_with_timeout(cmd, channel, execution_time_limit)
+            .await
+            .map_err(SshExecError::RemoteServerError)?;
+        Ok(output)
+    }
+}
+
+impl<AddrType: ToSocketAddrs> SshConnection<AddrType> {
+    // SSH接続の確立
+    async fn connect_with_timeout(
+        &self,
+        connection_time_limit: Duration,
+    ) -> Result<Channel, InternalServerError> {
+        let connect_future = async move {
+            let tcp = TcpStream::connect(&self.addrs)
+                .await
+                .context("Failed to connect to the SSH server")?;
+            let mut sess = Session::new().context("Failed to create a new SSH session")?;
+            sess.set_tcp_stream(tcp);
+            sess.handshake()
+                .context("Failed to perform SSH handshake")?;
+            sess.userauth_password(&self.username, &self.password)
+                .context("Failed to authenticate with the SSH server")?;
+            let chan = sess
+                .channel_session()
+                .context("Failed to open a new channel for SSH")?;
+            Ok(chan)
+        };
+        let timeout_future = async move {
+            timeout(connection_time_limit, connect_future)
+                .await
+                .map_err(anyhow::Error::from)
+                .context("Connection time limit exceeded")?
+        };
+        let result: Result<Channel, InternalServerError> =
+            timeout_future.await.map_err(InternalServerError);
+        result
+    }
+
+    // コマンドの実行
+    async fn exec_inner_with_timeout(
+        &self,
+        cmd: &str,
+        mut chan: Channel,
+        execution_time_limit: Duration,
+    ) -> Result<String, RemoteServerError> {
+        let exec_future = async move {
+            chan.exec(cmd)
+                .context("Failed to execute the command via SSH")?;
+            let mut output = String::new();
+            chan.read_to_string(&mut output)
+                .context("Failed to read the output from SSH")?;
+            Ok(output)
+        };
+        let timeout_future = async move {
+            let start_time = tokio::time::Instant::now();
+            let result = timeout(execution_time_limit + Duration::from_secs(1), exec_future)
+                .await
+                .map_err(anyhow::Error::from)
+                .context("Execution time limit exceeded")?;
+            let elapsed = tokio::time::Instant::now().duration_since(start_time);
+            if elapsed >= execution_time_limit {
+                return Err(anyhow::anyhow!("Execution time limit exceeded"));
+            }
+            result
+        };
+        let result: Result<String, RemoteServerError> =
+            timeout_future.await.map_err(RemoteServerError);
+        result
+    }
+}

judge-control-app/src/remote_exec/ssh/tests.rs

@@ -0,0 +1,137 @@
+#![allow(clippy::unwrap_used)]
+use super::super::{ssh::SshConnection, traits::RemoteExec};
+use anyhow::Result;
+use std::time::Duration;
+use uuid::Uuid;
+
+#[tokio::test]
+async fn test_ssh_connection() {
+    let uuid = Uuid::new_v4();
+    activate_container(uuid).await;
+    let mut ssh = SshConnection {
+        addrs: "localhost:2022",
+        username: "root".to_string(),
+        password: "password".to_string(),
+    };
+    let resp = ssh
+        .exec(
+            "cat /flag",
+            Duration::from_secs(60),
+            Duration::from_secs(60),
+        )
+        .await;
+    stop_ssh_docker_container(uuid).await.unwrap();
+    assert!(resp.is_ok());
+    assert_eq!(resp.unwrap(), "TEST_FLAG\n");
+}
+
+#[tokio::test]
+async fn test_ssh_connection_timeout() {
+    let uuid = Uuid::new_v4();
+    activate_container(uuid).await;
+    let mut ssh = SshConnection {
+        addrs: "localhost:2022",
+        username: "root".to_string(),
+        password: "password".to_string(),
+    };
+    let resp = ssh
+        .exec(
+            "sleep 1",
+            Duration::from_millis(1),
+            Duration::from_millis(1),
+        )
+        .await;
+    stop_ssh_docker_container(uuid).await.unwrap();
+    assert!(resp.is_err_and(|e| match e {
+        super::SshExecError::RemoteServerError(_) => true,
+        _ => false,
+    }));
+    return;
+}
+
+async fn activate_container(uuid: Uuid) {
+    if !check_docker_installed().await {
+        return;
+    }
+    if !check_docker_running().await {
+        if check_su_privilege().await {
+            start_docker_daemon().await.unwrap();
+        } else {
+            eprintln!("Neither docker is running nor you have su privilege");
+            return;
+        }
+    }
+    build_ssh_docker_image(uuid).await.unwrap();
+    let result = run_ssh_docker_container(uuid).await;
+    remove_docker_image(uuid).await.unwrap();
+    assert!(result.is_ok());
+}
+
+async fn check_docker_installed() -> bool {
+    let output = std::process::Command::new("docker")
+        .arg("--version")
+        .output()
+        .unwrap();
+    output.status.success()
+}
+
+async fn start_docker_daemon() -> Result<()> {
+    let _ = std::process::Command::new("systemctl")
+        .args(["start", "docker"])
+        .output()?;
+    Ok(())
+}
+
+async fn check_docker_running() -> bool {
+    let output = std::process::Command::new("systemctl")
+        .args(["is-active", "docker"])
+        .output()
+        .unwrap();
+    output.status.success()
+}
+
+async fn check_su_privilege() -> bool {
+    let output = std::process::Command::new("su")
+        .arg("-c")
+        .arg("whoami")
+        .output()
+        .unwrap();
+    output.status.success()
+}
+
+async fn build_ssh_docker_image(uuid: Uuid) -> Result<()> {
+    let _ = std::process::Command::new("docker")
+        .args(["build", "-t", &format!("ssh-server-test-{}", uuid), "."])
+        .current_dir("tests/ssh_server")
+        .output()?;
+    Ok(())
+}
+
+async fn remove_docker_image(uuid: Uuid) -> Result<()> {
+    let _ = std::process::Command::new("docker")
+        .args(["rmi", &format!("ssh-server-test-{}", uuid)])
+        .output()?;
+    Ok(())
+}
+
+async fn run_ssh_docker_container(uuid: Uuid) -> Result<()> {
+    let _ = std::process::Command::new("docker")
+        .args([
+            "run",
+            "-d",
+            "-p",
+            "2022:2022",
+            "--name",
+            &format!("ssh-server-test-{}", uuid),
+            &format!("ssh-server-test-{}", uuid),
+        ])
+        .output()?;
+    Ok(())
+}
+
+async fn stop_ssh_docker_container(uuid: Uuid) -> Result<()> {
+    let _ = std::process::Command::new("docker")
+        .args(["stop", &format!("ssh-server-test-{}", uuid)])
+        .output()?;
+    Ok(())
+}

judge-control-app/src/remote_exec/traits.rs

@@ -0,0 +1,12 @@
+use std::error::Error;
+use std::time::Duration;
+use tokio::net::ToSocketAddrs;
+
+pub trait RemoteExec<AddrType: ToSocketAddrs, ErrorType: Error> {
+    async fn exec(
+        &mut self,
+        cmd: &str,
+        connection_time_limit: Duration,
+        execution_time_limit: Duration,
+    ) -> Result<String, ErrorType>;
+}

judge-control-app/tests/ssh_server/Dockerfile

@@ -0,0 +1,27 @@
+FROM ubuntu:latest
+
+# 必要なパッケージをインストール
+RUN apt-get update && apt-get install -y \
+    openssh-server
+
+# SSH用のディレクトリを作成
+RUN mkdir /var/run/sshd
+
+# rootユーザーのパスワードを設定
+RUN echo 'root:password' | chpasswd
+
+# rootでのSSHアクセスを許可、パスワード認証を有効に
+RUN sed -i 's/#PermitRootLogin prohibit-password/PermitRootLogin yes/' /etc/ssh/sshd_config
+RUN sed -i 's/#PasswordAuthentication yes/PasswordAuthentication yes/' /etc/ssh/sshd_config
+
+# SSHデーモンのリッスンポートを2022に変更
+RUN sed -i 's/#Port 22/Port 2022/' /etc/ssh/sshd_config
+
+# テスト用のフラッグを設定
+RUN echo 'TEST_FLAG' > /flag
+
+# 2022番ポートを開放
+EXPOSE 2022
+
+# SSHサービスを起動
+CMD ["/usr/sbin/sshd", "-D"]