Mohammed Naser | 9a6c35d | 2025-02-13 13:23:51 -0500 | [diff] [blame] | 1 | extern crate tar; |
| 2 | |
| 3 | use bollard::container::{Config, CreateContainerOptions, StartContainerOptions}; |
| 4 | use bollard::exec::{CreateExecOptions, StartExecResults}; |
| 5 | use bollard::Docker; |
| 6 | use bytes::{BufMut, BytesMut}; |
| 7 | use futures_util::stream::StreamExt; |
| 8 | use futures_util::TryStreamExt; |
| 9 | use passwd::PasswdEntry; |
| 10 | use rand::Rng; |
| 11 | use std::collections::HashMap; |
| 12 | use std::io::Read; |
| 13 | use thiserror::Error; |
| 14 | |
| 15 | #[derive(Debug, Error)] |
| 16 | pub 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)] |
| 37 | pub struct DockerContainerGuard { |
| 38 | pub id: String, |
| 39 | pub image: String, |
| 40 | docker: Docker, |
| 41 | } |
| 42 | |
| 43 | impl 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 | |
| 192 | impl 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)] |
| 212 | mod tests { |
| 213 | use super::*; |
| 214 | |
| 215 | #[tokio::test] |
| 216 | async fn test_container_exec() -> Result<(), DockerContainerGuardError> { |
Mohammed Naser | 6a5e7d7 | 2025-02-17 16:06:24 -0500 | [diff] [blame] | 217 | let guard = DockerContainerGuard::spawn("registry.atmosphere.dev/docker.io/library/alpine:latest").await?; |
Mohammed Naser | 9a6c35d | 2025-02-13 13:23:51 -0500 | [diff] [blame] | 218 | |
| 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 Naser | 6a5e7d7 | 2025-02-17 16:06:24 -0500 | [diff] [blame] | 227 | let guard = DockerContainerGuard::spawn("registry.atmosphere.dev/docker.io/library/alpine:latest").await?; |
Mohammed Naser | 9a6c35d | 2025-02-13 13:23:51 -0500 | [diff] [blame] | 228 | |
| 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 Naser | 6a5e7d7 | 2025-02-17 16:06:24 -0500 | [diff] [blame] | 237 | let guard = DockerContainerGuard::spawn("registry.atmosphere.dev/docker.io/library/alpine:latest").await?; |
Mohammed Naser | 9a6c35d | 2025-02-13 13:23:51 -0500 | [diff] [blame] | 238 | |
| 239 | let user = guard.get_user("root").await?; |
| 240 | assert_eq!(user.name, "root"); |
| 241 | |
| 242 | Ok(()) |
| 243 | } |
| 244 | } |