From 78979574a082f1428578d51ae86106766bd4265f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=98=B3=E5=85=89=E5=B0=91=E5=B9=B4?= <849317537@qq.com>
Date: Mon, 2 Dec 2024 08:07:29 +0000
Subject: [PATCH] =?UTF-8?q?rust=E8=AF=AD=E8=A8=80=E5=8A=9F=E8=83=BD?=
 =?UTF-8?q?=E5=A2=9E=E5=BC=BA?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 test/service/handler.py | 27 +++++++++++++++++++++------
 1 file changed, 21 insertions(+), 6 deletions(-)

diff --git a/test/service/handler.py b/test/service/handler.py
index 6168191..029a7ff 100644
--- a/test/service/handler.py
+++ b/test/service/handler.py
@@ -5,11 +5,12 @@ import traceback
 
 TARGET_LANGUAGE = set(["python", "cpp", "go", "rust", "node", "javascript"])
 PROJECT_PATH = "/home/rocker/project"
-BUILD_TIMEOUT_S = 5
+BUILD_TIMEOUT_S = 3
+BASE_CODE_PATH = "/tmp/code"
 
 def go_handler():
     source_path = f"{PROJECT_PATH}/go_project/main.go"
-    subprocess.run(["cp", "/tmp/code", source_path], capture_output=True)
+    subprocess.run(["cp", BASE_CODE_PATH, source_path], capture_output=True)
     main_path = f"{PROJECT_PATH}/go_project/main"
     ret = subprocess.run(["/usr/bin/timeout", f"{BUILD_TIMEOUT_S}s", "go", "build", "-o", main_path], cwd=f"{PROJECT_PATH}/go_project", capture_output=True)
     std_err = ret.stderr.decode("utf-8")
@@ -19,7 +20,7 @@ def go_handler():
 
 def cpp_handler():
     source_path = f"{PROJECT_PATH}/cpp_project/main.cpp"
-    subprocess.run(["cp", "/tmp/code", source_path], capture_output=True)
+    subprocess.run(["cp", BASE_CODE_PATH, source_path], capture_output=True)
     main_path = f"{PROJECT_PATH}/cpp_project/main"
     ret = subprocess.run(["/usr/bin/timeout", f"{BUILD_TIMEOUT_S}s", "g++", source_path, "-o", main_path], capture_output=True)
     std_err = ret.stderr.decode("utf-8")
@@ -29,12 +30,25 @@ def cpp_handler():
 
 def node_handler():
     main_path = f"{PROJECT_PATH}/node_project/main.js"
-    subprocess.run(["cp", "/tmp/code", main_path], capture_output=True)
+    subprocess.run(["cp", BASE_CODE_PATH, main_path], capture_output=True)
     return "", main_path
 
 def rust_handler():
+    lines = []
+    is_have_main = False
+    with open(BASE_CODE_PATH, "r") as f:
+        for line in f.readlines():
+            lines.append(line)
+            if "fn" in line and "main" in line:
+                is_have_main = True
+    if is_have_main == False:
+        lines.insert(0, "fn main() {\n")
+        lines.append("}\n")
+    
     source_path = f"{PROJECT_PATH}/rust_project/src/main.rs"
-    subprocess.run(["cp", "/tmp/code", source_path], capture_output=True)
+    with open(source_path, "w") as f:
+        f.writelines(lines)
+
     ret = subprocess.run(["/usr/bin/timeout", f"{BUILD_TIMEOUT_S}s", "cargo", "build"], cwd=f"{PROJECT_PATH}/rust_project", capture_output=True)
     std_err = ret.stderr.decode("utf-8")
     if len(std_err) > 0 and ret.returncode != 0:
@@ -45,7 +59,7 @@ def rust_handler():
 
 def python_handler():
     main_path = f"{PROJECT_PATH}/python_project/main.py"
-    subprocess.run(["cp", "/tmp/code", main_path], capture_output=True)
+    subprocess.run(["cp", BASE_CODE_PATH, main_path], capture_output=True)
     return "", main_path
 
 
@@ -98,3 +112,4 @@ def main():
 
 if __name__ == "__main__":
     main()
+