From 99003849cacb70957f307215f075b1590876bf9e Mon Sep 17 00:00:00 2001
From: York Chen <ychen@d2iq.com>
Date: Fri, 4 Aug 2023 14:46:09 -0400
Subject: [PATCH] feat: handle config file

new optional field config in hostRun spec
---
 .../v1beta2/hostcollector_shared.go           |  7 ++--
 pkg/collect/host_run.go                       | 40 +++++++++++++++----
 2 files changed, 36 insertions(+), 11 deletions(-)

diff --git a/pkg/apis/troubleshoot/v1beta2/hostcollector_shared.go b/pkg/apis/troubleshoot/v1beta2/hostcollector_shared.go
index 317d0b3e0..2217989b9 100644
--- a/pkg/apis/troubleshoot/v1beta2/hostcollector_shared.go
+++ b/pkg/apis/troubleshoot/v1beta2/hostcollector_shared.go
@@ -180,9 +180,10 @@ type HostServices struct {
 
 type HostRun struct {
 	HostCollectorMeta `json:",inline" yaml:",inline"`
-	Command           string   `json:"command"`
-	Args              []string `json:"args"`
-	OutputDir         string   `json:"outputDir" yaml:"outputDir"`
+	Command           string            `json:"command"`
+	Args              []string          `json:"args"`
+	OutputDir         string            `json:"outputDir" yaml:"outputDir"`
+	Config            map[string]string `json:"config" yaml:"config"`
 }
 
 type HostCollect struct {
diff --git a/pkg/collect/host_run.go b/pkg/collect/host_run.go
index 2cea2b977..56fc3ca9a 100644
--- a/pkg/collect/host_run.go
+++ b/pkg/collect/host_run.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"encoding/json"
 	"fmt"
+	"gopkg.in/yaml.v2"
 	"os"
 	"os/exec"
 	"path"
@@ -19,6 +20,7 @@ type HostRunInfo struct {
 	ExitCode  string `json:"exitCode"`
 	Error     string `json:"error"`
 	OutputDir string `json:"outputDir"`
+	Config    string `json:"config"`
 }
 
 type CollectHostRun struct {
@@ -40,21 +42,44 @@ func (c *CollectHostRun) Collect(progressChan chan<- interface{}) (map[string][]
 	if collectorName == "" {
 		collectorName = "run-host"
 	}
-	var bundleOutputPath string
+	var (
+		bundleOutputPath          string
+		bundleOutputRelativePath  string
+		cmdConfigFilePath         string
+		cmdConfigFileRelativePath string
+	)
 
 	cmd := exec.Command(runHostCollector.Command, runHostCollector.Args...)
+
 	cmd.Env = append(os.Environ(),
 		fmt.Sprintf("BUNDLE_PATH=%s", c.BundlePath),
 	)
+	collectorRelativePath := filepath.Join("host-collectors/run-host", collectorName)
 
 	// if we choose to save result for the command run
 	if runHostCollector.OutputDir != "" {
-		bundleOutputPath = path.Join(c.BundlePath, "host-collectors/run-host", collectorName, runHostCollector.OutputDir)
+		bundleOutputRelativePath = filepath.Join(collectorRelativePath, runHostCollector.OutputDir)
+		bundleOutputPath = path.Join(c.BundlePath, bundleOutputRelativePath)
 		cmd.Env = append(cmd.Env,
 			fmt.Sprintf("WORKSPACE_DIR=%s", bundleOutputPath),
 		)
 	}
 
+	output := NewResult()
+
+	if runHostCollector.Config != nil {
+		cmdConfigFileRelativePath = filepath.Join(collectorRelativePath, "command-config.yaml")
+		cmdConfigFilePath = path.Join(c.BundlePath, cmdConfigFileRelativePath)
+		configBytes, err := yaml.Marshal(runHostCollector.Config)
+		if err != nil {
+			return nil, errors.Wrap(err, fmt.Sprintf("failed to marshal the command config in collector: %s", collectorName))
+		}
+		output.SaveResult(c.BundlePath, cmdConfigFileRelativePath, bytes.NewBuffer(configBytes))
+		cmd.Env = append(cmd.Env,
+			fmt.Sprintf("CONFIG=%s", cmdConfigFilePath),
+		)
+	}
+
 	var stdout, stderr bytes.Buffer
 	cmd.Stdout = &stdout
 	cmd.Stderr = &stderr
@@ -62,7 +87,8 @@ func (c *CollectHostRun) Collect(progressChan chan<- interface{}) (map[string][]
 	runInfo := HostRunInfo{
 		Command:   cmd.String(),
 		ExitCode:  "0",
-		OutputDir: bundleOutputPath,
+		OutputDir: bundleOutputRelativePath,
+		Config:    cmdConfigFilePath,
 	}
 
 	err := cmd.Run()
@@ -75,21 +101,19 @@ func (c *CollectHostRun) Collect(progressChan chan<- interface{}) (map[string][]
 		}
 	}
 
-	resultInfo := filepath.Join("host-collectors/run-host", collectorName, "collector-info.json")
-	result := filepath.Join("host-collectors/run-host", collectorName, "collector.txt")
+	resultInfo := filepath.Join(collectorRelativePath, "collector-info.json")
+	result := filepath.Join(collectorRelativePath, "collector.txt")
 
 	b, err := json.Marshal(runInfo)
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to marshal run host result")
 	}
 
-	output := NewResult()
 	output.SaveResult(c.BundlePath, resultInfo, bytes.NewBuffer(b))
 	output.SaveResult(c.BundlePath, result, bytes.NewBuffer(stdout.Bytes()))
 	// walkthrough the output directory and save result for each file
 	if runHostCollector.OutputDir != "" {
-		output.SaveResults(c.BundlePath,
-			filepath.Join("host-collectors/run-host", collectorName, runHostCollector.OutputDir))
+		output.SaveResults(c.BundlePath, bundleOutputRelativePath)
 	}
 
 	return output, nil