From a81dda1f9869ad9c330c686dc1805276ac61a770 Mon Sep 17 00:00:00 2001
From: Cezar Rata <ceradev@amazon.com>
Date: Wed, 11 Sep 2024 18:19:47 +0000
Subject: [PATCH] fix: image load should close stream after copy

Signed-off-by: Cezar Rata <ceradev@amazon.com>
---
 internal/service/image/load.go      | 9 +++++++--
 internal/service/image/load_test.go | 1 +
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/internal/service/image/load.go b/internal/service/image/load.go
index febd5a4..7c7ea49 100644
--- a/internal/service/image/load.go
+++ b/internal/service/image/load.go
@@ -42,11 +42,16 @@ func (s *service) Load(ctx context.Context, inStream io.Reader, outStream io.Wri
 		return err
 	}
 	defer func() {
-		rw.Close()
 		os.Remove(img)
 	}()
 	go func() {
-		io.Copy(rw, inStream)
+		written, err := io.Copy(rw, inStream)
+		if err != nil {
+			s.logger.Errorf("failed to copy: %s", err)
+		} else {
+			s.logger.Debugf("copied %d bytes", written)
+		}
+		rw.Close()
 	}()
 	if err = s.nctlImageSvc.LoadImage(ctx, img, outStream, quiet); err != nil {
 		s.logger.Errorf("failed to load image %s: %s", img, err)
diff --git a/internal/service/image/load_test.go b/internal/service/image/load_test.go
index f6fc394..1fd331f 100644
--- a/internal/service/image/load_test.go
+++ b/internal/service/image/load_test.go
@@ -42,6 +42,7 @@ var _ = Describe("Image Load API", func() {
 	})
 	Context("service", func() {
 		It("should return no errors upon success", func() {
+			logger.EXPECT().Debugf(gomock.Any(), gomock.Any())
 			ncClient.EXPECT().GetDataStore().
 				Return(name, nil)
 			ncClient.EXPECT().LoadImage(gomock.Any(), gomock.Any(), nil, gomock.Any()).