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 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}