From 48ee60761b453bb8b70d56f3d2f0471dff25fcad Mon Sep 17 00:00:00 2001 From: York Chen Date: Tue, 17 Oct 2023 16:22:46 -0400 Subject: [PATCH] refactor: mv other envs to processEnvVars --- pkg/collect/host_run.go | 105 +++++++++++++++++++++------------------- 1 file changed, 55 insertions(+), 50 deletions(-) diff --git a/pkg/collect/host_run.go b/pkg/collect/host_run.go index 2c9106235..769815cd6 100644 --- a/pkg/collect/host_run.go +++ b/pkg/collect/host_run.go @@ -43,8 +43,6 @@ func (c *CollectHostRun) Collect(progressChan chan<- interface{}) (map[string][] } var ( bundleOutputRelativePath string - cmdOutputTempDir string - cmdInputTempDir string err error ) @@ -55,47 +53,14 @@ func (c *CollectHostRun) Collect(progressChan chan<- interface{}) (map[string][] ExitCode: "0", } - err = c.processEnvVars(cmd, runInfo) + cmdOutputTempDir, err := c.processEnvVars(cmd) if err != nil { return nil, errors.Wrap(err, "failed to parse env variable") } collectorRelativePath := filepath.Join("host-collectors/run-host", collectorName) - // if we choose to save result for the command run - if runHostCollector.OutputDir != "" { - cmdOutputTempDir, err = os.MkdirTemp("", runHostCollector.OutputDir) - defer os.RemoveAll(cmdOutputTempDir) - if err != nil { - return nil, errors.New(fmt.Sprintf("failed to created temp dir for: %s", runHostCollector.OutputDir)) - } - cmd.Env = append(cmd.Env, - fmt.Sprintf("TS_WORKSPACE_DIR=%s", cmdOutputTempDir), - ) - } - - output := NewResult() - - if runHostCollector.Input != nil { - cmdInputTempDir, err = os.MkdirTemp("", "input") - defer os.RemoveAll(cmdInputTempDir) - if err != nil { - return nil, errors.New("failed to created temp dir for host run input") - } - for inFilename, inFileContent := range runHostCollector.Input { - if strings.Contains(inFileContent, "/") { - return nil, errors.New("Input filename contains '/'") - } - cmdInputFilePath := filepath.Join(cmdInputTempDir, inFilename) - err = os.WriteFile(cmdInputFilePath, []byte(inFileContent), 0644) - if err != nil { - return nil, errors.Wrap(err, fmt.Sprintf("failed to write input file: %s to temp directory", inFilename)) - } - } - cmd.Env = append(cmd.Env, - fmt.Sprintf("TS_INPUT_DIR=%s", cmdInputTempDir), - ) - } + runInfo.Env = cmd.Env var stdout, stderr bytes.Buffer cmd.Stdout = &stdout @@ -111,6 +76,7 @@ func (c *CollectHostRun) Collect(progressChan chan<- interface{}) (map[string][] } } + output := NewResult() resultInfo := filepath.Join("host-collectors/run-host", collectorName+"-info.json") result := filepath.Join("host-collectors/run-host", collectorName+".txt") @@ -131,7 +97,13 @@ func (c *CollectHostRun) Collect(progressChan chan<- interface{}) (map[string][] return output, nil } -func (c *CollectHostRun) processEnvVars(cmd *exec.Cmd, runInfo *HostRunInfo) error { +func (c *CollectHostRun) processEnvVars(cmd *exec.Cmd) (string, error) { + var ( + cmdOutputTempDir string + err error + ) + // clears the parent env vars + cmd.Env = []string{} guaranteedEnvs := []string{"PATH", "KUBECONFIG"} for _, key := range guaranteedEnvs { guaranteedEnvVal, found := os.LookupEnv(key) @@ -142,16 +114,18 @@ func (c *CollectHostRun) processEnvVars(cmd *exec.Cmd, runInfo *HostRunInfo) err } runHostCollector := c.hostCollector - if runHostCollector.InheritParentEnvs == true { - cmd.Env = append(cmd.Env, os.Environ()...) - } else if runHostCollector.InheritEnvs != nil { - for _, key := range runHostCollector.InheritEnvs { - envVal, found := os.LookupEnv(key) - if !found { - return errors.New(fmt.Sprintf("inherit env variable is not found: %s", key)) + if runHostCollector.InheritParentEnvs { + if runHostCollector.InheritEnvs != nil { + for _, key := range runHostCollector.InheritEnvs { + envVal, found := os.LookupEnv(key) + if !found { + return "", errors.New(fmt.Sprintf("inherit env variable is not found: %s", key)) + } + cmd.Env = append(cmd.Env, + fmt.Sprintf("%s=%s", key, envVal)) } - cmd.Env = append(cmd.Env, - fmt.Sprintf("%s=%s", key, envVal)) + } else { + cmd.Env = append(cmd.Env, os.Environ()...) } } @@ -162,12 +136,43 @@ func (c *CollectHostRun) processEnvVars(cmd *exec.Cmd, runInfo *HostRunInfo) err cmd.Env = append(cmd.Env, fmt.Sprintf("%s=%s", runHostCollector.Env[i], os.Getenv(runHostCollector.Env[i]))) } else { - return errors.New(fmt.Sprintf("env variable entry is missing '=' : %s", runHostCollector.Env[i])) + return "", errors.New(fmt.Sprintf("env variable entry is missing '=' : %s", runHostCollector.Env[i])) } } } - runInfo.Env = os.Environ() + // if we choose to save result for the command run + if runHostCollector.OutputDir != "" { + cmdOutputTempDir, err = os.MkdirTemp("", runHostCollector.OutputDir) + defer os.RemoveAll(cmdOutputTempDir) + if err != nil { + return "", errors.New(fmt.Sprintf("failed to created temp dir for: %s", runHostCollector.OutputDir)) + } + cmd.Env = append(cmd.Env, + fmt.Sprintf("TS_WORKSPACE_DIR=%s", cmdOutputTempDir), + ) + } + + if runHostCollector.Input != nil { + cmdInputTempDir, err := os.MkdirTemp("", "input") + defer os.RemoveAll(cmdInputTempDir) + if err != nil { + return "", errors.New("failed to created temp dir for host run input") + } + for inFilename, inFileContent := range runHostCollector.Input { + if strings.Contains(inFileContent, "/") { + return "", errors.New("Input filename contains '/'") + } + cmdInputFilePath := filepath.Join(cmdInputTempDir, inFilename) + err = os.WriteFile(cmdInputFilePath, []byte(inFileContent), 0644) + if err != nil { + return "", errors.Wrap(err, fmt.Sprintf("failed to write input file: %s to temp directory", inFilename)) + } + } + cmd.Env = append(cmd.Env, + fmt.Sprintf("TS_INPUT_DIR=%s", cmdInputTempDir), + ) + } - return nil + return cmdOutputTempDir, nil }