From 2fb11606de83d8e53ce3c37bb586e153ceee7aac Mon Sep 17 00:00:00 2001
From: slhmy <slhmy.zzy@gmail.com>
Date: Wed, 2 Oct 2024 08:52:21 +0800
Subject: [PATCH] Add compiler timeout

---
 judge-core/Cargo.toml      |  1 +
 judge-core/src/compiler.rs | 40 ++++++++++++++++++++++++++++----------
 2 files changed, 31 insertions(+), 10 deletions(-)

diff --git a/judge-core/Cargo.toml b/judge-core/Cargo.toml
index e7dd281..35925ee 100644
--- a/judge-core/Cargo.toml
+++ b/judge-core/Cargo.toml
@@ -17,6 +17,7 @@ serde = "1"
 serde_derive = "1"
 serde_json = "1"
 serde_yaml = "0.9"
+wait-timeout = "0.2"
 
 [dev-dependencies]
 # Need to lock the version of env_logger to 0.10.0
diff --git a/judge-core/src/compiler.rs b/judge-core/src/compiler.rs
index b00bde0..8e7dfef 100644
--- a/judge-core/src/compiler.rs
+++ b/judge-core/src/compiler.rs
@@ -3,8 +3,11 @@ use crate::utils::get_pathbuf_str;
 use anyhow::anyhow;
 use serde_derive::{Deserialize, Serialize};
 use std::path::PathBuf;
+use std::process::Stdio;
+use std::time::Duration;
 use std::{fmt, fs};
 use std::{process::Command, str::FromStr};
+use wait_timeout::ChildExt;
 
 const TEMPLATE_ARG_SRC_PATH: &str = "{src_path}";
 const TEMPLATE_ARG_TARGET_PATH: &str = "{target_path}";
@@ -13,6 +16,8 @@ const RUST_COMPILE_COMMAND_TEMPLATE: &str = "rustc {src_path} -o {target_path}";
 const CPP_COMPILE_COMMAND_TEMPLATE: &str = "g++ {src_path} -o {target_path} -O2 -static";
 const PYTHON_COMPILE_COMMAND_TEMPLATE: &str = "cp {src_path} {target_path}";
 
+const COMPILE_TIMEOUT: Duration = Duration::from_secs(15);
+
 #[derive(Clone)]
 struct CommandBuilder {
     command_template: String,
@@ -162,22 +167,37 @@ impl Compiler {
             std::fs::remove_file(target_path)?;
         }
 
-        let output = Command::new("sh")
+        let mut child = Command::new("sh")
             .arg("-c")
             .arg(
                 self.command_builder
                     .get_command(vec![src_path_string, target_path_string]),
             )
             .args(self.compiler_args.iter())
-            .output()?;
-        if output.status.success() {
-            let compile_output = String::from_utf8_lossy(&output.stdout).to_string();
-            log::debug!("Compile output: {}", compile_output);
-            Ok(compile_output)
-        } else {
-            let error_output = String::from_utf8_lossy(&output.stderr).to_string();
-            log::error!("Compile error: {}", error_output);
-            Err(JudgeCoreError::CompileError(error_output))
+            .stdout(Stdio::piped())
+            .stderr(Stdio::piped())
+            .spawn()?;
+
+        match child.wait_timeout(COMPILE_TIMEOUT)? {
+            Some(status) => {
+                if status.success() {
+                    let output = child.wait_with_output()?;
+                    let compile_output = String::from_utf8_lossy(&output.stdout).to_string();
+                    log::debug!("Compile output: {}", compile_output);
+                    Ok(compile_output)
+                } else {
+                    let output = child.wait_with_output()?;
+                    let error_output = String::from_utf8_lossy(&output.stderr).to_string();
+                    log::error!("Compile error: {}", error_output);
+                    Err(JudgeCoreError::CompileError(error_output))
+                }
+            }
+            None => {
+                child.kill()?;
+                let error_output = "Compile process timed out".to_string();
+                log::error!("Compile error: {}", error_output);
+                Err(JudgeCoreError::CompileError(error_output))
+            }
         }
     }
 }