cargo_test_support/
containers.rs

1//! Support for testing using Docker containers.
2//!
3//! The [`Container`] type is a builder for configuring a container to run.
4//! After you call `launch`, you can use the [`ContainerHandle`] to interact
5//! with the running container.
6//!
7//! Tests using containers must use `#[cargo_test(container_test)]` to disable
8//! them unless the `CARGO_CONTAINER_TESTS` environment variable is set.
9
10use cargo_util::ProcessBuilder;
11use std::collections::HashMap;
12use std::io::Read;
13use std::path::PathBuf;
14use std::process::Command;
15use std::sync::atomic::{AtomicUsize, Ordering};
16use std::sync::Mutex;
17use tar::Header;
18
19/// A builder for configuring a container to run.
20pub struct Container {
21    /// The host directory that forms the basis of the Docker image.
22    build_context: PathBuf,
23    /// Files to copy over to the image.
24    files: Vec<MkFile>,
25}
26
27/// A handle to a running container.
28///
29/// You can use this to interact with the container.
30pub struct ContainerHandle {
31    /// The name of the container.
32    name: String,
33    /// The IP address of the container.
34    ///
35    /// NOTE: This is currently unused, but may be useful so I left it in.
36    /// This can only be used on Linux. macOS and Windows docker doesn't allow
37    /// direct connection to the container.
38    pub ip_address: String,
39    /// Port mappings of `container_port` to `host_port` for ports exposed via EXPOSE.
40    pub port_mappings: HashMap<u16, u16>,
41}
42
43impl Container {
44    pub fn new(context_dir: &str) -> Container {
45        assert!(std::env::var_os("CARGO_CONTAINER_TESTS").is_some());
46        let mut build_context = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
47        build_context.push("containers");
48        build_context.push(context_dir);
49        Container {
50            build_context,
51            files: Vec::new(),
52        }
53    }
54
55    /// Adds a file to be copied into the container.
56    pub fn file(mut self, file: MkFile) -> Self {
57        self.files.push(file);
58        self
59    }
60
61    /// Starts the container.
62    pub fn launch(mut self) -> ContainerHandle {
63        static NEXT_ID: AtomicUsize = AtomicUsize::new(0);
64
65        let id = NEXT_ID.fetch_add(1, Ordering::SeqCst);
66        let name = format!("cargo_test_{id}");
67        remove_if_exists(&name);
68        self.create_container(&name);
69        self.copy_files(&name);
70        self.start_container(&name);
71        let info = self.container_inspect(&name);
72        let ip_address = if cfg!(target_os = "linux") {
73            info[0]["NetworkSettings"]["IPAddress"]
74                .as_str()
75                .unwrap()
76                .to_string()
77        } else {
78            // macOS and Windows can't make direct connections to the
79            // container. It only works through exposed ports or mapped ports.
80            "127.0.0.1".to_string()
81        };
82        let port_mappings = self.port_mappings(&info);
83        self.wait_till_ready(&port_mappings);
84
85        ContainerHandle {
86            name,
87            ip_address,
88            port_mappings,
89        }
90    }
91
92    fn create_container(&self, name: &str) {
93        static BUILD_LOCK: Mutex<()> = Mutex::new(());
94
95        let image_base = self.build_context.file_name().unwrap();
96        let image_name = format!("cargo-test-{}", image_base.to_str().unwrap());
97        let _lock = BUILD_LOCK
98            .lock()
99            .map_err(|_| panic!("previous docker build failed, unable to run test"));
100        ProcessBuilder::new("docker")
101            .args(&["build", "--tag", image_name.as_str()])
102            .arg(&self.build_context)
103            .exec_with_output()
104            .unwrap();
105
106        ProcessBuilder::new("docker")
107            .args(&[
108                "container",
109                "create",
110                "--publish-all",
111                "--rm",
112                "--name",
113                name,
114            ])
115            .arg(image_name)
116            .exec_with_output()
117            .unwrap();
118    }
119
120    fn copy_files(&mut self, name: &str) {
121        if self.files.is_empty() {
122            return;
123        }
124        let mut ar = tar::Builder::new(Vec::new());
125        ar.sparse(false);
126        let files = std::mem::replace(&mut self.files, Vec::new());
127        for mut file in files {
128            ar.append_data(&mut file.header, &file.path, file.contents.as_slice())
129                .unwrap();
130        }
131        let ar = ar.into_inner().unwrap();
132        ProcessBuilder::new("docker")
133            .args(&["cp", "-"])
134            .arg(format!("{name}:/"))
135            .stdin(ar)
136            .exec_with_output()
137            .unwrap();
138    }
139
140    fn start_container(&self, name: &str) {
141        ProcessBuilder::new("docker")
142            .args(&["container", "start"])
143            .arg(name)
144            .exec_with_output()
145            .unwrap();
146    }
147
148    fn container_inspect(&self, name: &str) -> serde_json::Value {
149        let output = ProcessBuilder::new("docker")
150            .args(&["inspect", name])
151            .exec_with_output()
152            .unwrap();
153        serde_json::from_slice(&output.stdout).unwrap()
154    }
155
156    /// Returns the mapping of container_port->host_port for ports that were
157    /// exposed with EXPOSE.
158    fn port_mappings(&self, info: &serde_json::Value) -> HashMap<u16, u16> {
159        info[0]["NetworkSettings"]["Ports"]
160            .as_object()
161            .unwrap()
162            .iter()
163            .map(|(key, value)| {
164                let key = key
165                    .strip_suffix("/tcp")
166                    .expect("expected TCP only ports")
167                    .parse()
168                    .unwrap();
169                let values = value.as_array().unwrap();
170                let value = values
171                    .iter()
172                    .find(|value| value["HostIp"].as_str().unwrap() == "0.0.0.0")
173                    .expect("expected localhost IP");
174                let host_port = value["HostPort"].as_str().unwrap().parse().unwrap();
175                (key, host_port)
176            })
177            .collect()
178    }
179
180    fn wait_till_ready(&self, port_mappings: &HashMap<u16, u16>) {
181        for port in port_mappings.values() {
182            let mut ok = false;
183            for _ in 0..30 {
184                match std::net::TcpStream::connect(format!("127.0.0.1:{port}")) {
185                    Ok(_) => {
186                        ok = true;
187                        break;
188                    }
189                    Err(e) => {
190                        if e.kind() != std::io::ErrorKind::ConnectionRefused {
191                            panic!("unexpected localhost connection error: {e:?}");
192                        }
193                        std::thread::sleep(std::time::Duration::new(1, 0));
194                    }
195                }
196            }
197            if !ok {
198                panic!("no listener on localhost port {port}");
199            }
200        }
201    }
202}
203
204impl ContainerHandle {
205    /// Executes a program inside a running container.
206    pub fn exec(&self, args: &[&str]) -> std::process::Output {
207        ProcessBuilder::new("docker")
208            .args(&["container", "exec", &self.name])
209            .args(args)
210            .exec_with_output()
211            .unwrap()
212    }
213
214    /// Returns the contents of a file inside the container.
215    pub fn read_file(&self, path: &str) -> String {
216        let output = ProcessBuilder::new("docker")
217            .args(&["cp", &format!("{}:{}", self.name, path), "-"])
218            .exec_with_output()
219            .unwrap();
220        let mut ar = tar::Archive::new(output.stdout.as_slice());
221        let mut entry = ar.entries().unwrap().next().unwrap().unwrap();
222        let mut contents = String::new();
223        entry.read_to_string(&mut contents).unwrap();
224        contents
225    }
226}
227
228impl Drop for ContainerHandle {
229    fn drop(&mut self) {
230        // To help with debugging, this will keep the container alive.
231        if std::env::var_os("CARGO_CONTAINER_TEST_KEEP").is_some() {
232            return;
233        }
234        remove_if_exists(&self.name);
235    }
236}
237
238fn remove_if_exists(name: &str) {
239    if let Err(e) = Command::new("docker")
240        .args(&["container", "rm", "--force", name])
241        .output()
242    {
243        panic!("failed to run docker: {e}");
244    }
245}
246
247/// Builder for configuring a file to copy into a container.
248pub struct MkFile {
249    path: String,
250    contents: Vec<u8>,
251    header: Header,
252}
253
254impl MkFile {
255    /// Defines a file to add to the container.
256    ///
257    /// This should be passed to `Container::file`.
258    ///
259    /// The path is the path inside the container to create the file.
260    pub fn path(path: &str) -> MkFile {
261        MkFile {
262            path: path.to_string(),
263            contents: Vec::new(),
264            header: Header::new_gnu(),
265        }
266    }
267
268    pub fn contents(mut self, contents: impl Into<Vec<u8>>) -> Self {
269        self.contents = contents.into();
270        self.header.set_size(self.contents.len() as u64);
271        self
272    }
273
274    pub fn mode(mut self, mode: u32) -> Self {
275        self.header.set_mode(mode);
276        self
277    }
278
279    pub fn uid(mut self, uid: u64) -> Self {
280        self.header.set_uid(uid);
281        self
282    }
283
284    pub fn gid(mut self, gid: u64) -> Self {
285        self.header.set_gid(gid);
286        self
287    }
288}