From cf7864cd97340f5cafce0c465c8401946c867912 Mon Sep 17 00:00:00 2001
From: Ethan Mosbaugh <ethan@replicated.com>
Date: Fri, 23 Jul 2021 13:37:57 +0000
Subject: [PATCH] Copy collectors extractArchive property

---
 .../troubleshoot/v1beta2/collector_shared.go  | 14 +++--
 pkg/collect/copy.go                           | 60 ++++++++++++++++---
 pkg/collect/copy_from_host.go                 | 13 +++-
 3 files changed, 71 insertions(+), 16 deletions(-)

diff --git a/pkg/apis/troubleshoot/v1beta2/collector_shared.go b/pkg/apis/troubleshoot/v1beta2/collector_shared.go
index b2165cd6..9c2df30a 100644
--- a/pkg/apis/troubleshoot/v1beta2/collector_shared.go
+++ b/pkg/apis/troubleshoot/v1beta2/collector_shared.go
@@ -92,12 +92,13 @@ type Exec struct {
 }
 
 type Copy struct {
-	CollectorMeta `json:",inline" yaml:",inline"`
-	Name          string   `json:"name,omitempty" yaml:"name,omitempty"`
-	Selector      []string `json:"selector" yaml:"selector"`
-	Namespace     string   `json:"namespace" yaml:"namespace"`
-	ContainerPath string   `json:"containerPath" yaml:"containerPath"`
-	ContainerName string   `json:"containerName,omitempty" yaml:"containerName,omitempty"`
+	CollectorMeta  `json:",inline" yaml:",inline"`
+	Name           string   `json:"name,omitempty" yaml:"name,omitempty"`
+	Selector       []string `json:"selector" yaml:"selector"`
+	Namespace      string   `json:"namespace" yaml:"namespace"`
+	ContainerPath  string   `json:"containerPath" yaml:"containerPath"`
+	ContainerName  string   `json:"containerName,omitempty" yaml:"containerName,omitempty"`
+	ExtractArchive bool     `json:"extractArchive,omitempty" yaml:"extractArchive,omitempty"`
 }
 
 type CopyFromHost struct {
@@ -109,6 +110,7 @@ type CopyFromHost struct {
 	ImagePullSecret *ImagePullSecrets `json:"imagePullSecret,omitempty" yaml:"imagePullSecret,omitempty"`
 	Timeout         string            `json:"timeout,omitempty" yaml:"timeout,omitempty"`
 	HostPath        string            `json:"hostPath" yaml:"hostPath"`
+	ExtractArchive  bool              `json:"extractArchive,omitempty" yaml:"extractArchive,omitempty"`
 }
 
 type HTTP struct {
diff --git a/pkg/collect/copy.go b/pkg/collect/copy.go
index 1cfc9d2a..cafdc9c0 100644
--- a/pkg/collect/copy.go
+++ b/pkg/collect/copy.go
@@ -1,9 +1,12 @@
 package collect
 
 import (
+	"archive/tar"
 	"bytes"
 	"context"
 	"fmt"
+	"io"
+	"io/ioutil"
 	"path/filepath"
 
 	"github.com/pkg/errors"
@@ -64,23 +67,35 @@ func copyFiles(ctx context.Context, client *kubernetes.Clientset, c *Collector,
 		containerName = copyCollector.ContainerName
 	}
 
+	errs := map[string]string{}
+
 	stdout, stderr, err := getFilesFromPod(ctx, c.ClientConfig, client, pod.Name, containerName, pod.Namespace, copyCollector.ContainerPath)
 	if err != nil {
-		errors := map[string]string{
-			filepath.Join(copyCollector.ContainerPath, "error"): err.Error(),
-		}
+		errs[filepath.Join(copyCollector.ContainerPath, "error")] = err.Error()
 		if len(stdout) > 0 {
-			errors[filepath.Join(copyCollector.ContainerPath, "stdout")] = string(stdout)
+			errs[filepath.Join(copyCollector.ContainerPath, "stdout")] = string(stdout)
 		}
 		if len(stderr) > 0 {
-			errors[filepath.Join(copyCollector.ContainerPath, "stderr")] = string(stderr)
+			errs[filepath.Join(copyCollector.ContainerPath, "stderr")] = string(stderr)
 		}
-		return nil, errors
+		return nil, errs
 	}
 
-	return map[string][]byte{
-		filepath.Base(copyCollector.ContainerPath) + ".tar": stdout,
-	}, nil
+	runOutput := map[string][]byte{}
+
+	if copyCollector.ExtractArchive {
+		files, err := extractTar(bytes.NewReader(stdout))
+		if err != nil {
+			errs[filepath.Join(copyCollector.ContainerPath, "error")] = errors.Wrap(err, "extract tar").Error()
+		}
+		for name, data := range files {
+			runOutput[filepath.Join(filepath.Base(copyCollector.ContainerPath), name)] = data
+		}
+	} else {
+		runOutput[filepath.Base(copyCollector.ContainerPath)+".tar"] = stdout
+	}
+
+	return runOutput, errs
 }
 
 func getFilesFromPod(ctx context.Context, clientConfig *restclient.Config, client kubernetes.Interface, podName string, containerName string, namespace string, containerPath string) ([]byte, []byte, error) {
@@ -131,3 +146,30 @@ func getCopyErrosFileName(copyCollector *troubleshootv1beta2.Copy) string {
 	// TODO: random part
 	return "errors.json"
 }
+
+func extractTar(reader io.Reader) (map[string][]byte, error) {
+	files := map[string][]byte{}
+
+	tr := tar.NewReader(reader)
+	for {
+		header, err := tr.Next()
+		if err == io.EOF {
+			break
+		} else if err != nil {
+			return files, errors.Wrap(err, "read header")
+		}
+
+		switch header.Typeflag {
+		case tar.TypeReg:
+			data, err := ioutil.ReadAll(tr)
+			if err != nil {
+				return files, errors.Wrapf(err, "read file %s", header.Name)
+			}
+			files[header.Name] = data
+		default:
+			continue
+		}
+	}
+
+	return files, nil
+}
diff --git a/pkg/collect/copy_from_host.go b/pkg/collect/copy_from_host.go
index ae7c8913..4c74d460 100644
--- a/pkg/collect/copy_from_host.go
+++ b/pkg/collect/copy_from_host.go
@@ -1,6 +1,7 @@
 package collect
 
 import (
+	"bytes"
 	"context"
 	"path/filepath"
 	"time"
@@ -225,7 +226,17 @@ func copyFromHostGetFilesFromPods(ctx context.Context, clientConfig *restclient.
 				runOutput[filepath.Join(outputNodeFilename, "stderr.txt")] = stderr
 			}
 		} else {
-			runOutput[filepath.Join(outputNodeFilename, "archive.tar")] = stdout
+			if collector.ExtractArchive {
+				files, err := extractTar(bytes.NewReader(stdout))
+				if err != nil {
+					runOutput[filepath.Join(outputNodeFilename, "error.txt")] = []byte(errors.Wrap(err, "extract tar").Error())
+				}
+				for name, data := range files {
+					runOutput[filepath.Join(outputNodeFilename, name)] = data
+				}
+			} else {
+				runOutput[filepath.Join(outputNodeFilename, "archive.tar")] = stdout
+			}
 		}
 	}