1use std::{env, io, marker::PhantomData, path::PathBuf, process::Stdio, time::Duration};
2
3use bon::bon;
4use state_shift::{impl_state, type_state};
5use tokio::{
6 fs,
7 io::{AsyncReadExt, AsyncWriteExt},
8};
9use uuid::Uuid;
10
11use crate::{Language, Metrics, Resource, Sandbox, Verdict};
12
13const MAIN: &str = "main";
14const CHECKER: &str = "checker";
15const BUFFER_SIZE: usize = 8 * 1024;
16
17pub struct Code<'a> {
18 pub language: Language,
19 pub content: &'a [u8],
20}
21
22#[type_state(
23 states = (Created, Compiled),
24 slots = (Created)
25)]
26#[derive(Default)]
27pub struct Judge {
28 pub project_path: PathBuf,
29 pub language: Language,
30 pub checker_language: Option<Language>,
31 pub is_interactive: bool,
32 pub resource: Resource,
33 pub time_limit: Duration,
34}
35
36#[bon]
37impl Judge<Created> {
38 #[builder]
39 pub async fn new<'a>(
40 main: Code<'a>,
41 checker: Option<Code<'a>>,
42 #[builder(default = false, name = "interactive")] is_interactive: bool,
43 #[builder(default)] resource: Resource,
44 #[builder(default)] time_limit: Duration,
45 ) -> io::Result<Judge<Created>> {
46 let project_path = env::temp_dir().join(Uuid::new_v4().to_string());
47 fs::create_dir(&project_path).await?;
48
49 let main_path = project_path
50 .join(MAIN)
51 .with_extension(main.language.extension);
52 fs::write(&main_path, main.content).await?;
53 if let Some(checker) = &checker {
54 let mut checker_path = project_path.join(CHECKER);
55 if checker.language.is_interpreted() {
56 checker_path.set_extension(checker.language.extension);
57 }
58 let mut checker_file = fs::OpenOptions::new()
59 .create(true)
60 .write(true)
61 .truncate(true)
62 .mode(0o755)
63 .open(&checker_path)
64 .await?;
65 checker_file.write_all(checker.content).await?;
66 checker_file.sync_all().await?;
67 }
68
69 Ok(Judge {
70 project_path,
71 language: main.language,
72 checker_language: checker.map(|checker| checker.language),
73 is_interactive,
74 resource,
75 time_limit,
76 _state: PhantomData,
77 })
78 }
79}
80
81#[impl_state]
82impl Judge {
83 #[require(Created)]
84 #[switch_to(Compiled)]
85 pub async fn compile(self) -> io::Result<Result<Judge<Compiled>, Verdict>> {
86 if let Some(mut cmd) = self.language.get_compile_command(MAIN) {
87 let mut process = cmd.current_dir(&self.project_path).spawn()?;
88 let status = process.wait().await?;
89 if !status.success() {
90 return Ok(Err(Verdict::CompilationError));
91 }
92 }
93
94 Ok(Ok(Judge {
95 project_path: self.project_path,
96 language: self.language,
97 checker_language: self.checker_language,
98 is_interactive: self.is_interactive,
99 resource: self.resource,
100 time_limit: self.time_limit,
101 }))
102 }
103
104 #[require(Compiled)]
105 pub async fn read_executable(&self) -> io::Result<Vec<u8>> {
106 let mut path = self.project_path.join(MAIN);
107 if self.language.is_interpreted() {
108 path.set_extension(self.language.extension);
109 }
110
111 fs::read(path).await
112 }
113
114 #[require(Compiled)]
115 pub async fn run(&self, input: &[u8]) -> io::Result<Metrics> {
116 let checker_language = self
117 .checker_language
118 .ok_or(io::Error::other("Missing checker"))?;
119 let mut checker = checker_language
120 .get_run_command(CHECKER)
121 .current_dir(&self.project_path)
122 .stdin(Stdio::piped())
123 .stdout(Stdio::piped())
124 .stderr(Stdio::null())
125 .spawn()?;
126 let mut cstdin = checker.stdin.take().unwrap();
127 let mut cstdout = checker.stdout.take().unwrap();
128 cstdin.write_all(input).await?;
129 cstdin.write_all(b"\n").await?;
130 cstdin.flush().await?;
131
132 let sandbox = Sandbox::new(self.resource, self.time_limit)?;
133 let mut cmd = self.language.get_run_command(MAIN);
134 cmd.current_dir(&self.project_path)
135 .stdin(Stdio::piped())
136 .stdout(Stdio::piped())
137 .stderr(Stdio::piped());
138 let mut main = sandbox.spawn(cmd)?;
139 let mut stdin = main.stdin.take().unwrap();
140 let mut stdout = main.stdout.take().unwrap();
141 let mut stderr = main.stderr.take().unwrap();
142
143 let monitor = tokio::spawn(async move { sandbox.monitor(main).await });
144 if !self.is_interactive {
145 stdin.write_all(input).await?;
146 stdin.write_all(b"\n").await?;
147 stdin.flush().await?;
148 }
149 let stdin_thread =
150 tokio::spawn(async move { tokio::io::copy(&mut cstdout, &mut stdin).await });
151 let stdout_thread = tokio::spawn(async move {
152 let mut out = vec![];
153 let mut buffer = [0u8; BUFFER_SIZE];
154 loop {
155 let n = stdout.read(&mut buffer).await?;
156 if n == 0 {
157 break;
158 }
159 if cstdin.write_all(&buffer[..n]).await.is_err() {
160 break;
161 }
162 cstdin.flush().await?;
163 out.extend_from_slice(&buffer[0..n]);
164 }
165
166 Ok::<_, io::Error>(out)
167 });
168
169 let (verdict, run_time, memory_usage) = monitor.await.unwrap()?;
170 let checker_status = checker.wait().await?;
171 drop(checker);
172
173 let _ = stdin_thread.await;
174 let stdout = stdout_thread.await.unwrap()?;
175 let mut err = vec![];
176 stderr.read_to_end(&mut err).await?;
177
178 if let Some(verdict) = verdict {
179 return Ok(Metrics {
180 verdict,
181 run_time,
182 stdout,
183 stderr: err,
184 memory_usage,
185 });
186 }
187
188 let verdict = if checker_status.success() {
189 Verdict::Accepted
190 } else {
191 Verdict::WrongAnswer
192 };
193
194 Ok(Metrics {
195 verdict,
196 run_time,
197 stdout,
198 stderr: err,
199 memory_usage,
200 })
201 }
202}