diff --git a/enclave_build/src/docker.rs b/enclave_build/src/docker.rs index 55358738..df3fbbfc 100644 --- a/enclave_build/src/docker.rs +++ b/enclave_build/src/docker.rs @@ -2,14 +2,14 @@ // SPDX-License-Identifier: Apache-2.0 use crate::docker::DockerError::CredentialsError; +use crate::utils::handle_stream_output; use base64::{engine::general_purpose, Engine as _}; use bollard::auth::DockerCredentials; use bollard::image::{BuildImageOptions, CreateImageOptions}; use bollard::secret::ImageInspect; use bollard::Docker; use flate2::{write::GzEncoder, Compression}; -use futures::stream::StreamExt; -use log::{debug, error, info}; +use log::{debug, error}; use serde_json::{json, Value}; use std::fs::File; use std::io::Write; @@ -196,30 +196,11 @@ impl DockerUtil { } }; - let mut stream = - self.docker - .create_image(Some(create_image_options), None, credentials); - - loop { - if let Some(item) = stream.next().await { - match item { - Ok(output) => { - if let Some(err_msg) = &output.error { - error!("{:?}", err_msg); - break Err(DockerError::PullError); - } else { - info!("{:?}", output); - } - } - Err(e) => { - error!("{:?}", e); - break Err(DockerError::PullError); - } - } - } else { - break Ok(()); - } - } + let stream = self + .docker + .create_image(Some(create_image_options), None, credentials); + + handle_stream_output(stream, DockerError::PullError).await }) } @@ -244,7 +225,7 @@ impl DockerUtil { let runtime = Runtime::new().map_err(|_| DockerError::RuntimeError)?; runtime.block_on(async move { - let mut stream = self.docker.build_image( + let stream = self.docker.build_image( BuildImageOptions { dockerfile: "Dockerfile".to_string(), t: self.docker_image.clone(), @@ -254,26 +235,7 @@ impl DockerUtil { Some(Self::build_tarball(dockerfile_dir)?.into()), ); - loop { - if let Some(item) = stream.next().await { - match item { - Ok(output) => { - if let Some(err_msg) = &output.error { - error!("{:?}", err_msg.clone()); - break Err(DockerError::BuildError); - } else { - info!("{:?}", output); - } - } - Err(e) => { - error!("{:?}", e); - break Err(DockerError::BuildError); - } - } - } else { - break Ok(()); - } - } + handle_stream_output(stream, DockerError::BuildError).await }) } diff --git a/enclave_build/src/lib.rs b/enclave_build/src/lib.rs index 12d9ccf5..66069b48 100644 --- a/enclave_build/src/lib.rs +++ b/enclave_build/src/lib.rs @@ -7,6 +7,7 @@ use std::path::Path; use std::process::Command; mod docker; +mod utils; mod yaml_generator; use aws_nitro_enclaves_image_format::defs::{EifBuildInfo, EifIdentityInfo, EIF_HDR_ARCH_ARM64}; diff --git a/enclave_build/src/utils.rs b/enclave_build/src/utils.rs new file mode 100644 index 00000000..530ab7bb --- /dev/null +++ b/enclave_build/src/utils.rs @@ -0,0 +1,53 @@ +// Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 + +use bollard::errors::Error; +use bollard::secret::{BuildInfo, CreateImageInfo}; +use futures::stream::StreamExt; +use futures::Stream; +use log::{error, info}; + +pub trait StreamItem { + fn error(&self) -> Option; +} + +// Implement StreamItem for CreateImageInfo +impl StreamItem for CreateImageInfo { + fn error(&self) -> Option { + self.error.clone() + } +} + +// Implement StreamItem for BuildInfo +impl StreamItem for BuildInfo { + fn error(&self) -> Option { + self.error.clone() + } +} + +pub async fn handle_stream_output( + mut stream: impl Stream> + Unpin, + error_type: U, +) -> Result<(), U> +where + T: StreamItem + std::fmt::Debug, +{ + while let Some(item) = stream.next().await { + match item { + Ok(output) => { + if let Some(err_msg) = output.error() { + error!("{:?}", err_msg); + return Err(error_type); + } else { + info!("{:?}", output); + } + } + Err(e) => { + error!("{:?}", e); + return Err(error_type); + } + } + } + + Ok(()) +}