Skip to main content

yggdrasil/api/code_runner/
languages.rs

1//! 语言注册表与围栏代码块 info string 解析。
2//!
3//! 支持的语言通过 [`LANGUAGES`] 注册(镜像名、运行命令、默认资源限制、是否允许网络)。
4//! [`parse_fence_info`] 解析 markdown 围栏代码块的信息串,识别 `runnable` 标记与
5//! 可选的 JSON 资源覆盖(如 `python runnable {"timeout_secs":10}`)。
6//!
7//! 实际可用语言默认即注册表里的全部;若设置了 `CODE_RUNNER_LANGUAGES`
8//! 环境变量,则进一步收窄到该白名单内([`is_supported_lang`])。
9//!
10//! **别名归一化**:[`normalize_lang`] 把 `js`/`javascript` 归一为 `node`、
11//! `rs` 归一为 `rust`、`ts`/`typescript` 归一为 `bun`。归一化在 [`parse_fence_info`]
12//! 与 [`is_supported_lang`] 内完成,markdown 渲染期就把别名换成 canonical key,
13//! 下游 `data-lang` / `ExecRequest.language` / `LANGUAGES.get` 只见到 canonical
14//! 名(`python`/`node`/`go`/`rust`/`bun`),不再出现别名,避免查表落空。
15
16use std::collections::HashMap;
17use std::sync::LazyLock;
18
19use crate::infra::runner_config::{ResourceLimits, RUNNER_CONFIG};
20
21/// 单个语言的运行定义。语言名即 [`LANGUAGES`] 的 key,不再冗余存字段。
22pub struct LanguageDef {
23    /// 容器镜像(task 12 的 docker build 产出,如 `yggdrasil-runner-python:latest`)。
24    pub image: String,
25    /// 容器内执行命令(源码会注入到 `/code/main.{ext}`)。
26    pub run_cmd: String,
27    pub extension: String,
28    pub default_limits: ResourceLimits,
29    /// 该语言本身是否允许网络(与全局/请求级 allow_network 取与)。
30    pub allow_network: bool,
31}
32
33/// 内置语言注册表。新增语言时在此 `insert` 即可默认启用;
34/// 若运维需要收窄,设置 `CODE_RUNNER_LANGUAGES` 为逗号分隔列表。
35pub static LANGUAGES: LazyLock<HashMap<String, LanguageDef>> = LazyLock::new(|| {
36    let mut m = HashMap::new();
37
38    m.insert(
39        "python".to_string(),
40        LanguageDef {
41            image: "yggdrasil-runner-python:latest".to_string(),
42            // -u (unbuffered):强制 stdout/stderr 行刷新。
43            // 容器 attach 用 pipe(非 TTY),Python 默认对 pipe 做块缓冲(4KB),
44            // 导致流式输出失效——print 的内容攒在缓冲区,进程退出才一次性刷出。
45            // -u 等价于 PYTHONUNBUFFERED=1,让每行 print 立即写出。
46            run_cmd: "python -u /code/main.py".to_string(),
47            extension: "py".to_string(),
48            default_limits: ResourceLimits {
49                cpu_cores: 1.0,
50                memory_mb: 256,
51                timeout_secs: 5,
52                output_bytes: 1_048_576,
53                allow_network: false,
54            },
55            allow_network: false,
56        },
57    );
58
59    m.insert(
60        "node".to_string(),
61        LanguageDef {
62            image: "yggdrasil-runner-node:latest".to_string(),
63            run_cmd: "node /code/main.js".to_string(),
64            extension: "js".to_string(),
65            default_limits: ResourceLimits {
66                cpu_cores: 1.0,
67                memory_mb: 256,
68                timeout_secs: 5,
69                output_bytes: 1_048_576,
70                allow_network: false,
71            },
72            allow_network: false,
73        },
74    );
75
76    // 编译型语言:go run 是单条命令(内部编译 + 运行),可直接作为 run_cmd。
77    // 只读根文件系统下 $HOME/.cache 不可写,镜像已把 GOCACHE/GOTMPDIR/GOPATH
78    // 重定向到可写的 /tmp tmpfs。编译冷启动比解释型慢,timeout 提到 10s。
79    m.insert(
80        "go".to_string(),
81        LanguageDef {
82            image: "yggdrasil-runner-go:latest".to_string(),
83            run_cmd: "go run /code/main.go".to_string(),
84            extension: "go".to_string(),
85            default_limits: ResourceLimits {
86                cpu_cores: 1.0,
87                memory_mb: 256,
88                timeout_secs: 10,
89                output_bytes: 1_048_576,
90                allow_network: false,
91            },
92            allow_network: false,
93        },
94    );
95
96    // rustc 编译 + 运行是两步,但 docker.rs 注入脚本用 exec 执行 run_cmd,
97    // exec 替换 shell 进程后 "A && B" 后半段不执行,故镜像内置 run-rust.sh wrapper。
98    // rustc 内存开销大、编译慢,memory 提到 512MB、timeout 提到 15s。
99    m.insert(
100        "rust".to_string(),
101        LanguageDef {
102            image: "yggdrasil-runner-rust:latest".to_string(),
103            run_cmd: "/usr/local/bin/run-rust.sh".to_string(),
104            extension: "rs".to_string(),
105            default_limits: ResourceLimits {
106                cpu_cores: 1.0,
107                memory_mb: 512,
108                timeout_secs: 15,
109                output_bytes: 1_048_576,
110                allow_network: false,
111            },
112            allow_network: false,
113        },
114    );
115
116    // bun:原生 TypeScript 运行器(无需 tsc 预编译)。`ts`/`typescript` 别名
117    // 归一到 `bun`(见 LANG_ALIASES)。bun 镜像用官方安装脚本装在 alpine 上,
118    // 单步 `bun run /code/main.ts` 直接执行,无需 wrapper。
119    // 默认资源与 node 对齐:JS 运行时开销相近,256MB/5s 足够。
120    m.insert(
121        "bun".to_string(),
122        LanguageDef {
123            image: "yggdrasil-runner-bun:latest".to_string(),
124            run_cmd: "bun run /code/main.ts".to_string(),
125            extension: "ts".to_string(),
126            default_limits: ResourceLimits {
127                cpu_cores: 1.0,
128                memory_mb: 256,
129                timeout_secs: 5,
130                output_bytes: 1_048_576,
131                allow_network: false,
132            },
133            allow_network: false,
134        },
135    );
136
137    m
138});
139
140/// 语言别名 → canonical key 映射。
141///
142/// canonical key(`python`/`node`/`go`/`rust`/`bun`)本身不在表里——它们原样
143/// 通过。仅收录读者/作者习惯的简写:`js`/`javascript`→`node`、`rs`→`rust`、
144/// `ts`/`typescript`→`bun`。比较大小写不敏感(见 [`normalize_lang`])。
145const LANG_ALIASES: &[(&str, &str)] = &[
146    ("js", "node"),
147    ("javascript", "node"),
148    ("rs", "rust"),
149    ("ts", "bun"),
150    ("typescript", "bun"),
151];
152
153/// 把语言标识归一化为 `LANGUAGES` 注册表里的 canonical key。
154///
155/// 步骤:trim → lowercase → 查 [`LANG_ALIASES`],命中返回映射值,未命中返回
156/// trim+lowercase 后的原值(保证 canonical 名与未注册字符串都能被调用方处理)。
157///
158/// 归一化是 [`is_supported_lang`] / [`parse_fence_info`] 的前置步骤,也是
159/// `execute.rs` 在 `LANGUAGES.get` 前必须调用的——否则用别名查表会落空。
160pub fn normalize_lang(lang: &str) -> String {
161    let clean = lang.trim().to_lowercase();
162    for &(from, to) in LANG_ALIASES {
163        if clean == from {
164            return to.to_string();
165        }
166    }
167    clean
168}
169
170/// 是否支持该语言:先 [`normalize_lang`] 归一化,再查 `LANGUAGES` 注册表。
171/// 若设置了 `CODE_RUNNER_LANGUAGES`,还需同时在该白名单内(用于收窄可用语言,
172/// 白名单用 canonical key 比较——运维写 `node` 而非 `js`);未设置则全部放行。
173pub fn is_supported_lang(lang: &str) -> bool {
174    let canonical = normalize_lang(lang);
175    LANGUAGES.contains_key(&canonical)
176        && RUNNER_CONFIG
177            .languages
178            .as_ref()
179            .is_none_or(|list| list.iter().any(|l| l == &canonical))
180}
181
182/// 解析围栏代码块的 info string。
183///
184/// 格式:`<lang> [runnable|run] [ {<ResourceLimits JSON>} ]`
185///
186/// 返回 `(lang, runnable, overrides)`。`lang` 已经过 [`normalize_lang`] 归一化
187/// (`js`→`node`、`ts`→`bun` 等),下游 markdown 渲染 / ExecRequest.language /
188/// `LANGUAGES.get` 都拿到 canonical key。未知 token 静默忽略;JSON 解析失败时
189/// overrides 为 None。
190pub fn parse_fence_info(info: &str) -> (String, bool, Option<ResourceLimits>) {
191    let tokens: Vec<&str> = info.split_whitespace().collect();
192    if tokens.is_empty() {
193        return ("".to_string(), false, None);
194    }
195    let lang = normalize_lang(tokens[0]);
196    let mut runnable = false;
197    let mut overrides = None;
198
199    if tokens.len() > 1 {
200        for &tok in &tokens[1..] {
201            if tok == "runnable" || tok == "run" {
202                runnable = true;
203            } else if tok.starts_with('{') {
204                if let Ok(limits) = serde_json::from_str::<ResourceLimits>(tok) {
205                    overrides = Some(limits);
206                }
207            }
208        }
209    }
210
211    (lang, runnable, overrides)
212}
213
214#[cfg(all(test, feature = "server"))]
215mod tests {
216    use super::*;
217
218    #[test]
219    fn parse_fence_info_runnable_marker() {
220        let (lang, runnable, overrides) = parse_fence_info("python runnable");
221        assert_eq!(lang, "python");
222        assert!(runnable);
223        assert!(overrides.is_none());
224    }
225
226    #[test]
227    fn parse_fence_info_run_alias() {
228        let (lang, runnable, _) = parse_fence_info("node run");
229        assert_eq!(lang, "node");
230        assert!(runnable);
231    }
232
233    #[test]
234    fn parse_fence_info_with_overrides() {
235        let (lang, runnable, overrides) = parse_fence_info(
236            r#"node runnable {"timeout_secs":10,"memory_mb":512,"allow_network":true,"cpu_cores":1.0,"output_bytes":1024}"#,
237        );
238        assert_eq!(lang, "node");
239        assert!(runnable);
240        let limits = overrides.unwrap();
241        assert_eq!(limits.timeout_secs, 10);
242        assert_eq!(limits.memory_mb, 512);
243        assert!(limits.allow_network);
244    }
245
246    #[test]
247    fn parse_fence_info_not_runnable() {
248        let (lang, runnable, _) = parse_fence_info("rust");
249        assert_eq!(lang, "rust");
250        assert!(!runnable);
251    }
252
253    #[test]
254    fn parse_fence_info_empty() {
255        let (lang, runnable, overrides) = parse_fence_info("");
256        assert_eq!(lang, "");
257        assert!(!runnable);
258        assert!(overrides.is_none());
259    }
260
261    #[test]
262    fn parse_fence_info_case_insensitive_lang() {
263        let (lang, _, _) = parse_fence_info("PYTHON runnable");
264        assert_eq!(lang, "python");
265    }
266
267    #[test]
268    fn parse_fence_info_malformed_json_yields_none() {
269        let (lang, runnable, overrides) = parse_fence_info(r#"python runnable {not valid json}"#);
270        assert_eq!(lang, "python");
271        assert!(runnable);
272        assert!(overrides.is_none(), "malformed JSON should yield None");
273    }
274
275    #[test]
276    fn is_supported_lang_default_all_open() {
277        // 默认未设 CODE_RUNNER_LANGUAGES:注册表里的语言全部放行。
278        assert!(is_supported_lang("python"));
279        assert!(is_supported_lang("node"));
280        assert!(is_supported_lang("go"));
281        assert!(is_supported_lang("rust"));
282        // 未注册的语言仍不支持。
283        assert!(!is_supported_lang("ruby"));
284        assert!(!is_supported_lang(""));
285    }
286
287    #[test]
288    fn is_supported_lang_case_and_whitespace_insensitive() {
289        assert!(is_supported_lang(" Python "));
290        assert!(is_supported_lang("NODE"));
291    }
292
293    // ---- 别名归一化(normalize_lang / parse_fence_info / is_supported_lang 三处协同)----
294
295    #[test]
296    fn normalize_lang_canonical_passthrough() {
297        // canonical key 原样通过(仅 trim+lowercase)。
298        assert_eq!(normalize_lang("python"), "python");
299        assert_eq!(normalize_lang("node"), "node");
300        assert_eq!(normalize_lang("go"), "go");
301        assert_eq!(normalize_lang("rust"), "rust");
302        assert_eq!(normalize_lang("bun"), "bun");
303    }
304
305    #[test]
306    fn normalize_lang_aliases() {
307        // 别名 → canonical。
308        assert_eq!(normalize_lang("js"), "node");
309        assert_eq!(normalize_lang("javascript"), "node");
310        assert_eq!(normalize_lang("rs"), "rust");
311        assert_eq!(normalize_lang("ts"), "bun");
312        assert_eq!(normalize_lang("typescript"), "bun");
313    }
314
315    #[test]
316    fn normalize_lang_case_and_whitespace_insensitive() {
317        // 输入容忍:大小写、首尾空白。匹配 is_supported_lang 的既有契约。
318        assert_eq!(normalize_lang("  JS  "), "node");
319        assert_eq!(normalize_lang("JavaScript"), "node");
320        assert_eq!(normalize_lang("\tRS\t"), "rust");
321        assert_eq!(normalize_lang("TypeScript"), "bun");
322    }
323
324    #[test]
325    fn normalize_lang_unregistered_passthrough() {
326        // 未在别名表、未在注册表的语言原样返回(trim+lowercase 后),
327        // 由调用方决定是否拒绝。normalize_lang 不做存在性判断。
328        assert_eq!(normalize_lang("ruby"), "ruby");
329        assert_eq!(normalize_lang("Brainfuck"), "brainfuck");
330        assert_eq!(normalize_lang(""), "");
331    }
332
333    #[test]
334    fn is_supported_lang_accepts_aliases() {
335        // 别名经 normalize_lang 归一后命中注册表,应放行。
336        assert!(is_supported_lang("js"));
337        assert!(is_supported_lang("javascript"));
338        assert!(is_supported_lang("rs"));
339        assert!(is_supported_lang("ts"));
340        assert!(is_supported_lang("typescript"));
341    }
342
343    #[test]
344    fn is_supported_lang_accepts_bun_canonical() {
345        // bun 作为新增 canonical key,默认放行。
346        assert!(is_supported_lang("bun"));
347    }
348
349    #[test]
350    fn parse_fence_info_normalizes_alias_to_canonical() {
351        // 关键契约:parse_fence_info 返回的 lang 是 canonical key,
352        // 下游 markdown.rs 的 data-lang / class="language-{lang}" 直接用此值。
353        let (lang, runnable, _) = parse_fence_info("js runnable");
354        assert_eq!(lang, "node");
355        assert!(runnable);
356
357        let (lang, _, _) = parse_fence_info("typescript runnable");
358        assert_eq!(lang, "bun");
359
360        let (lang, _, _) = parse_fence_info("rs runnable");
361        assert_eq!(lang, "rust");
362
363        // 大小写不敏感。
364        let (lang, _, _) = parse_fence_info("JavaScript runnable");
365        assert_eq!(lang, "node");
366    }
367
368    #[test]
369    fn parse_fence_info_bun_canonical_runs_as_ts() {
370        // bun 自身是 canonical,runnable 块以 bun 执行。
371        let (lang, runnable, _) = parse_fence_info("bun runnable");
372        assert_eq!(lang, "bun");
373        assert!(runnable);
374    }
375}