yggdrasil/infra/
runner_config.rs

1use serde::{Deserialize, Serialize};
2#[cfg(feature = "server")]
3use std::env;
4#[cfg(feature = "server")]
5use std::sync::LazyLock;
6
7#[derive(Clone, Serialize, Deserialize, Debug, PartialEq)]
8pub struct ResourceLimits {
9    pub cpu_cores: f64,
10    pub memory_mb: u64,
11    pub timeout_secs: u64,
12    pub output_bytes: u64,
13    pub allow_network: bool,
14}
15
16#[cfg(feature = "server")]
17pub struct RunnerConfig {
18    pub max_cpu_cores: f64,
19    pub max_memory_mb: u64,
20    pub max_timeout_secs: u64,
21    pub max_output_bytes: u64,
22    pub max_source_bytes: u64,
23    pub allow_network: bool,
24    pub max_concurrent: usize,
25    pub queue_timeout_secs: u64,
26    pub task_ttl_secs: u64,
27    pub docker_socket_path: String,
28    /// 语言白名单。`None` 表示不限制——注册表里的所有语言均视为支持；
29    /// `Some(list)` 表示收窄到列表内（仍须同时在 LANGUAGES 注册表存在）。
30    pub languages: Option<Vec<String>>,
31}
32
33#[cfg(feature = "server")]
34fn parse_allow_network(v: &str) -> bool {
35    let l = v.to_lowercase();
36    l == "true" || l == "1" || l == "yes"
37}
38
39#[cfg(feature = "server")]
40pub static RUNNER_CONFIG: LazyLock<RunnerConfig> = LazyLock::new(|| {
41    // CODE_RUNNER_LANGUAGES 未设置时默认全开（None）：注册表里的语言均可用，
42    // 新增语言无需同步白名单。设置为逗号分隔列表则收窄到这些语言。
43    let languages = env::var("CODE_RUNNER_LANGUAGES").ok().map(|s| {
44        s.split(',')
45            .map(|t| t.trim().to_lowercase())
46            .filter(|t| !t.is_empty())
47            .collect()
48    });
49
50    RunnerConfig {
51        max_cpu_cores: env::var("CODE_RUNNER_MAX_CPU_CORES")
52            .ok()
53            .and_then(|v| v.parse().ok())
54            .unwrap_or(2.0),
55        max_memory_mb: env::var("CODE_RUNNER_MAX_MEMORY_MB")
56            .ok()
57            .and_then(|v| v.parse().ok())
58            .unwrap_or(1024),
59        max_timeout_secs: env::var("CODE_RUNNER_MAX_TIMEOUT_SECS")
60            .ok()
61            .and_then(|v| v.parse().ok())
62            .unwrap_or(30),
63        max_output_bytes: env::var("CODE_RUNNER_MAX_OUTPUT_BYTES")
64            .ok()
65            .and_then(|v| v.parse().ok())
66            .unwrap_or(1048576),
67        max_source_bytes: env::var("CODE_RUNNER_MAX_SOURCE_BYTES")
68            .ok()
69            .and_then(|v| v.parse().ok())
70            .unwrap_or(65536),
71        allow_network: env::var("CODE_RUNNER_ALLOW_NETWORK")
72            .ok()
73            .map(|v| parse_allow_network(&v))
74            .unwrap_or(false),
75        max_concurrent: env::var("CODE_RUNNER_MAX_CONCURRENT")
76            .ok()
77            .and_then(|v| v.parse().ok())
78            .unwrap_or(4),
79        queue_timeout_secs: env::var("CODE_RUNNER_QUEUE_TIMEOUT_SECS")
80            .ok()
81            .and_then(|v| v.parse().ok())
82            .unwrap_or(30),
83        task_ttl_secs: env::var("CODE_RUNNER_TASK_TTL_SECS")
84            .ok()
85            .and_then(|v| v.parse().ok())
86            .unwrap_or(300),
87        docker_socket_path: env::var("DOCKER_SOCKET_PATH")
88            .unwrap_or_else(|_| "/var/run/docker.sock".to_string()),
89        languages,
90    }
91});
92
93#[cfg(feature = "server")]
94pub fn clamp_limits(merged: ResourceLimits, lang_allows_network: bool) -> ResourceLimits {
95    clamp_limits_impl(merged, lang_allows_network, &RUNNER_CONFIG)
96}
97
98#[cfg(feature = "server")]
99fn clamp_limits_impl(
100    merged: ResourceLimits,
101    lang_allows_network: bool,
102    config: &RunnerConfig,
103) -> ResourceLimits {
104    let max_cpu = if config.max_cpu_cores.is_nan() {
105        2.0
106    } else {
107        config.max_cpu_cores
108    };
109    let min_cpu = 0.1f64.min(max_cpu);
110    let cpu_cores = if merged.cpu_cores.is_nan() {
111        min_cpu
112    } else {
113        merged.cpu_cores.clamp(min_cpu, max_cpu)
114    };
115
116    let min_mem = 16.min(config.max_memory_mb);
117    let memory_mb = merged.memory_mb.clamp(min_mem, config.max_memory_mb);
118
119    let min_timeout = 1.min(config.max_timeout_secs);
120    let timeout_secs = merged
121        .timeout_secs
122        .clamp(min_timeout, config.max_timeout_secs);
123
124    ResourceLimits {
125        cpu_cores,
126        memory_mb,
127        timeout_secs,
128        output_bytes: merged.output_bytes.min(config.max_output_bytes),
129        allow_network: merged.allow_network && config.allow_network && lang_allows_network,
130    }
131}
132
133#[cfg(all(test, feature = "server"))]
134mod tests {
135    use super::*;
136
137    #[test]
138    fn test_clamp_limits() {
139        let raw = ResourceLimits {
140            cpu_cores: 5.0,
141            memory_mb: 4096,
142            timeout_secs: 120,
143            output_bytes: 9999999,
144            allow_network: true,
145        };
146        let clamped = clamp_limits(raw, true);
147        assert!(clamped.cpu_cores <= 2.0);
148        assert!(clamped.memory_mb <= 1024);
149        assert!(clamped.timeout_secs <= 30);
150        assert!(clamped.output_bytes <= 1048576);
151        assert!(!clamped.allow_network);
152    }
153
154    #[test]
155    fn test_clamp_limits_safeguarded() {
156        let config = RunnerConfig {
157            max_cpu_cores: 0.05,
158            max_memory_mb: 8,
159            max_timeout_secs: 0,
160            max_output_bytes: 100,
161            max_source_bytes: 100,
162            allow_network: true,
163            max_concurrent: 1,
164            queue_timeout_secs: 1,
165            task_ttl_secs: 1,
166            docker_socket_path: "".to_string(),
167            languages: None,
168        };
169        let raw = ResourceLimits {
170            cpu_cores: 1.0,
171            memory_mb: 64,
172            timeout_secs: 10,
173            output_bytes: 50,
174            allow_network: true,
175        };
176        let clamped = clamp_limits_impl(raw, true, &config);
177        assert_eq!(clamped.cpu_cores, 0.05);
178        assert_eq!(clamped.memory_mb, 8);
179        assert_eq!(clamped.timeout_secs, 0);
180        assert_eq!(clamped.output_bytes, 50);
181        assert!(clamped.allow_network);
182    }
183
184    #[test]
185    fn test_parse_allow_network() {
186        assert!(parse_allow_network("true"));
187        assert!(parse_allow_network("TRUE"));
188        assert!(parse_allow_network("True"));
189        assert!(parse_allow_network("1"));
190        assert!(parse_allow_network("yes"));
191        assert!(parse_allow_network("YES"));
192        assert!(parse_allow_network("Yes"));
193        assert!(!parse_allow_network("false"));
194        assert!(!parse_allow_network("0"));
195        assert!(!parse_allow_network("no"));
196    }
197}
yggdrasil/infra/runner_config.rs

yggdrasil/infra/
runner_config.rs
