Skip to main content

capsule_core/wasm/
execution_policy.rs

1use serde::{Deserialize, Serialize};
2use std::time::Duration;
3
4#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
5#[serde(rename_all = "snake_case")]
6pub enum Compute {
7    Low,
8    Medium,
9    High,
10    Custom(u64),
11}
12
13impl Compute {
14    pub fn as_fuel(&self) -> u64 {
15        match self {
16            Compute::Low => 100_000_000,
17            Compute::Medium => 2_000_000_000,
18            Compute::High => 50_000_000_000,
19            Compute::Custom(fuel) => *fuel,
20        }
21    }
22
23    pub fn to_u64(&self) -> u64 {
24        self.as_fuel()
25    }
26}
27
28#[derive(Clone, Debug, Serialize, Deserialize)]
29pub struct ExecutionPolicy {
30    pub name: String,
31    pub compute: Compute,
32    pub ram: Option<u64>,
33    pub timeout: Option<String>,
34    pub max_retries: u64,
35
36    #[serde(default)]
37    pub allowed_files: Vec<String>,
38
39    #[serde(default)]
40    pub env_variables: Vec<String>,
41}
42
43impl Default for ExecutionPolicy {
44    fn default() -> Self {
45        Self {
46            name: "default".to_string(),
47            compute: Compute::Medium,
48            ram: None,
49            timeout: None,
50            max_retries: 0,
51            allowed_files: Vec::new(),
52            env_variables: Vec::new(),
53        }
54    }
55}
56
57impl ExecutionPolicy {
58    pub fn new() -> Self {
59        Self::default()
60    }
61
62    pub fn name(mut self, name: Option<String>) -> Self {
63        if let Some(n) = name {
64            self.name = n;
65        }
66        self
67    }
68
69    pub fn compute(mut self, compute: Option<Compute>) -> Self {
70        if let Some(c) = compute {
71            self.compute = c;
72        }
73        self
74    }
75
76    pub fn ram(mut self, ram: Option<u64>) -> Self {
77        self.ram = ram;
78        self
79    }
80
81    pub fn timeout(mut self, timeout: Option<String>) -> Self {
82        self.timeout = timeout;
83        self
84    }
85
86    pub fn max_retries(mut self, max_retries: Option<u64>) -> Self {
87        if let Some(m) = max_retries {
88            self.max_retries = m;
89        }
90        self
91    }
92
93    pub fn timeout_duration(&self) -> Option<Duration> {
94        self.timeout
95            .as_ref()
96            .and_then(|s| humantime::parse_duration(s).ok())
97    }
98
99    pub fn allowed_files(mut self, files: Vec<String>) -> Self {
100        self.allowed_files = files;
101        self
102    }
103
104    pub fn env_variables(mut self, env_variables: Vec<String>) -> Self {
105        self.env_variables = env_variables;
106        self
107    }
108}
109
110#[cfg(test)]
111mod tests {
112    use super::*;
113
114    #[test]
115    fn test_execution_policy() {
116        let policy = ExecutionPolicy::new()
117            .name(Some("test".to_string()))
118            .compute(None)
119            .ram(Some(128))
120            .timeout(Some("60s".to_string()))
121            .max_retries(Some(3))
122            .allowed_files(vec!["/etc/passwd".to_string()])
123            .env_variables(vec!["API_KEY".to_string()]);
124
125        assert_eq!(policy.name, "test");
126        assert_eq!(policy.compute, Compute::Medium);
127        assert_eq!(policy.ram, Some(128));
128        assert_eq!(policy.timeout, Some("60s".to_string()));
129        assert_eq!(policy.max_retries, 3);
130        assert_eq!(policy.allowed_files, vec!["/etc/passwd".to_string()]);
131        assert_eq!(policy.env_variables, vec!["API_KEY".to_string()]);
132    }
133}