stormchaser_runner_docker/container_machine/
docker_utils.rs

1use super::{crypto, DockerContainerMachine};
2use anyhow::Result;
3use bollard::container::Config;
4use bollard::service::{HostConfig, Mount};
5use std::collections::HashMap;
6use stormchaser_model::dsl::CommonContainerSpec;
7use tracing::info;
8
9impl<S> DockerContainerMachine<S> {
10    /// Attempts to detect the Docker network mode the runner is operating in.
11    pub async fn get_network_mode(&self) -> Option<String> {
12        if let Ok(hostname) = std::env::var("HOSTNAME") {
13            if let Ok(inspect) = self.docker.inspect_container(&hostname, None).await {
14                if let Some(networks) = inspect.network_settings.and_then(|ns| ns.networks) {
15                    if let Some(network_name) = networks.keys().next() {
16                        info!("Detected runner network: {}", network_name);
17                        return Some(network_name.clone());
18                    }
19                }
20            }
21        }
22        info!("Could not detect runner network, using default");
23        None
24    }
25
26    /// Builds the Docker container configuration based on the provided specifications.
27    pub fn build_container_config(
28        &self,
29        spec: &CommonContainerSpec,
30        mounts: Vec<Mount>,
31        network_mode: Option<String>,
32        storage_names: &[String],
33    ) -> Result<Config<String>> {
34        let mut env: Vec<String> = spec
35            .env
36            .clone()
37            .unwrap_or_default()
38            .into_iter()
39            .map(|e| format!("{}={}", e.name, e.value))
40            .collect();
41
42        if !storage_names.is_empty() {
43            env.push(format!("STORMCHASER_STORAGES={}", storage_names.join(" ")));
44        }
45
46        let final_image = spec.image.clone();
47
48        let mut original_full_cmd = Vec::new();
49        if let Some(cmd) = &spec.command {
50            original_full_cmd.extend(cmd.clone());
51        }
52        if let Some(args) = &spec.args {
53            original_full_cmd.extend(args.clone());
54        }
55
56        let (final_command, final_args) = if !original_full_cmd.is_empty() {
57            let step_name = &self.metadata.step_dsl.name;
58            let wrapped_script = format!(
59                "echo '========================================'; \
60                 echo 'Step Metadata: {}'; \
61                 echo \"Command: $@\"; \
62                 echo '========================================'; \
63                 \"$@\"; \
64                 RET=$?; \
65                 echo '========================================'; \
66                 echo 'Completion Status: '$RET; \
67                 echo '========================================'; \
68                 exit $RET",
69                step_name
70            );
71            let mut new_args = vec!["-c".to_string(), wrapped_script, "--".to_string()];
72            new_args.extend(original_full_cmd);
73            (Some(vec!["/bin/sh".to_string()]), Some(new_args))
74        } else {
75            (spec.command.clone(), spec.args.clone())
76        };
77
78        let step_dsl_json = serde_json::to_string(&self.metadata.step_dsl).unwrap_or_default();
79        let step_dsl_val = if let Some(key) = &self.metadata.encryption_key {
80            crypto::encrypt_state(&step_dsl_json, key)?
81        } else {
82            step_dsl_json
83        };
84
85        let mut labels = HashMap::new();
86        labels.insert("managed-by".to_string(), "stormchaser".to_string());
87        labels.insert(
88            "stormchaser-run-id".to_string(),
89            self.metadata.run_id.to_string(),
90        );
91        labels.insert(
92            "stormchaser-step-id".to_string(),
93            self.metadata.step_id.to_string(),
94        );
95        labels.insert(
96            "stormchaser.v1.io/received-at".to_string(),
97            self.metadata.received_at.to_rfc3339(),
98        );
99        labels.insert("stormchaser.v1.io/step-dsl".to_string(), step_dsl_val);
100
101        if self.metadata.encryption_key.is_some() {
102            labels.insert(
103                "stormchaser.v1.io/state-encrypted".to_string(),
104                "true".to_string(),
105            );
106        }
107
108        Ok(Config {
109            image: Some(final_image),
110            cmd: final_args,
111            entrypoint: final_command,
112            env: Some(env),
113            labels: Some(labels),
114            host_config: Some(HostConfig {
115                cpu_quota: spec.cpu.as_ref().and_then(|c| self.parse_cpu_to_quota(c)),
116                memory: spec
117                    .memory
118                    .as_ref()
119                    .and_then(|m| self.parse_memory_to_bytes(m)),
120                privileged: spec.privileged,
121                mounts: Some(mounts),
122                network_mode,
123                ..Default::default()
124            }),
125            ..Default::default()
126        })
127    }
128
129    /// Parses a CPU string limit to Docker CPU quota representation.
130    pub fn parse_cpu_to_quota(&self, cpu: &str) -> Option<i64> {
131        if let Some(m_idx) = cpu.find('m') {
132            if let Ok(m_cores) = cpu[..m_idx].parse::<i64>() {
133                return Some(m_cores * 100);
134            }
135        } else if let Ok(cores) = cpu.parse::<f64>() {
136            return Some((cores * 100_000.0) as i64);
137        }
138        None
139    }
140
141    /// Parses a memory string limit into bytes.
142    pub fn parse_memory_to_bytes(&self, memory: &str) -> Option<i64> {
143        let mem = memory.to_lowercase();
144        if let Some(idx) = mem.find(|c: char| c.is_alphabetic()) {
145            let val = mem[..idx].trim().parse::<i64>().ok()?;
146            let suffix = &mem[idx..];
147            match suffix {
148                "gi" | "g" => Some(val * 1024 * 1024 * 1024),
149                "mi" | "m" => Some(val * 1024 * 1024),
150                "ki" | "k" => Some(val * 1024),
151                _ => Some(val),
152            }
153        } else {
154            memory.parse::<i64>().ok()
155        }
156    }
157}
158
159#[cfg(test)]
160mod tests {
161    use super::super::{state, ContainerMetadata, DockerContainerMachine};
162    use super::*;
163    use stormchaser_model::dsl::Step;
164    use uuid::Uuid;
165
166    fn get_dummy_machine() -> DockerContainerMachine<state::Initialized> {
167        let docker = bollard::Docker::connect_with_local_defaults().unwrap();
168        let step_dsl: Step = serde_json::from_str(
169            r#"{
170            "name": "test",
171            "type": "RunContainer",
172            "params": {},
173            "spec": {},
174            "aggregation": [],
175            "next": [],
176            "outputs": [],
177            "reports": []
178        }"#,
179        )
180        .unwrap();
181
182        let metadata = ContainerMetadata {
183            run_id: Uuid::new_v4(),
184            step_id: Uuid::new_v4(),
185            step_dsl,
186            received_at: chrono::Utc::now(),
187            encryption_key: None,
188            storage: None,
189            test_report_urls: None,
190        };
191        DockerContainerMachine::new(docker, metadata, None)
192    }
193
194    #[test]
195    fn test_parse_cpu_to_quota() {
196        let machine = get_dummy_machine();
197        assert_eq!(machine.parse_cpu_to_quota("500m"), Some(50000));
198        assert_eq!(machine.parse_cpu_to_quota("1.0"), Some(100000));
199    }
200
201    #[test]
202    fn test_parse_memory_to_bytes() {
203        let machine = get_dummy_machine();
204        assert_eq!(
205            machine.parse_memory_to_bytes("512Mi"),
206            Some(512 * 1024 * 1024)
207        );
208        assert_eq!(
209            machine.parse_memory_to_bytes("1Gi"),
210            Some(1024 * 1024 * 1024)
211        );
212    }
213
214    #[test]
215    fn test_build_container_config() {
216        let machine = get_dummy_machine();
217        let spec = CommonContainerSpec {
218            image: "alpine:latest".into(),
219            command: Some(vec!["echo".into()]),
220            args: Some(vec!["hello".into()]),
221            env: None,
222            cpu: Some("500m".into()),
223            memory: Some("512Mi".into()),
224            privileged: Some(false),
225            storage_mounts: None,
226        };
227
228        let config = machine
229            .build_container_config(&spec, vec![], None, &[])
230            .unwrap();
231        assert_eq!(config.image, Some("alpine:latest".into()));
232    }
233}
stormchaser_runner_docker/container_machine/docker_utils.rs

stormchaser_runner_docker/container_machine/
docker_utils.rs
