blob: 65f15051b6929100f6c46a6625422c9dd94e5278 [file] [log] [blame]
Mohammed Naser9a6c35d2025-02-13 13:23:51 -05001extern crate tar;
2
3use bollard::container::{Config, CreateContainerOptions, StartContainerOptions};
4use bollard::exec::{CreateExecOptions, StartExecResults};
5use bollard::Docker;
6use bytes::{BufMut, BytesMut};
7use futures_util::stream::StreamExt;
8use futures_util::TryStreamExt;
9use passwd::PasswdEntry;
10use rand::Rng;
11use std::collections::HashMap;
12use std::io::Read;
13use thiserror::Error;
14
15#[derive(Debug, Error)]
16pub enum DockerContainerGuardError {
17 #[error("Docker API error: {0}")]
18 DockerError(#[from] bollard::errors::Error),
19
20 #[error("File not found: {0}")]
21 FileNotFound(String),
22
23 #[error("Failed to extract file: {0}")]
24 FileExtractionFailed(#[from] std::io::Error),
25
26 #[error("Too many files in tarball")]
27 TooManyFilesInTarball,
28
29 #[error("Failed to parse password entry: {0}")]
30 FailedToParsePasswdEntry(#[from] passwd::PasswdEntryError),
31
32 #[error("User not found: {0}")]
33 UserNotFound(String),
34}
35
36#[derive(Debug)]
37pub struct DockerContainerGuard {
38 pub id: String,
39 pub image: String,
40 docker: Docker,
41}
42
43impl DockerContainerGuard {
44 // Spawns a new container using Bollard.
45 //
46 // The container is automatically cleaned up when the guard goes out of scope.
47 pub async fn spawn(image_name: &str) -> Result<Self, DockerContainerGuardError> {
48 let docker = Docker::connect_with_local_defaults()?;
49
50 let container_name: String = rand::thread_rng()
51 .sample_iter(&rand::distributions::Alphanumeric)
52 .take(10)
53 .map(char::from)
54 .collect();
55
56 docker
57 .create_image(
58 Some(bollard::image::CreateImageOptions {
59 from_image: image_name,
60 ..Default::default()
61 }),
62 None,
63 None,
64 )
65 .try_collect::<Vec<_>>()
66 .await?;
67
68 let container = docker
69 .create_container(
70 Some(CreateContainerOptions {
71 name: container_name,
72 ..Default::default()
73 }),
74 Config {
75 image: Some(image_name),
76 cmd: Some(vec!["sh"]),
77 tty: Some(true),
78 ..Default::default()
79 },
80 )
81 .await?;
82
83 docker
84 .start_container(&container.id, None::<StartContainerOptions<String>>)
85 .await?;
86
87 Ok(Self {
88 id: container.id,
89 image: image_name.to_string(),
90 docker,
91 })
92 }
93
94 /// Executes a command inside the container using Bollard.
95 ///
96 /// Returns the output as a String.
97 pub async fn exec(&self, cmd: Vec<&str>) -> Result<String, bollard::errors::Error> {
98 let exec_instance = self
99 .docker
100 .create_exec(
101 &self.id,
102 CreateExecOptions {
103 attach_stdout: Some(true),
104 attach_stderr: Some(true),
105 cmd: Some(cmd.iter().map(|s| s.to_string()).collect()),
106 ..Default::default()
107 },
108 )
109 .await?;
110 let start_exec_result = self.docker.start_exec(&exec_instance.id, None).await?;
111
112 if let StartExecResults::Attached {
113 output: out_stream, ..
114 } = start_exec_result
115 {
116 let output = out_stream
117 .filter_map(|chunk| async {
118 match chunk {
119 Ok(bollard::container::LogOutput::StdOut { message })
120 | Ok(bollard::container::LogOutput::StdErr { message }) => {
121 Some(String::from_utf8_lossy(&message).to_string())
122 }
123 _ => None,
124 }
125 })
126 .fold(String::new(), |mut acc, item| async move {
127 acc.push_str(&item);
128 acc
129 })
130 .await;
131
132 return Ok(output);
133 }
134
135 Ok(String::new())
136 }
137
138 // Read a file from the container.
139 pub async fn read_file(&self, path: &str) -> Result<String, DockerContainerGuardError> {
140 let bytes = self
141 .docker
142 .download_from_container::<String>(
143 &self.id,
144 Some(bollard::container::DownloadFromContainerOptions { path: path.into() }),
145 )
146 .try_fold(BytesMut::new(), |mut bytes, b| async move {
147 bytes.put(b);
148 Ok(bytes)
149 })
150 .await?;
151
152 if bytes.len() == 0 {
153 return Err(DockerContainerGuardError::FileNotFound(path.into()));
154 }
155
156 for file in tar::Archive::new(&bytes[..]).entries()? {
157 let mut s = String::new();
158 file?.read_to_string(&mut s).unwrap();
159 return Ok(s);
160 }
161
162 Err(DockerContainerGuardError::FileNotFound(path.into()))
163 }
164
165 // Get a HashMap of all users in the container.
166 pub async fn get_users(
167 &self,
168 ) -> Result<HashMap<String, PasswdEntry>, DockerContainerGuardError> {
169 let output = self.read_file("/etc/passwd").await?;
170
171 return output
172 .lines()
173 .map(|line| {
174 PasswdEntry::from_line(line)
175 .map(|entry| (entry.name.clone(), entry))
176 .map_err(DockerContainerGuardError::from)
177 })
178 .collect();
179 }
180
181 // Get a specific user from the container.
182 pub async fn get_user(&self, name: &str) -> Result<PasswdEntry, DockerContainerGuardError> {
183 let users = self.get_users().await?;
184 let user = users
185 .get(name)
186 .ok_or_else(|| DockerContainerGuardError::UserNotFound(name.into()))?;
187
188 Ok(user.clone())
189 }
190}
191
192impl Drop for DockerContainerGuard {
193 fn drop(&mut self) {
194 let docker = self.docker.clone();
195 let container_id = self.id.clone();
196
197 tokio::spawn(async move {
198 docker
199 .remove_container(
200 &container_id,
201 Some(bollard::container::RemoveContainerOptions {
202 force: true,
203 ..Default::default()
204 }),
205 )
206 .await
207 });
208 }
209}
210
211#[cfg(test)]
212mod tests {
213 use super::*;
214
215 #[tokio::test]
216 async fn test_container_exec() -> Result<(), DockerContainerGuardError> {
Mohammed Naser6a5e7d72025-02-17 16:06:24 -0500217 let guard = DockerContainerGuard::spawn("registry.atmosphere.dev/docker.io/library/alpine:latest").await?;
Mohammed Naser9a6c35d2025-02-13 13:23:51 -0500218
219 let output = guard.exec(vec!["echo", "hello from container"]).await?;
220 assert!(output.contains("hello from container"));
221
222 Ok(())
223 }
224
225 #[tokio::test]
226 async fn test_container_read_file() -> Result<(), DockerContainerGuardError> {
Mohammed Naser6a5e7d72025-02-17 16:06:24 -0500227 let guard = DockerContainerGuard::spawn("registry.atmosphere.dev/docker.io/library/alpine:latest").await?;
Mohammed Naser9a6c35d2025-02-13 13:23:51 -0500228
229 let file = guard.read_file("/usr/lib/os-release").await?;
230 assert!(file.len() > 0);
231
232 Ok(())
233 }
234
235 #[tokio::test]
236 async fn test_container_get_user() -> Result<(), DockerContainerGuardError> {
Mohammed Naser6a5e7d72025-02-17 16:06:24 -0500237 let guard = DockerContainerGuard::spawn("registry.atmosphere.dev/docker.io/library/alpine:latest").await?;
Mohammed Naser9a6c35d2025-02-13 13:23:51 -0500238
239 let user = guard.get_user("root").await?;
240 assert_eq!(user.name, "root");
241
242 Ok(())
243 }
244}