skill_context/
resources.rs

1//! Resource configuration types.
2//!
3//! This module defines resource limits and capabilities for execution contexts.
4
5use serde::{Deserialize, Serialize};
6use std::time::Duration;
7
8/// Resource limits and capabilities configuration.
9#[derive(Debug, Clone, Default, Serialize, Deserialize)]
10#[serde(rename_all = "snake_case")]
11pub struct ResourceConfig {
12    /// CPU limits.
13    #[serde(default, skip_serializing_if = "Option::is_none")]
14    pub cpu: Option<CpuConfig>,
15
16    /// Memory limits.
17    #[serde(default, skip_serializing_if = "Option::is_none")]
18    pub memory: Option<MemoryConfig>,
19
20    /// Network configuration.
21    #[serde(default)]
22    pub network: NetworkConfig,
23
24    /// Filesystem capabilities.
25    #[serde(default)]
26    pub filesystem: FilesystemConfig,
27
28    /// Execution limits.
29    #[serde(default)]
30    pub execution: ExecutionLimits,
31}
32
33impl ResourceConfig {
34    /// Create a new default resource configuration.
35    pub fn new() -> Self {
36        Self::default()
37    }
38
39    /// Set CPU limits.
40    pub fn with_cpu(mut self, cpu: CpuConfig) -> Self {
41        self.cpu = Some(cpu);
42        self
43    }
44
45    /// Set memory limits.
46    pub fn with_memory(mut self, memory: MemoryConfig) -> Self {
47        self.memory = Some(memory);
48        self
49    }
50
51    /// Set network configuration.
52    pub fn with_network(mut self, network: NetworkConfig) -> Self {
53        self.network = network;
54        self
55    }
56
57    /// Set filesystem configuration.
58    pub fn with_filesystem(mut self, filesystem: FilesystemConfig) -> Self {
59        self.filesystem = filesystem;
60        self
61    }
62
63    /// Set execution limits.
64    pub fn with_execution(mut self, execution: ExecutionLimits) -> Self {
65        self.execution = execution;
66        self
67    }
68
69    /// Enable network access with default settings.
70    pub fn with_network_enabled(mut self) -> Self {
71        self.network.enabled = true;
72        self
73    }
74
75    /// Disable network access.
76    pub fn with_network_disabled(mut self) -> Self {
77        self.network.enabled = false;
78        self
79    }
80
81    /// Set execution timeout.
82    pub fn with_timeout(mut self, seconds: u64) -> Self {
83        self.execution.timeout_seconds = Some(seconds);
84        self
85    }
86
87    /// Set memory limit using a size string (e.g., "512m", "2g").
88    pub fn with_memory_limit(mut self, limit: impl Into<String>) -> Self {
89        self.memory = Some(MemoryConfig {
90            limit: limit.into(),
91            swap: None,
92            reservation: None,
93        });
94        self
95    }
96
97    /// Set CPU limit.
98    pub fn with_cpu_limit(mut self, limit: impl Into<String>) -> Self {
99        self.cpu = Some(CpuConfig {
100            limit: limit.into(),
101            shares: None,
102        });
103        self
104    }
105}
106
107/// CPU configuration.
108#[derive(Debug, Clone, Serialize, Deserialize)]
109#[serde(rename_all = "snake_case")]
110pub struct CpuConfig {
111    /// CPU cores limit (e.g., "0.5", "2", "4").
112    pub limit: String,
113
114    /// CPU shares for relative priority (Docker cgroups).
115    #[serde(default, skip_serializing_if = "Option::is_none")]
116    pub shares: Option<u32>,
117}
118
119impl CpuConfig {
120    /// Create a new CPU configuration.
121    pub fn new(limit: impl Into<String>) -> Self {
122        Self {
123            limit: limit.into(),
124            shares: None,
125        }
126    }
127
128    /// Set CPU shares.
129    pub fn with_shares(mut self, shares: u32) -> Self {
130        self.shares = Some(shares);
131        self
132    }
133
134    /// Parse the limit as a float (number of cores).
135    pub fn limit_as_cores(&self) -> Option<f64> {
136        self.limit.parse().ok()
137    }
138
139    /// Convert to Docker CPU quota format (microseconds per 100ms period).
140    pub fn as_docker_quota(&self) -> Option<i64> {
141        self.limit_as_cores().map(|cores| (cores * 100_000.0) as i64)
142    }
143}
144
145/// Memory configuration.
146#[derive(Debug, Clone, Serialize, Deserialize)]
147#[serde(rename_all = "snake_case")]
148pub struct MemoryConfig {
149    /// Memory limit (e.g., "512m", "2g").
150    pub limit: String,
151
152    /// Swap limit (e.g., "1g", "0" to disable).
153    #[serde(default, skip_serializing_if = "Option::is_none")]
154    pub swap: Option<String>,
155
156    /// Memory reservation (soft limit).
157    #[serde(default, skip_serializing_if = "Option::is_none")]
158    pub reservation: Option<String>,
159}
160
161impl MemoryConfig {
162    /// Create a new memory configuration.
163    pub fn new(limit: impl Into<String>) -> Self {
164        Self {
165            limit: limit.into(),
166            swap: None,
167            reservation: None,
168        }
169    }
170
171    /// Set swap limit.
172    pub fn with_swap(mut self, swap: impl Into<String>) -> Self {
173        self.swap = Some(swap.into());
174        self
175    }
176
177    /// Disable swap.
178    pub fn without_swap(mut self) -> Self {
179        self.swap = Some("0".to_string());
180        self
181    }
182
183    /// Set memory reservation.
184    pub fn with_reservation(mut self, reservation: impl Into<String>) -> Self {
185        self.reservation = Some(reservation.into());
186        self
187    }
188
189    /// Parse the limit as bytes.
190    pub fn limit_as_bytes(&self) -> Option<u64> {
191        parse_size(&self.limit)
192    }
193
194    /// Parse the swap limit as bytes.
195    pub fn swap_as_bytes(&self) -> Option<u64> {
196        self.swap.as_ref().and_then(|s| parse_size(s))
197    }
198
199    /// Parse the reservation as bytes.
200    pub fn reservation_as_bytes(&self) -> Option<u64> {
201        self.reservation.as_ref().and_then(|s| parse_size(s))
202    }
203}
204
205/// Network configuration.
206#[derive(Debug, Clone, Serialize, Deserialize)]
207#[serde(rename_all = "snake_case")]
208pub struct NetworkConfig {
209    /// Whether network access is allowed.
210    #[serde(default)]
211    pub enabled: bool,
212
213    /// Network mode for Docker (none, bridge, host).
214    #[serde(default, skip_serializing_if = "Option::is_none")]
215    pub mode: Option<String>,
216
217    /// Allowed outbound hosts (whitelist).
218    #[serde(default, skip_serializing_if = "Option::is_none")]
219    pub allowed_hosts: Option<Vec<String>>,
220
221    /// Blocked hosts (blacklist).
222    #[serde(default, skip_serializing_if = "Option::is_none")]
223    pub blocked_hosts: Option<Vec<String>>,
224
225    /// DNS servers.
226    #[serde(default, skip_serializing_if = "Option::is_none")]
227    pub dns: Option<Vec<String>>,
228}
229
230impl Default for NetworkConfig {
231    fn default() -> Self {
232        Self {
233            enabled: false,
234            mode: None,
235            allowed_hosts: None,
236            blocked_hosts: None,
237            dns: None,
238        }
239    }
240}
241
242impl NetworkConfig {
243    /// Create a new network configuration with network enabled.
244    pub fn enabled() -> Self {
245        Self {
246            enabled: true,
247            ..Default::default()
248        }
249    }
250
251    /// Create a new network configuration with network disabled.
252    pub fn disabled() -> Self {
253        Self::default()
254    }
255
256    /// Set the network mode.
257    pub fn with_mode(mut self, mode: impl Into<String>) -> Self {
258        self.mode = Some(mode.into());
259        self
260    }
261
262    /// Set allowed hosts.
263    pub fn with_allowed_hosts(mut self, hosts: Vec<String>) -> Self {
264        self.allowed_hosts = Some(hosts);
265        self
266    }
267
268    /// Add an allowed host.
269    pub fn allow_host(mut self, host: impl Into<String>) -> Self {
270        self.allowed_hosts
271            .get_or_insert_with(Vec::new)
272            .push(host.into());
273        self
274    }
275
276    /// Set blocked hosts.
277    pub fn with_blocked_hosts(mut self, hosts: Vec<String>) -> Self {
278        self.blocked_hosts = Some(hosts);
279        self
280    }
281
282    /// Block a host.
283    pub fn block_host(mut self, host: impl Into<String>) -> Self {
284        self.blocked_hosts
285            .get_or_insert_with(Vec::new)
286            .push(host.into());
287        self
288    }
289
290    /// Set DNS servers.
291    pub fn with_dns(mut self, servers: Vec<String>) -> Self {
292        self.dns = Some(servers);
293        self
294    }
295
296    /// Check if a host is allowed.
297    pub fn is_host_allowed(&self, host: &str) -> bool {
298        if !self.enabled {
299            return false;
300        }
301
302        // Check blocked list first
303        if let Some(ref blocked) = self.blocked_hosts {
304            if blocked.iter().any(|b| host_matches(host, b)) {
305                return false;
306            }
307        }
308
309        // If there's a whitelist, host must be in it
310        if let Some(ref allowed) = self.allowed_hosts {
311            return allowed.iter().any(|a| host_matches(host, a));
312        }
313
314        // No whitelist means all non-blocked hosts are allowed
315        true
316    }
317}
318
319/// Filesystem configuration.
320#[derive(Debug, Clone, Default, Serialize, Deserialize)]
321#[serde(rename_all = "snake_case")]
322pub struct FilesystemConfig {
323    /// Read-only root filesystem.
324    #[serde(default)]
325    pub read_only_root: bool,
326
327    /// Paths that are writable (within read-only root).
328    #[serde(default, skip_serializing_if = "Vec::is_empty")]
329    pub writable_paths: Vec<String>,
330
331    /// Maximum file size that can be created.
332    #[serde(default, skip_serializing_if = "Option::is_none")]
333    pub max_file_size: Option<String>,
334
335    /// Maximum total disk usage.
336    #[serde(default, skip_serializing_if = "Option::is_none")]
337    pub max_disk_usage: Option<String>,
338}
339
340impl FilesystemConfig {
341    /// Create a new filesystem configuration.
342    pub fn new() -> Self {
343        Self::default()
344    }
345
346    /// Enable read-only root filesystem.
347    pub fn read_only(mut self) -> Self {
348        self.read_only_root = true;
349        self
350    }
351
352    /// Add a writable path.
353    pub fn with_writable_path(mut self, path: impl Into<String>) -> Self {
354        self.writable_paths.push(path.into());
355        self
356    }
357
358    /// Set maximum file size.
359    pub fn with_max_file_size(mut self, size: impl Into<String>) -> Self {
360        self.max_file_size = Some(size.into());
361        self
362    }
363
364    /// Set maximum disk usage.
365    pub fn with_max_disk_usage(mut self, size: impl Into<String>) -> Self {
366        self.max_disk_usage = Some(size.into());
367        self
368    }
369
370    /// Parse max file size as bytes.
371    pub fn max_file_size_bytes(&self) -> Option<u64> {
372        self.max_file_size.as_ref().and_then(|s| parse_size(s))
373    }
374
375    /// Parse max disk usage as bytes.
376    pub fn max_disk_usage_bytes(&self) -> Option<u64> {
377        self.max_disk_usage.as_ref().and_then(|s| parse_size(s))
378    }
379}
380
381/// Execution limits.
382#[derive(Debug, Clone, Default, Serialize, Deserialize)]
383#[serde(rename_all = "snake_case")]
384pub struct ExecutionLimits {
385    /// Maximum execution time in seconds.
386    #[serde(default, skip_serializing_if = "Option::is_none")]
387    pub timeout_seconds: Option<u64>,
388
389    /// Maximum concurrent executions.
390    #[serde(default, skip_serializing_if = "Option::is_none")]
391    pub max_concurrent: Option<u32>,
392
393    /// Rate limiting.
394    #[serde(default, skip_serializing_if = "Option::is_none")]
395    pub rate_limit: Option<RateLimit>,
396}
397
398impl ExecutionLimits {
399    /// Create new execution limits.
400    pub fn new() -> Self {
401        Self::default()
402    }
403
404    /// Set timeout.
405    pub fn with_timeout(mut self, seconds: u64) -> Self {
406        self.timeout_seconds = Some(seconds);
407        self
408    }
409
410    /// Set max concurrent executions.
411    pub fn with_max_concurrent(mut self, max: u32) -> Self {
412        self.max_concurrent = Some(max);
413        self
414    }
415
416    /// Set rate limit.
417    pub fn with_rate_limit(mut self, requests: u32, window_seconds: u32) -> Self {
418        self.rate_limit = Some(RateLimit {
419            requests,
420            window_seconds,
421        });
422        self
423    }
424
425    /// Get timeout as Duration.
426    pub fn timeout(&self) -> Option<Duration> {
427        self.timeout_seconds.map(Duration::from_secs)
428    }
429}
430
431/// Rate limiting configuration.
432#[derive(Debug, Clone, Serialize, Deserialize)]
433#[serde(rename_all = "snake_case")]
434pub struct RateLimit {
435    /// Maximum requests per window.
436    pub requests: u32,
437
438    /// Window duration in seconds.
439    pub window_seconds: u32,
440}
441
442impl RateLimit {
443    /// Create a new rate limit.
444    pub fn new(requests: u32, window_seconds: u32) -> Self {
445        Self {
446            requests,
447            window_seconds,
448        }
449    }
450
451    /// Get window as Duration.
452    pub fn window(&self) -> Duration {
453        Duration::from_secs(self.window_seconds as u64)
454    }
455
456    /// Calculate requests per second.
457    pub fn requests_per_second(&self) -> f64 {
458        if self.window_seconds == 0 {
459            0.0
460        } else {
461            self.requests as f64 / self.window_seconds as f64
462        }
463    }
464}
465
466/// Parse a size string (e.g., "512m", "2g", "1024k") into bytes.
467pub fn parse_size(s: &str) -> Option<u64> {
468    let s = s.trim().to_lowercase();
469    if s.is_empty() {
470        return None;
471    }
472
473    let (num_str, multiplier) = if s.ends_with("gb") || s.ends_with("g") {
474        let num = s.trim_end_matches(|c| c == 'g' || c == 'b');
475        (num, 1024 * 1024 * 1024)
476    } else if s.ends_with("mb") || s.ends_with("m") {
477        let num = s.trim_end_matches(|c| c == 'm' || c == 'b');
478        (num, 1024 * 1024)
479    } else if s.ends_with("kb") || s.ends_with("k") {
480        let num = s.trim_end_matches(|c| c == 'k' || c == 'b');
481        (num, 1024)
482    } else if s.ends_with('b') {
483        let num = s.trim_end_matches('b');
484        (num, 1)
485    } else {
486        // Assume bytes if no suffix
487        (s.as_str(), 1)
488    };
489
490    num_str.trim().parse::<u64>().ok().map(|n| n * multiplier)
491}
492
493/// Check if a hostname matches a pattern (supports wildcards like *.example.com).
494fn host_matches(host: &str, pattern: &str) -> bool {
495    if pattern.starts_with("*.") {
496        let suffix = &pattern[1..]; // Keep the dot
497        host.ends_with(suffix) || host == &pattern[2..]
498    } else {
499        host == pattern
500    }
501}
502
503#[cfg(test)]
504mod tests {
505    use super::*;
506
507    #[test]
508    fn test_resource_config_builder() {
509        let config = ResourceConfig::new()
510            .with_cpu_limit("2")
511            .with_memory_limit("1g")
512            .with_network_enabled()
513            .with_timeout(300);
514
515        assert!(config.cpu.is_some());
516        assert!(config.memory.is_some());
517        assert!(config.network.enabled);
518        assert_eq!(config.execution.timeout_seconds, Some(300));
519    }
520
521    #[test]
522    fn test_cpu_config() {
523        let cpu = CpuConfig::new("2.5").with_shares(1024);
524
525        assert_eq!(cpu.limit_as_cores(), Some(2.5));
526        assert_eq!(cpu.shares, Some(1024));
527        assert_eq!(cpu.as_docker_quota(), Some(250_000));
528    }
529
530    #[test]
531    fn test_memory_config() {
532        let mem = MemoryConfig::new("512m")
533            .with_swap("1g")
534            .with_reservation("256m");
535
536        assert_eq!(mem.limit_as_bytes(), Some(512 * 1024 * 1024));
537        assert_eq!(mem.swap_as_bytes(), Some(1024 * 1024 * 1024));
538        assert_eq!(mem.reservation_as_bytes(), Some(256 * 1024 * 1024));
539    }
540
541    #[test]
542    fn test_network_config() {
543        let net = NetworkConfig::enabled()
544            .with_mode("bridge")
545            .allow_host("api.example.com")
546            .allow_host("*.amazonaws.com")
547            .block_host("blocked.example.com");
548
549        assert!(net.enabled);
550        assert!(net.is_host_allowed("api.example.com"));
551        assert!(net.is_host_allowed("s3.amazonaws.com"));
552        assert!(!net.is_host_allowed("blocked.example.com"));
553        assert!(!net.is_host_allowed("other.com"));
554    }
555
556    #[test]
557    fn test_network_disabled() {
558        let net = NetworkConfig::disabled();
559        assert!(!net.is_host_allowed("any.com"));
560    }
561
562    #[test]
563    fn test_filesystem_config() {
564        let fs = FilesystemConfig::new()
565            .read_only()
566            .with_writable_path("/tmp")
567            .with_max_file_size("100m");
568
569        assert!(fs.read_only_root);
570        assert!(fs.writable_paths.contains(&"/tmp".to_string()));
571        assert_eq!(fs.max_file_size_bytes(), Some(100 * 1024 * 1024));
572    }
573
574    #[test]
575    fn test_execution_limits() {
576        let limits = ExecutionLimits::new()
577            .with_timeout(60)
578            .with_max_concurrent(10)
579            .with_rate_limit(100, 60);
580
581        assert_eq!(limits.timeout(), Some(Duration::from_secs(60)));
582        assert_eq!(limits.max_concurrent, Some(10));
583
584        let rate = limits.rate_limit.unwrap();
585        assert_eq!(rate.requests_per_second(), 100.0 / 60.0);
586    }
587
588    #[test]
589    fn test_parse_size() {
590        assert_eq!(parse_size("1024"), Some(1024));
591        assert_eq!(parse_size("1k"), Some(1024));
592        assert_eq!(parse_size("1kb"), Some(1024));
593        assert_eq!(parse_size("1m"), Some(1024 * 1024));
594        assert_eq!(parse_size("1mb"), Some(1024 * 1024));
595        assert_eq!(parse_size("1g"), Some(1024 * 1024 * 1024));
596        assert_eq!(parse_size("1gb"), Some(1024 * 1024 * 1024));
597        assert_eq!(parse_size("512M"), Some(512 * 1024 * 1024));
598        assert_eq!(parse_size(""), None);
599        assert_eq!(parse_size("invalid"), None);
600    }
601
602    #[test]
603    fn test_host_matches() {
604        assert!(host_matches("api.example.com", "api.example.com"));
605        assert!(host_matches("api.example.com", "*.example.com"));
606        assert!(host_matches("sub.api.example.com", "*.example.com"));
607        assert!(host_matches("example.com", "*.example.com"));
608        assert!(!host_matches("other.com", "*.example.com"));
609    }
610
611    #[test]
612    fn test_resource_config_serialization() {
613        let config = ResourceConfig::new()
614            .with_cpu_limit("2")
615            .with_memory_limit("1g")
616            .with_network_enabled();
617
618        let json = serde_json::to_string(&config).unwrap();
619        let deserialized: ResourceConfig = serde_json::from_str(&json).unwrap();
620
621        assert!(deserialized.cpu.is_some());
622        assert!(deserialized.memory.is_some());
623        assert!(deserialized.network.enabled);
624    }
625}
skill_context/resources.rs

skill_context/
resources.rs
