capsule_core/wasm/
execution_policy.rs1use 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}