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 allowed_hosts: Vec<String>,
41
42    #[serde(default)]
43    pub env_variables: Vec<String>,
44}
45
46impl Default for ExecutionPolicy {
47    fn default() -> Self {
48        Self {
49            name: "default".to_string(),
50            compute: Compute::Medium,
51            ram: None,
52            timeout: None,
53            max_retries: 0,
54            allowed_files: Vec::new(),
55            allowed_hosts: vec!["*".to_string()],
56            env_variables: Vec::new(),
57        }
58    }
59}
60
61impl ExecutionPolicy {
62    pub fn new() -> Self {
63        Self::default()
64    }
65
66    pub fn name(mut self, name: Option<String>) -> Self {
67        if let Some(n) = name {
68            self.name = n;
69        }
70        self
71    }
72
73    pub fn compute(mut self, compute: Option<Compute>) -> Self {
74        if let Some(c) = compute {
75            self.compute = c;
76        }
77        self
78    }
79
80    pub fn ram(mut self, ram: Option<u64>) -> Self {
81        self.ram = ram;
82        self
83    }
84
85    pub fn timeout(mut self, timeout: Option<String>) -> Self {
86        self.timeout = timeout;
87        self
88    }
89
90    pub fn max_retries(mut self, max_retries: Option<u64>) -> Self {
91        if let Some(m) = max_retries {
92            self.max_retries = m;
93        }
94        self
95    }
96
97    pub fn timeout_duration(&self) -> Option<Duration> {
98        self.timeout
99            .as_ref()
100            .and_then(|s| humantime::parse_duration(s).ok())
101    }
102
103    pub fn allowed_files(mut self, files: Vec<String>) -> Self {
104        self.allowed_files = files;
105        self
106    }
107
108    pub fn allowed_hosts(mut self, allowed_hosts: Vec<String>) -> Self {
109        self.allowed_hosts = allowed_hosts;
110        self
111    }
112
113    pub fn env_variables(mut self, env_variables: Vec<String>) -> Self {
114        self.env_variables = env_variables;
115        self
116    }
117}
118
119#[cfg(test)]
120mod tests {
121    use super::*;
122
123    #[test]
124    fn test_execution_policy() {
125        let policy = ExecutionPolicy::new()
126            .name(Some("test".to_string()))
127            .compute(None)
128            .ram(Some(128))
129            .timeout(Some("60s".to_string()))
130            .max_retries(Some(3))
131            .allowed_files(vec!["/etc/passwd".to_string()])
132            .env_variables(vec!["API_KEY".to_string()]);
133
134        assert_eq!(policy.name, "test");
135        assert_eq!(policy.compute, Compute::Medium);
136        assert_eq!(policy.ram, Some(128));
137        assert_eq!(policy.timeout, Some("60s".to_string()));
138        assert_eq!(policy.max_retries, 3);
139        assert_eq!(policy.allowed_files, vec!["/etc/passwd".to_string()]);
140        assert_eq!(policy.env_variables, vec!["API_KEY".to_string()]);
141        assert_eq!(policy.allowed_hosts, vec!["*".to_string()]);
142    }
143}