From ca51e92878a22b5284cf714e4e5c98e34c2678c3 Mon Sep 17 00:00:00 2001
From: divolgin <dmitriy@replicated.com>
Date: Thu, 30 Sep 2021 18:24:52 +0000
Subject: [PATCH] Allow memory writers

---
 pkg/collect/copy.go           |  2 ++
 pkg/collect/copy_from_host.go |  2 ++
 pkg/collect/logs.go           |  4 ++--
 pkg/collect/result.go         | 24 +++++++++++++++++++++---
 4 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/pkg/collect/copy.go b/pkg/collect/copy.go
index 8ebdd3e8..a3e43715 100644
--- a/pkg/collect/copy.go
+++ b/pkg/collect/copy.go
@@ -138,6 +138,8 @@ func copyFilesFromPod(ctx context.Context, dstPath string, clientConfig *restcli
 		if err != nil {
 			return nil, nil, errors.Wrap(err, "failed to craete dest file")
 		}
+		defer result.CloseWriter(dstPath, filepath.Base(containerPath)+".tar", w)
+
 		stdoutWriter = w
 	}
 
diff --git a/pkg/collect/copy_from_host.go b/pkg/collect/copy_from_host.go
index 806384e6..f40230a8 100644
--- a/pkg/collect/copy_from_host.go
+++ b/pkg/collect/copy_from_host.go
@@ -315,6 +315,8 @@ func copyFilesFromHost(ctx context.Context, dstPath string, clientConfig *restcl
 		if err != nil {
 			return nil, nil, errors.Wrap(err, "failed to craete dest file")
 		}
+		defer result.CloseWriter(dstPath, "archive.tar", w)
+
 		stdoutWriter = w
 	}
 
diff --git a/pkg/collect/logs.go b/pkg/collect/logs.go
index b25412dc..8e33d6b6 100644
--- a/pkg/collect/logs.go
+++ b/pkg/collect/logs.go
@@ -125,7 +125,7 @@ func savePodLogs(ctx context.Context, bundlePath string, client *kubernetes.Clie
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to get log writer")
 	}
-	defer logWriter.Close()
+	defer result.CloseWriter(bundlePath, fileKey+".log", logWriter)
 
 	_, err = io.Copy(logWriter, podLogs)
 	if err != nil {
@@ -145,7 +145,7 @@ func savePodLogs(ctx context.Context, bundlePath string, client *kubernetes.Clie
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to get previous log writer")
 	}
-	defer logWriter.Close()
+	defer result.CloseWriter(bundlePath, fileKey+"-previous.log", logWriter)
 
 	_, err = io.Copy(prevLogWriter, podLogs)
 	if err != nil {
diff --git a/pkg/collect/result.go b/pkg/collect/result.go
index 01e537b5..890e12c1 100644
--- a/pkg/collect/result.go
+++ b/pkg/collect/result.go
@@ -89,6 +89,10 @@ func (r CollectorResult) GetReader(bundlePath string, relativePath string) (io.R
 		return bytes.NewReader(r[relativePath]), nil
 	}
 
+	if bundlePath == "" {
+		return nil, errors.New("cannot create reader, bundle path is empty")
+	}
+
 	filename := filepath.Join(bundlePath, relativePath)
 	f, err := os.Open(filename)
 	if err != nil {
@@ -98,9 +102,10 @@ func (r CollectorResult) GetReader(bundlePath string, relativePath string) (io.R
 	return f, nil
 }
 
-func (r CollectorResult) GetWriter(bundlePath string, relativePath string) (io.WriteCloser, error) {
-	if r[relativePath] != nil {
-		return nil, errors.New("writing to memory is not implemented")
+func (r CollectorResult) GetWriter(bundlePath string, relativePath string) (io.Writer, error) {
+	if bundlePath == "" {
+		var b bytes.Buffer
+		return &b, nil
 	}
 
 	fileDir, _ := filepath.Split(relativePath)
@@ -120,6 +125,19 @@ func (r CollectorResult) GetWriter(bundlePath string, relativePath string) (io.W
 	return f, nil
 }
 
+func (r CollectorResult) CloseWriter(bundlePath string, relativePath string, writer interface{}) error {
+	if c, ok := writer.(io.Closer); ok {
+		return errors.Wrap(c.Close(), "failed to close writer")
+	}
+
+	if b, ok := writer.(*bytes.Buffer); ok {
+		r[relativePath] = b.Bytes()
+		return nil
+	}
+
+	return errors.Errorf("cannot close writer of type %T", writer)
+}
+
 func TarSupportBundleDir(bundlePath string, input CollectorResult, outputFilename string) error {
 	fileWriter, err := os.Create(outputFilename)
 	if err != nil {