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")]
34pub static RUNNER_CONFIG: LazyLock<RunnerConfig> = LazyLock::new(|| {
35    let cfg = crate::config::runner();
36    RunnerConfig {
37        max_cpu_cores: cfg.max_cpu_cores,
38        max_memory_mb: cfg.max_memory_mb as u64,
39        max_timeout_secs: cfg.max_timeout_secs as u64,
40        max_output_bytes: cfg.max_output_bytes,
41        max_source_bytes: cfg.max_source_bytes,
42        allow_network: cfg.allow_network,
43        max_concurrent: cfg.max_concurrent as usize,
44        queue_timeout_secs: cfg.queue_timeout_secs as u64,
45        task_ttl_secs: cfg.task_ttl_secs as u64,
46        docker_socket_path: env::var("DOCKER_SOCKET_PATH")
47            .unwrap_or_else(|_| "/var/run/docker.sock".to_string()),
48        // languages 是逗号分隔字符串，拆成 Vec<String>。
49        languages: cfg.languages.map(|s| {
50            s.split(',')
51                .map(|t| t.trim().to_lowercase())
52                .filter(|t| !t.is_empty())
53                .collect()
54        }),
55    }
56});
57
58#[cfg(feature = "server")]
59pub fn clamp_limits(merged: ResourceLimits, lang_allows_network: bool) -> ResourceLimits {
60    clamp_limits_impl(merged, lang_allows_network, &RUNNER_CONFIG)
61}
62
63#[cfg(feature = "server")]
64fn clamp_limits_impl(
65    merged: ResourceLimits,
66    lang_allows_network: bool,
67    config: &RunnerConfig,
68) -> ResourceLimits {
69    let max_cpu = if config.max_cpu_cores.is_nan() {
70        2.0
71    } else {
72        config.max_cpu_cores
73    };
74    let min_cpu = 0.1f64.min(max_cpu);
75    let cpu_cores = if merged.cpu_cores.is_nan() {
76        min_cpu
77    } else {
78        merged.cpu_cores.clamp(min_cpu, max_cpu)
79    };
80
81    let min_mem = 16.min(config.max_memory_mb);
82    let memory_mb = merged.memory_mb.clamp(min_mem, config.max_memory_mb);
83
84    let min_timeout = 1.min(config.max_timeout_secs);
85    let timeout_secs = merged
86        .timeout_secs
87        .clamp(min_timeout, config.max_timeout_secs);
88
89    ResourceLimits {
90        cpu_cores,
91        memory_mb,
92        timeout_secs,
93        output_bytes: merged.output_bytes.min(config.max_output_bytes),
94        allow_network: merged.allow_network && config.allow_network && lang_allows_network,
95    }
96}
97
98#[cfg(all(test, feature = "server"))]
99mod tests {
100    use super::*;
101
102    #[test]
103    fn test_clamp_limits() {
104        let raw = ResourceLimits {
105            cpu_cores: 5.0,
106            memory_mb: 4096,
107            timeout_secs: 120,
108            output_bytes: 9999999,
109            allow_network: true,
110        };
111        let clamped = clamp_limits(raw, true);
112        assert!(clamped.cpu_cores <= 2.0);
113        assert!(clamped.memory_mb <= 1024);
114        assert!(clamped.timeout_secs <= 30);
115        assert!(clamped.output_bytes <= 1048576);
116        assert!(!clamped.allow_network);
117    }
118
119    #[test]
120    fn test_clamp_limits_safeguarded() {
121        let config = RunnerConfig {
122            max_cpu_cores: 0.05,
123            max_memory_mb: 8,
124            max_timeout_secs: 0,
125            max_output_bytes: 100,
126            max_source_bytes: 100,
127            allow_network: true,
128            max_concurrent: 1,
129            queue_timeout_secs: 1,
130            task_ttl_secs: 1,
131            docker_socket_path: "".to_string(),
132            languages: None,
133        };
134        let raw = ResourceLimits {
135            cpu_cores: 1.0,
136            memory_mb: 64,
137            timeout_secs: 10,
138            output_bytes: 50,
139            allow_network: true,
140        };
141        let clamped = clamp_limits_impl(raw, true, &config);
142        assert_eq!(clamped.cpu_cores, 0.05);
143        assert_eq!(clamped.memory_mb, 8);
144        assert_eq!(clamped.timeout_secs, 0);
145        assert_eq!(clamped.output_bytes, 50);
146        assert!(clamped.allow_network);
147    }
148}
yggdrasil/infra/runner_config.rs

yggdrasil/infra/
runner_config.rs
