Skip to main content

testcontainers/core/
copy.rs

1use std::path::{Path, PathBuf};
2
3use async_trait::async_trait;
4use tokio::io::{AsyncRead, AsyncReadExt, AsyncWriteExt};
5use tokio_tar::EntryType;
6
7#[derive(Debug, Clone)]
8pub struct CopyToContainerCollection(Vec<CopyToContainer>);
9
10#[derive(Debug, Clone)]
11pub struct CopyToContainer {
12    target: CopyTargetOptions,
13    source: CopyDataSource,
14}
15
16#[derive(Debug, Clone)]
17pub struct CopyTargetOptions {
18    target: String,
19    mode: Option<u32>,
20}
21
22#[derive(Debug, Clone)]
23pub enum CopyDataSource {
24    File(PathBuf),
25    Data(Vec<u8>),
26}
27
28/// Errors that can occur while materializing data copied from a container.
29#[derive(Debug, thiserror::Error)]
30pub enum CopyFromContainerError {
31    #[error("io failed with error: {0}")]
32    Io(#[from] std::io::Error),
33    #[error("archive did not contain any regular files")]
34    EmptyArchive,
35    #[error("requested container path is a directory")]
36    IsDirectory,
37    #[error("archive entry type '{0:?}' is not supported for requested target")]
38    UnsupportedEntry(EntryType),
39}
40
41/// Abstraction for materializing the bytes read from a source into a concrete destination.
42///
43/// Implementors typically persist the incoming bytes to disk or buffer them in memory. Some return
44/// a value that callers can work with (for example, the collected bytes), while others simply
45/// report success with `()`. Implementations must consume the provided reader until EOF or return
46/// an error. Destinations are allowed to discard any existing data to make room for the incoming
47/// bytes.
48#[async_trait(?Send)]
49pub trait CopyFileFromContainer {
50    type Output;
51
52    /// Writes all bytes from the reader into `self`, returning a value that represents the completed operation (or `()` for sinks that only confirm success).
53    ///
54    /// Implementations may mutate `self` and must propagate I/O errors via [`CopyFromContainerError`].
55    async fn copy_from_reader<R>(self, reader: R) -> Result<Self::Output, CopyFromContainerError>
56    where
57        R: AsyncRead + Unpin;
58}
59
60#[async_trait(?Send)]
61impl CopyFileFromContainer for Vec<u8> {
62    type Output = Vec<u8>;
63
64    async fn copy_from_reader<R>(
65        mut self,
66        reader: R,
67    ) -> Result<Self::Output, CopyFromContainerError>
68    where
69        R: AsyncRead + Unpin,
70    {
71        let mut_ref = &mut self;
72        mut_ref.copy_from_reader(reader).await?;
73        Ok(self)
74    }
75}
76
77#[async_trait(?Send)]
78impl CopyFileFromContainer for &mut Vec<u8> {
79    type Output = ();
80
81    async fn copy_from_reader<R>(
82        mut self,
83        mut reader: R,
84    ) -> Result<Self::Output, CopyFromContainerError>
85    where
86        R: AsyncRead + Unpin,
87    {
88        self.clear();
89        reader
90            .read_to_end(&mut self)
91            .await
92            .map_err(CopyFromContainerError::Io)?;
93        Ok(())
94    }
95}
96
97#[async_trait(?Send)]
98impl CopyFileFromContainer for PathBuf {
99    type Output = ();
100
101    async fn copy_from_reader<R>(self, reader: R) -> Result<Self::Output, CopyFromContainerError>
102    where
103        R: AsyncRead + Unpin,
104    {
105        self.as_path().copy_from_reader(reader).await
106    }
107}
108
109#[async_trait(?Send)]
110impl CopyFileFromContainer for &Path {
111    type Output = ();
112
113    async fn copy_from_reader<R>(
114        self,
115        mut reader: R,
116    ) -> Result<Self::Output, CopyFromContainerError>
117    where
118        R: AsyncRead + Unpin,
119    {
120        if let Some(parent) = self.parent() {
121            if !parent.as_os_str().is_empty() {
122                tokio::fs::create_dir_all(parent)
123                    .await
124                    .map_err(CopyFromContainerError::Io)?;
125            }
126        }
127
128        let mut file = tokio::fs::File::create(self)
129            .await
130            .map_err(CopyFromContainerError::Io)?;
131
132        tokio::io::copy(&mut reader, &mut file)
133            .await
134            .map_err(CopyFromContainerError::Io)?;
135
136        file.flush().await.map_err(CopyFromContainerError::Io)?;
137        Ok(())
138    }
139}
140
141#[derive(Debug, thiserror::Error)]
142pub enum CopyToContainerError {
143    #[error("io failed with error: {0}")]
144    IoError(std::io::Error),
145    #[error("failed to get the path name: {0}")]
146    PathNameError(String),
147}
148
149impl CopyToContainerCollection {
150    pub fn new(collection: Vec<CopyToContainer>) -> Self {
151        Self(collection)
152    }
153
154    pub fn add(&mut self, entry: CopyToContainer) {
155        self.0.push(entry);
156    }
157
158    pub(crate) async fn tar(&self) -> Result<bytes::Bytes, CopyToContainerError> {
159        let mut ar = tokio_tar::Builder::new(Vec::new());
160
161        for copy_to_container in &self.0 {
162            copy_to_container.append_tar(&mut ar).await?
163        }
164
165        let bytes = ar
166            .into_inner()
167            .await
168            .map_err(CopyToContainerError::IoError)?;
169
170        Ok(bytes::Bytes::copy_from_slice(bytes.as_slice()))
171    }
172}
173
174impl CopyToContainer {
175    pub fn new(source: impl Into<CopyDataSource>, target: impl Into<CopyTargetOptions>) -> Self {
176        Self {
177            source: source.into(),
178            target: target.into(),
179        }
180    }
181
182    pub(crate) async fn tar(&self) -> Result<bytes::Bytes, CopyToContainerError> {
183        let mut ar = tokio_tar::Builder::new(Vec::new());
184
185        self.append_tar(&mut ar).await?;
186
187        let bytes = ar
188            .into_inner()
189            .await
190            .map_err(CopyToContainerError::IoError)?;
191
192        Ok(bytes::Bytes::copy_from_slice(bytes.as_slice()))
193    }
194
195    pub(crate) async fn append_tar(
196        &self,
197        ar: &mut tokio_tar::Builder<Vec<u8>>,
198    ) -> Result<(), CopyToContainerError> {
199        self.source.append_tar(ar, &self.target).await
200    }
201}
202
203impl CopyTargetOptions {
204    pub fn new(target: impl Into<String>) -> Self {
205        Self {
206            target: target.into(),
207            mode: None,
208        }
209    }
210
211    pub fn with_mode(mut self, mode: u32) -> Self {
212        self.mode = Some(mode);
213        self
214    }
215
216    pub fn target(&self) -> &str {
217        &self.target
218    }
219
220    pub fn mode(&self) -> Option<u32> {
221        self.mode
222    }
223}
224
225impl<T> From<T> for CopyTargetOptions
226where
227    T: Into<String>,
228{
229    fn from(value: T) -> Self {
230        CopyTargetOptions::new(value.into())
231    }
232}
233
234impl From<&Path> for CopyDataSource {
235    fn from(value: &Path) -> Self {
236        CopyDataSource::File(value.to_path_buf())
237    }
238}
239
240impl From<PathBuf> for CopyDataSource {
241    fn from(value: PathBuf) -> Self {
242        CopyDataSource::File(value)
243    }
244}
245impl From<Vec<u8>> for CopyDataSource {
246    fn from(value: Vec<u8>) -> Self {
247        CopyDataSource::Data(value)
248    }
249}
250
251impl CopyDataSource {
252    pub(crate) async fn append_tar(
253        &self,
254        ar: &mut tokio_tar::Builder<Vec<u8>>,
255        target: &CopyTargetOptions,
256    ) -> Result<(), CopyToContainerError> {
257        let target_path = target.target();
258
259        match self {
260            CopyDataSource::File(source_file_path) => {
261                if let Err(e) = append_tar_file(ar, source_file_path, target).await {
262                    log::error!(
263                        "Could not append file/dir to tar: {source_file_path:?}:{target_path}"
264                    );
265                    return Err(e);
266                }
267            }
268            CopyDataSource::Data(data) => {
269                if let Err(e) = append_tar_bytes(ar, data, target).await {
270                    log::error!("Could not append data to tar: {target_path}");
271                    return Err(e);
272                }
273            }
274        };
275
276        Ok(())
277    }
278}
279
280async fn append_tar_file(
281    ar: &mut tokio_tar::Builder<Vec<u8>>,
282    source_file_path: &Path,
283    target: &CopyTargetOptions,
284) -> Result<(), CopyToContainerError> {
285    let target_path = make_path_relative(target.target());
286    let meta = tokio::fs::metadata(source_file_path)
287        .await
288        .map_err(CopyToContainerError::IoError)?;
289
290    if meta.is_dir() {
291        ar.append_dir_all(target_path, source_file_path)
292            .await
293            .map_err(CopyToContainerError::IoError)?;
294    } else {
295        let f = &mut tokio::fs::File::open(source_file_path)
296            .await
297            .map_err(CopyToContainerError::IoError)?;
298
299        let mut header = tokio_tar::Header::new_gnu();
300        header.set_size(meta.len());
301
302        #[cfg(unix)]
303        {
304            use std::os::unix::fs::PermissionsExt;
305            let mode = target.mode().unwrap_or_else(|| meta.permissions().mode());
306            header.set_mode(mode);
307        }
308
309        #[cfg(not(unix))]
310        {
311            let mode = target.mode().unwrap_or(0o644);
312            header.set_mode(mode);
313        }
314
315        header.set_cksum();
316
317        ar.append_data(&mut header, target_path, f)
318            .await
319            .map_err(CopyToContainerError::IoError)?;
320    };
321
322    Ok(())
323}
324
325async fn append_tar_bytes(
326    ar: &mut tokio_tar::Builder<Vec<u8>>,
327    data: &Vec<u8>,
328    target: &CopyTargetOptions,
329) -> Result<(), CopyToContainerError> {
330    let relative_target_path = make_path_relative(target.target());
331
332    let mut header = tokio_tar::Header::new_gnu();
333    header.set_size(data.len() as u64);
334    header.set_mode(target.mode().unwrap_or(0o0644));
335    header.set_cksum();
336
337    ar.append_data(&mut header, relative_target_path, data.as_slice())
338        .await
339        .map_err(CopyToContainerError::IoError)?;
340
341    Ok(())
342}
343
344fn make_path_relative(path: &str) -> String {
345    // TODO support also absolute windows paths like "C:\temp\foo.txt"
346    if path.starts_with("/") {
347        path.trim_start_matches("/").to_string()
348    } else {
349        path.to_string()
350    }
351}
352
353#[cfg(test)]
354mod tests {
355    use std::{fs::File, io::Write};
356
357    use futures::StreamExt;
358    use tempfile::tempdir;
359    use tokio_tar::Archive;
360
361    use super::*;
362
363    #[tokio::test]
364    async fn copytocontainer_tar_file_success() {
365        let temp_dir = tempdir().unwrap();
366        let file_path = temp_dir.path().join("file.txt");
367        let mut file = File::create(&file_path).unwrap();
368        writeln!(file, "TEST").unwrap();
369
370        let copy_to_container = CopyToContainer::new(file_path, "file.txt");
371        let result = copy_to_container.tar().await;
372
373        assert!(result.is_ok());
374        let bytes = result.unwrap();
375        assert!(!bytes.is_empty());
376    }
377
378    #[tokio::test]
379    async fn copytocontainer_tar_data_success() {
380        let data = vec![1, 2, 3, 4, 5];
381        let copy_to_container = CopyToContainer::new(data, "data.bin");
382        let result = copy_to_container.tar().await;
383
384        assert!(result.is_ok());
385        let bytes = result.unwrap();
386        assert!(!bytes.is_empty());
387    }
388
389    #[tokio::test]
390    async fn copytocontainer_tar_file_not_found() {
391        let temp_dir = tempdir().unwrap();
392        let non_existent_file_path = temp_dir.path().join("non_existent_file.txt");
393
394        let copy_to_container = CopyToContainer::new(non_existent_file_path, "file.txt");
395        let result = copy_to_container.tar().await;
396
397        assert!(result.is_err());
398        if let Err(CopyToContainerError::IoError(err)) = result {
399            assert_eq!(err.kind(), std::io::ErrorKind::NotFound);
400        } else {
401            panic!("Expected IoError");
402        }
403    }
404
405    #[tokio::test]
406    async fn copytocontainercollection_tar_file_and_data() {
407        let temp_dir = tempdir().unwrap();
408        let file_path = temp_dir.path().join("file.txt");
409        let mut file = File::create(&file_path).unwrap();
410        writeln!(file, "TEST").unwrap();
411
412        let copy_to_container_collection = CopyToContainerCollection::new(vec![
413            CopyToContainer::new(file_path, "file.txt"),
414            CopyToContainer::new(vec![1, 2, 3, 4, 5], "data.bin"),
415        ]);
416
417        let result = copy_to_container_collection.tar().await;
418
419        assert!(result.is_ok());
420        let bytes = result.unwrap();
421        assert!(!bytes.is_empty());
422    }
423
424    #[tokio::test]
425    async fn tar_bytes_respects_custom_mode() {
426        let data = vec![1, 2, 3];
427        let target = CopyTargetOptions::new("data.bin").with_mode(0o600);
428        let copy_to_container = CopyToContainer::new(data, target);
429
430        let tar_bytes = copy_to_container.tar().await.unwrap();
431        let mut archive = Archive::new(std::io::Cursor::new(tar_bytes));
432        let mut entries = archive.entries().unwrap();
433        let entry = entries.next().await.unwrap().unwrap();
434
435        assert_eq!(entry.header().mode().unwrap(), 0o600);
436    }
437
438    #[tokio::test]
439    async fn tar_file_respects_custom_mode() {
440        let temp_dir = tempdir().unwrap();
441        let file_path = temp_dir.path().join("file.txt");
442        let mut file = File::create(&file_path).unwrap();
443        writeln!(file, "TEST").unwrap();
444
445        let target = CopyTargetOptions::new("file.txt").with_mode(0o640);
446        let copy_to_container = CopyToContainer::new(file_path, target);
447
448        let tar_bytes = copy_to_container.tar().await.unwrap();
449        let mut archive = Archive::new(std::io::Cursor::new(tar_bytes));
450        let mut entries = archive.entries().unwrap();
451        let entry = entries.next().await.unwrap().unwrap();
452
453        assert_eq!(entry.header().mode().unwrap(), 0o640);
454    }
455}