Merge pull request #250 from secrethub/feature/refactor-run-command-file-reading

Refactor run command

Author: SimonBarendse

internals/secrethub/run.go

@@ -37,15 +37,16 @@ var (
 	ErrSignalFailed           = errRun.Code("signal_failed").ErrorPref("error while propagating signal to process: %s")
 	ErrReadEnvDir             = errRun.Code("env_dir_read_error").ErrorPref("could not read the environment directory: %s")
 	ErrReadEnvFile            = errRun.Code("env_file_read_error").ErrorPref("could not read the environment file %s: %s")
-	ErrEnvDirNotFound         = errRun.Code("env_dir_not_found").Error(fmt.Sprintf("could not find specified environment. Make sure you have executed `%s set`.", ApplicationName))
+	ErrReadDefaultEnvFile     = errRun.Code("default_env_file_read_error").ErrorPref("could not read default run env-file %s: %s")
 	ErrTemplate               = errRun.Code("invalid_template").ErrorPref("could not parse template at line %d: %s")
 	ErrParsingTemplate        = errRun.Code("template_parsing_failed").ErrorPref("error while processing template file '%s': %s")
 	ErrInvalidTemplateVar     = errRun.Code("invalid_template_var").ErrorPref("template variable '%s' is invalid: template variables may only contain uppercase letters, digits, and the '_' (underscore) and are not allowed to start with a number")
 	ErrSecretsNotAllowedInKey = errRun.Code("secret_in_key").Error("secrets are not allowed in run template keys")
 )
 
 const (
-	maskString = "<redacted by SecretHub>"
+	defaultEnvFile = "secrethub.env"
+	maskString     = "<redacted by SecretHub>"
 	// templateVarEnvVarPrefix is used to prefix environment variables
 	// that should be used as template variables.
 	templateVarEnvVarPrefix = "SECRETHUB_VAR_"
@@ -58,9 +59,11 @@ const (
 // defined with --envar or --env-file flags and secrets.yml files.
 // The yml files write to .secretsenv/<env-name> when running the set command.
 type RunCommand struct {
-	command                      []string
 	io                           ui.IO
 	osEnv                        []string
+	readFile                     func(filename string) ([]byte, error)
+	osStat                       func(filename string) (os.FileInfo, error)
+	command                      []string
 	envar                        map[string]string
 	envFile                      string
 	templateVars                 map[string]string
@@ -78,6 +81,8 @@ func NewRunCommand(io ui.IO, newClient newClientFunc) *RunCommand {
 	return &RunCommand{
 		io:           io,
 		osEnv:        os.Environ(),
+		readFile:     ioutil.ReadFile,
+		osStat:       os.Stat,
 		envar:        make(map[string]string),
 		templateVars: make(map[string]string),
 		newClient:    newClient,
@@ -112,66 +117,155 @@ func (cmd *RunCommand) Register(r command.Registerer) {
 // Run reads files from the .secretsenv/<env-name> directory, sets them as environment variables and runs the given command.
 // Note that the environment variables are only passed to the child process and not exported globally, which is nice.
 func (cmd *RunCommand) Run() error {
-	// Parse
-	envSources := []EnvSource{}
+	environment, secrets, err := cmd.sourceEnvironment()
+	if err != nil {
+		return err
+	}
+
+	// This makes sure commands encapsulated in quotes also work.
+	if len(cmd.command) == 1 {
+		cmd.command = strings.Split(cmd.command[0], " ")
+	}
+
+	valuesToMask := make([][]byte, 0, len(secrets))
+	for _, val := range secrets {
+		if val != "" {
+			valuesToMask = append(valuesToMask, []byte(val))
+		}
+	}
+
+	maskedStdout := masker.NewMaskedWriter(cmd.io.Stdout(), valuesToMask, maskString, cmd.maskingTimeout)
+	maskedStderr := masker.NewMaskedWriter(os.Stderr, valuesToMask, maskString, cmd.maskingTimeout)
+
+	command := exec.Command(cmd.command[0], cmd.command[1:]...)
+	command.Env = environment
+	command.Stdin = os.Stdin
+	if cmd.noMasking {
+		command.Stdout = cmd.io.Stdout()
+		command.Stderr = os.Stderr
+	} else {
+		command.Stdout = maskedStdout
+		command.Stderr = maskedStderr
+
+		go maskedStdout.Run()
+		go maskedStderr.Run()
+	}
+
+	err = command.Start()
+	if err != nil {
+		return ErrStartFailed(err)
+	}
+
+	done := make(chan bool, 1)
+
+	// Pass all signals to child process
+	signals := make(chan os.Signal, 1)
+	signal.Notify(signals)
+
+	go func() {
+		select {
+		case s := <-signals:
+			err := command.Process.Signal(s)
+			if err != nil && !strings.Contains(err.Error(), "process already finished") {
+				fmt.Fprintln(os.Stderr, ErrSignalFailed(err))
+			}
+		case <-done:
+			signal.Stop(signals)
+			return
+		}
+	}()
 
+	commandErr := command.Wait()
+	done <- true
+
+	if !cmd.noMasking {
+		err = maskedStdout.Flush()
+		if err != nil {
+			fmt.Fprintln(os.Stderr, err)
+		}
+		err = maskedStderr.Flush()
+		if err != nil {
+			fmt.Fprintln(os.Stderr, err)
+		}
+	}
+
+	if commandErr != nil {
+		// Check if the program exited with an error
+		exitErr, ok := commandErr.(*exec.ExitError)
+		if ok {
+			waitStatus, ok := exitErr.Sys().(syscall.WaitStatus)
+			if ok {
+				// Return the status code returned by the process
+				os.Exit(waitStatus.ExitStatus())
+				return nil
+			}
+
+		}
+		return commandErr
+	}
+
+	return nil
+}
+
+// sourceEnvironment returns the environment of the subcommand, with all the secrets sourced
+// and the secret values that need to be masked.
+func (cmd *RunCommand) sourceEnvironment() ([]string, []string, error) {
 	osEnv, passthroughEnv := parseKeyValueStringsToMap(cmd.osEnv)
 
+	envSources := []EnvSource{}
+
 	referenceEnv := newReferenceEnv(osEnv)
 	envSources = append(envSources, referenceEnv)
 
 	// TODO: Validate the flags when parsing by implementing the Flag interface for EnvFlags.
 	flagSource, err := NewEnvFlags(cmd.envar)
 	if err != nil {
-		return err
+		return nil, nil, err
 	}
 	envSources = append(envSources, flagSource)
 
 	if cmd.envFile == "" {
-		const defaultEnvFile = "secrethub.env"
-		_, err := os.Stat(defaultEnvFile)
-		if err != nil {
-			if !os.IsNotExist(err) {
-				return fmt.Errorf("could not read default run env-file %s: %s", defaultEnvFile, err)
-			}
-		} else {
+		_, err := cmd.osStat(defaultEnvFile)
+		if err == nil {
 			cmd.envFile = defaultEnvFile
+		} else if !os.IsNotExist(err) {
+			return nil, nil, ErrReadDefaultEnvFile(defaultEnvFile, err)
 		}
 	}
 
 	if cmd.envFile != "" {
 		templateVariableReader, err := newVariableReader(osEnv, cmd.templateVars)
 		if err != nil {
-			return err
+			return nil, nil, err
 		}
 
 		if !cmd.dontPromptMissingTemplateVar {
 			templateVariableReader = newPromptMissingVariableReader(templateVariableReader, cmd.io)
 		}
 
-		raw, err := ioutil.ReadFile(cmd.envFile)
+		raw, err := cmd.readFile(cmd.envFile)
 		if err != nil {
-			return ErrCannotReadFile(cmd.envFile, err)
+			return nil, nil, ErrCannotReadFile(cmd.envFile, err)
 		}
 
 		parser, err := getTemplateParser(raw, cmd.templateVersion)
 		if err != nil {
-			return err
+			return nil, nil, err
 		}
 
-		envFile, err := ReadEnvFile(cmd.envFile, templateVariableReader, parser)
+		envFile, err := ReadEnvFile(cmd.envFile, bytes.NewReader(raw), templateVariableReader, parser)
 		if err != nil {
-			return err
+			return nil, nil, err
 		}
 		envSources = append(envSources, envFile)
 	}
 
 	envDir := filepath.Join(secretspec.SecretEnvPath, cmd.env)
-	_, err = os.Stat(envDir)
+	_, err = cmd.osStat(envDir)
 	if err == nil {
 		dirSource, err := NewEnvDir(envDir)
 		if err != nil {
-			return err
+			return nil, nil, err
 		}
 		envSources = append(envSources, dirSource)
 	}
@@ -193,7 +287,7 @@ func (cmd *RunCommand) Run() error {
 	for path := range secrets {
 		secret, err := secretReader.ReadSecret(path)
 		if err != nil {
-			return err
+			return nil, nil, err
 		}
 		secrets[path] = secret
 	}
@@ -203,7 +297,7 @@ func (cmd *RunCommand) Run() error {
 	for _, source := range envSources {
 		pairs, err := source.Env(secrets, secretReader)
 		if err != nil {
-			return err
+			return nil, nil, err
 		}
 
 		for key, value := range pairs {
@@ -215,7 +309,7 @@ func (cmd *RunCommand) Run() error {
 		}
 	}
 
-	// Finally, source the remaining envars from the OS environment.
+	// Source the remaining envars from the OS environment.
 	for key, value := range osEnv {
 		// Only set a variable if it wasn't set by a configured source.
 		_, found := environment[key]
@@ -224,91 +318,10 @@ func (cmd *RunCommand) Run() error {
 		}
 	}
 
-	// This makes sure commands encapsulated in quotes also work.
-	if len(cmd.command) == 1 {
-		cmd.command = strings.Split(cmd.command[0], " ")
-	}
-
-	values := secretReader.Values()
+	// Finally add the unparsed variables
+	processedOsEnv := append(passthroughEnv, mapToKeyValueStrings(environment)...)
 
-	valuesToMask := make([][]byte, 0, len(values))
-	for _, val := range values {
-		if val != "" {
-			valuesToMask = append(valuesToMask, []byte(val))
-		}
-	}
-
-	maskedStdout := masker.NewMaskedWriter(cmd.io.Stdout(), valuesToMask, maskString, cmd.maskingTimeout)
-	maskedStderr := masker.NewMaskedWriter(os.Stderr, valuesToMask, maskString, cmd.maskingTimeout)
-
-	command := exec.Command(cmd.command[0], cmd.command[1:]...)
-	command.Env = append(passthroughEnv, mapToKeyValueStrings(environment)...)
-	command.Stdin = os.Stdin
-	if cmd.noMasking {
-		command.Stdout = cmd.io.Stdout()
-		command.Stderr = os.Stderr
-	} else {
-		command.Stdout = maskedStdout
-		command.Stderr = maskedStderr
-
-		go maskedStdout.Run()
-		go maskedStderr.Run()
-	}
-
-	err = command.Start()
-	if err != nil {
-		return ErrStartFailed(err)
-	}
-
-	done := make(chan bool, 1)
-
-	// Pass all signals to child process
-	signals := make(chan os.Signal, 1)
-	signal.Notify(signals)
-
-	go func() {
-		select {
-		case s := <-signals:
-			err := command.Process.Signal(s)
-			if err != nil && !strings.Contains(err.Error(), "process already finished") {
-				fmt.Fprintln(os.Stderr, ErrSignalFailed(err))
-			}
-		case <-done:
-			signal.Stop(signals)
-			return
-		}
-	}()
-
-	commandErr := command.Wait()
-	done <- true
-
-	if !cmd.noMasking {
-		err = maskedStdout.Flush()
-		if err != nil {
-			fmt.Fprintln(os.Stderr, err)
-		}
-		err = maskedStderr.Flush()
-		if err != nil {
-			fmt.Fprintln(os.Stderr, err)
-		}
-	}
-
-	if commandErr != nil {
-		// Check if the program exited with an error
-		exitErr, ok := commandErr.(*exec.ExitError)
-		if ok {
-			waitStatus, ok := exitErr.Sys().(syscall.WaitStatus)
-			if ok {
-				// Return the status code returned by the process
-				os.Exit(waitStatus.ExitStatus())
-				return nil
-			}
-
-		}
-		return commandErr
-	}
-
-	return nil
+	return processedOsEnv, secretReader.Values(), nil
 }
 
 // mapToKeyValueStrings converts a map to a slice of key=value pairs.
@@ -412,12 +425,8 @@ func (t envTemplate) Secrets() []string {
 }
 
 // ReadEnvFile reads and parses a .env file.
-func ReadEnvFile(filepath string, varReader tpl.VariableReader, parser tpl.Parser) (EnvFile, error) {
-	r, err := os.Open(filepath)
-	if err != nil {
-		return EnvFile{}, ErrCannotReadFile(filepath, err)
-	}
-	env, err := NewEnv(r, varReader, parser)
+func ReadEnvFile(filepath string, reader io.Reader, varReader tpl.VariableReader, parser tpl.Parser) (EnvFile, error) {
+	env, err := NewEnv(reader, varReader, parser)
 	if err != nil {
 		return EnvFile{}, ErrParsingTemplate(filepath, err)
 	}
@@ -667,13 +676,6 @@ type EnvDir map[string]string
 // NewEnvDir sources environment variables from files in a given directory,
 // using the file name as key and contents as value.
 func NewEnvDir(path string) (EnvDir, error) {
-	_, err := os.Stat(path)
-	if os.IsNotExist(err) {
-		return nil, ErrEnvDirNotFound
-	} else if err != nil {
-		return nil, ErrReadEnvDir(err)
-	}
-
 	files, err := ioutil.ReadDir(path)
 	if err != nil {
 		return nil, ErrReadEnvDir(err)