From e5f2c2fbba4ede11cf366f1e7eed31cbf6cfbe4d Mon Sep 17 00:00:00 2001
From: Erdem Meydanli <meydanli@amazon.com>
Date: Thu, 21 Mar 2024 22:44:34 +0000
Subject: [PATCH] enclave_build: Extract stream output handling

Extract stream output handling code into handle_stream_output function
and move to utils.rs

Signed-off-by: Erdem Meydanli <meydanli@amazon.com>
---
 enclave_build/src/docker.rs | 56 ++++++-------------------------------
 enclave_build/src/lib.rs    |  1 +
 enclave_build/src/utils.rs  | 52 ++++++++++++++++++++++++++++++++++
 3 files changed, 62 insertions(+), 47 deletions(-)
 create mode 100644 enclave_build/src/utils.rs

diff --git a/enclave_build/src/docker.rs b/enclave_build/src/docker.rs
index 553587385..df3fbbfc8 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 12d9ccf52..66069b48b 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 000000000..3c78faad9
--- /dev/null
+++ b/enclave_build/src/utils.rs
@@ -0,0 +1,52 @@
+// 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<String>;
+}
+
+// Implement StreamItem for CreateImageInfo
+impl StreamItem for CreateImageInfo {
+    fn error(&self) -> Option<String> {
+        self.error.clone()
+    }
+}
+
+// Implement StreamItem for BuildInfo
+impl StreamItem for BuildInfo {
+    fn error(&self) -> Option<String> {
+        self.error.clone()
+    }
+}
+
+pub async fn handle_stream_output<T, U>(
+    mut stream: impl Stream<Item = Result<T, Error>> + 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(())
+}