From c2bc901a414ac640e88a06ee2ddc0af7cc97a341 Mon Sep 17 00:00:00 2001
From: Jan De Dobbeleer <jan.de.dobbeleer@gmail.com>
Date: Tue, 5 Jan 2021 13:05:37 +0100
Subject: [PATCH] refactor: command caching without leaking

---
 src/environment.go            | 48 ++++++++++++++++++++---------------
 src/main.go                   |  4 ++-
 src/segment_az.go             |  6 ++---
 src/segment_az_test.go        |  2 +-
 src/segment_command.go        |  3 +--
 src/segment_command_test.go   | 40 ++++++++++++++++++++---------
 src/segment_dotnet_test.go    |  2 +-
 src/segment_git.go            | 17 ++++++-------
 src/segment_git_test.go       | 17 +++++--------
 src/segment_kubectl.go        |  6 ++---
 src/segment_kubectl_test.go   |  2 +-
 src/segment_language.go       |  5 ++--
 src/segment_language_test.go  |  2 +-
 src/segment_path_test.go      |  4 +--
 src/segment_python_test.go    |  2 +-
 src/segment_terraform.go      |  6 ++---
 src/segment_terraform_test.go |  2 +-
 17 files changed, 92 insertions(+), 76 deletions(-)

diff --git a/src/environment.go b/src/environment.go
index 9065a43a..b7df4e8a 100644
--- a/src/environment.go
+++ b/src/environment.go
@@ -3,7 +3,6 @@ package main
 import (
 	"context"
 	"io/ioutil"
-	"log"
 	"net/http"
 	"os"
 	"os/exec"
@@ -22,6 +21,15 @@ const (
 	windowsPlatform = "windows"
 )
 
+type commandError struct {
+	err      string
+	exitCode int
+}
+
+func (e *commandError) Error() string {
+	return e.err
+}
+
 type environmentInfo interface {
 	getenv(key string) string
 	getcwd() string
@@ -36,7 +44,7 @@ type environmentInfo interface {
 	getHostName() (string, error)
 	getRuntimeGOOS() string
 	getPlatform() string
-	hasCommand(command string) (string, bool)
+	hasCommand(command string) bool
 	runCommand(command string, args ...string) (string, error)
 	runShellCommand(shell, command string) string
 	lastErrorCode() int
@@ -49,17 +57,9 @@ type environmentInfo interface {
 }
 
 type environment struct {
-	args *args
-	cwd  string
-}
-
-type commandError struct {
-	err      string
-	exitCode int
-}
-
-func (e *commandError) Error() string {
-	return e.err
+	args     *args
+	cwd      string
+	commands map[string]string
 }
 
 func (env *environment) getenv(key string) string {
@@ -152,6 +152,9 @@ func (env *environment) getPlatform() string {
 }
 
 func (env *environment) runCommand(command string, args ...string) (string, error) {
+	if cmd, ok := env.commands[command]; ok {
+		command = cmd
+	}
 	out, err := exec.Command(command, args...).Output()
 	if err != nil {
 		return "", err
@@ -160,17 +163,20 @@ func (env *environment) runCommand(command string, args ...string) (string, erro
 }
 
 func (env *environment) runShellCommand(shell, command string) string {
-	out, err := exec.Command(shell, "-c", command).Output()
-	if err != nil {
-		log.Println(err)
-		return ""
-	}
-	return strings.TrimSpace(string(out))
+	out, _ := env.runCommand(shell, "-c", command)
+	return out
 }
 
-func (env *environment) hasCommand(command string) (string, bool) {
+func (env *environment) hasCommand(command string) bool {
+	if _, ok := env.commands[command]; ok {
+		return true
+	}
 	path, err := exec.LookPath(command)
-	return path, err == nil
+	if err == nil {
+		env.commands[command] = path
+		return true
+	}
+	return false
 }
 
 func (env *environment) lastErrorCode() int {
diff --git a/src/main.go b/src/main.go
index 7ee8fefc..059e0a64 100644
--- a/src/main.go
+++ b/src/main.go
@@ -97,7 +97,9 @@ func main() {
 	}
 	flag.Parse()
 	env := &environment{
-		args: args,
+		args:     args,
+		commands: make(map[string]string),
+		cwd:      *args.PWD,
 	}
 	if *args.Millis {
 		fmt.Print(time.Now().UnixNano() / 1000000)
diff --git a/src/segment_az.go b/src/segment_az.go
index 8a3baa7c..c047bbe3 100644
--- a/src/segment_az.go
+++ b/src/segment_az.go
@@ -36,12 +36,12 @@ func (a *az) init(props *properties, env environmentInfo) {
 }
 
 func (a *az) enabled() bool {
-	commandPath, commandExists := a.env.hasCommand("az")
-	if (!a.idEnabled() && !a.nameEnabled()) || !commandExists {
+	cmd := "az"
+	if (!a.idEnabled() && !a.nameEnabled()) || !a.env.hasCommand(cmd) {
 		return false
 	}
 
-	output, _ := a.env.runCommand(commandPath, "account", "show", "--query=[name,id]", "-o=tsv")
+	output, _ := a.env.runCommand(cmd, "account", "show", "--query=[name,id]", "-o=tsv")
 	if output == "" {
 		return false
 	}
diff --git a/src/segment_az_test.go b/src/segment_az_test.go
index fe6b7365..98c7a1eb 100644
--- a/src/segment_az_test.go
+++ b/src/segment_az_test.go
@@ -18,7 +18,7 @@ type azArgs struct {
 
 func bootStrapAzTest(args *azArgs) *az {
 	env := new(MockedEnvironment)
-	env.On("hasCommand", "az").Return("az", args.enabled)
+	env.On("hasCommand", "az").Return(args.enabled)
 	env.On("runCommand", "az", []string{"account", "show", "--query=[name,id]", "-o=tsv"}).Return(fmt.Sprintf("%s\n%s\n", args.subscriptionName, args.subscriptionID), nil)
 	props := &properties{
 		values: map[Property]interface{}{
diff --git a/src/segment_command.go b/src/segment_command.go
index 8c378752..bbc26ac0 100644
--- a/src/segment_command.go
+++ b/src/segment_command.go
@@ -17,8 +17,7 @@ const (
 
 func (c *command) enabled() bool {
 	shell := c.props.getString(ExecutableShell, "bash")
-	shell, commandExists := c.env.hasCommand(shell)
-	if !commandExists {
+	if !c.env.hasCommand(shell) {
 		return false
 	}
 	command := c.props.getString(Command, "echo no command specified")
diff --git a/src/segment_command_test.go b/src/segment_command_test.go
index 79de1c4d..9bb86c7b 100644
--- a/src/segment_command_test.go
+++ b/src/segment_command_test.go
@@ -7,7 +7,9 @@ import (
 )
 
 func TestExecuteCommand(t *testing.T) {
-	env := &environment{}
+	env := &environment{
+		commands: make(map[string]string),
+	}
 	props := &properties{
 		values: map[Property]interface{}{
 			Command: "echo hello",
@@ -19,11 +21,13 @@ func TestExecuteCommand(t *testing.T) {
 	}
 	enabled := c.enabled()
 	assert.True(t, enabled)
-	assert.Equal(t, c.string(), "hello")
+	assert.Equal(t, "hello", c.string())
 }
 
 func TestExecuteMultipleCommandsOrFirst(t *testing.T) {
-	env := &environment{}
+	env := &environment{
+		commands: make(map[string]string),
+	}
 	props := &properties{
 		values: map[Property]interface{}{
 			Command: "exit 1 || echo hello",
@@ -35,11 +39,13 @@ func TestExecuteMultipleCommandsOrFirst(t *testing.T) {
 	}
 	enabled := c.enabled()
 	assert.True(t, enabled)
-	assert.Equal(t, c.string(), "hello")
+	assert.Equal(t, "hello", c.string())
 }
 
 func TestExecuteMultipleCommandsOrSecond(t *testing.T) {
-	env := &environment{}
+	env := &environment{
+		commands: make(map[string]string),
+	}
 	props := &properties{
 		values: map[Property]interface{}{
 			Command: "echo hello || echo world",
@@ -51,11 +57,13 @@ func TestExecuteMultipleCommandsOrSecond(t *testing.T) {
 	}
 	enabled := c.enabled()
 	assert.True(t, enabled)
-	assert.Equal(t, c.string(), "hello")
+	assert.Equal(t, "hello", c.string())
 }
 
 func TestExecuteMultipleCommandsAnd(t *testing.T) {
-	env := &environment{}
+	env := &environment{
+		commands: make(map[string]string),
+	}
 	props := &properties{
 		values: map[Property]interface{}{
 			Command: "echo hello && echo world",
@@ -67,11 +75,13 @@ func TestExecuteMultipleCommandsAnd(t *testing.T) {
 	}
 	enabled := c.enabled()
 	assert.True(t, enabled)
-	assert.Equal(t, c.string(), "helloworld")
+	assert.Equal(t, "helloworld", c.string())
 }
 
 func TestExecuteSingleCommandEmpty(t *testing.T) {
-	env := &environment{}
+	env := &environment{
+		commands: make(map[string]string),
+	}
 	props := &properties{
 		values: map[Property]interface{}{
 			Command: "",
@@ -86,7 +96,9 @@ func TestExecuteSingleCommandEmpty(t *testing.T) {
 }
 
 func TestExecuteSingleCommandNoCommandProperty(t *testing.T) {
-	env := &environment{}
+	env := &environment{
+		commands: make(map[string]string),
+	}
 	props := &properties{}
 	c := &command{
 		props: props,
@@ -98,7 +110,9 @@ func TestExecuteSingleCommandNoCommandProperty(t *testing.T) {
 }
 
 func TestExecuteMultipleCommandsAndDisabled(t *testing.T) {
-	env := &environment{}
+	env := &environment{
+		commands: make(map[string]string),
+	}
 	props := &properties{
 		values: map[Property]interface{}{
 			Command: "echo && echo",
@@ -113,7 +127,9 @@ func TestExecuteMultipleCommandsAndDisabled(t *testing.T) {
 }
 
 func TestExecuteMultipleCommandsOrDisabled(t *testing.T) {
-	env := &environment{}
+	env := &environment{
+		commands: make(map[string]string),
+	}
 	props := &properties{
 		values: map[Property]interface{}{
 			Command: "echo|| echo",
diff --git a/src/segment_dotnet_test.go b/src/segment_dotnet_test.go
index f8f7f88d..6e501d00 100644
--- a/src/segment_dotnet_test.go
+++ b/src/segment_dotnet_test.go
@@ -16,7 +16,7 @@ type dotnetArgs struct {
 
 func bootStrapDotnetTest(args *dotnetArgs) *dotnet {
 	env := new(MockedEnvironment)
-	env.On("hasCommand", "dotnet").Return("dotnet", args.enabled)
+	env.On("hasCommand", "dotnet").Return(args.enabled)
 	if args.unsupported {
 		err := &commandError{exitCode: 145}
 		env.On("runCommand", "dotnet", []string{"--version"}).Return("", err)
diff --git a/src/segment_git.go b/src/segment_git.go
index c27c8135..cc08f0ce 100644
--- a/src/segment_git.go
+++ b/src/segment_git.go
@@ -47,10 +47,9 @@ func (s *gitStatus) string(prefix, color string) string {
 }
 
 type git struct {
-	props       *properties
-	env         environmentInfo
-	repo        *gitRepo
-	commandPath string
+	props *properties
+	env   environmentInfo
+	repo  *gitRepo
 }
 
 const (
@@ -114,15 +113,15 @@ const (
 	BehindColor Property = "behind_color"
 	// AheadColor if set, the color to use when the branch is ahead and behind the remote
 	AheadColor Property = "ahead_color"
+
+	gitCommand = "git"
 )
 
 func (g *git) enabled() bool {
-	commandPath, commandExists := g.env.hasCommand("git")
-	if !commandExists {
+	if !g.env.hasCommand("git") {
 		return false
 	}
-	g.commandPath = commandPath
-	output, _ := g.env.runCommand(g.commandPath, "rev-parse", "--is-inside-work-tree")
+	output, _ := g.env.runCommand(gitCommand, "rev-parse", "--is-inside-work-tree")
 	return output == "true"
 }
 
@@ -242,7 +241,7 @@ func (g *git) getStatusColor(defaultValue string) string {
 
 func (g *git) getGitCommandOutput(args ...string) string {
 	args = append([]string{"-c", "core.quotepath=false", "-c", "color.status=false"}, args...)
-	val, _ := g.env.runCommand(g.commandPath, args...)
+	val, _ := g.env.runCommand(gitCommand, args...)
 	return val
 }
 
diff --git a/src/segment_git_test.go b/src/segment_git_test.go
index 8f831bbc..dd4159ad 100644
--- a/src/segment_git_test.go
+++ b/src/segment_git_test.go
@@ -12,7 +12,7 @@ const (
 
 func TestEnabledGitNotFound(t *testing.T) {
 	env := new(MockedEnvironment)
-	env.On("hasCommand", "git").Return("", false)
+	env.On("hasCommand", "git").Return(false)
 	g := &git{
 		env: env,
 	}
@@ -21,7 +21,7 @@ func TestEnabledGitNotFound(t *testing.T) {
 
 func TestEnabledInWorkingDirectory(t *testing.T) {
 	env := new(MockedEnvironment)
-	env.On("hasCommand", "git").Return("git", true)
+	env.On("hasCommand", "git").Return(true)
 	env.On("runCommand", "git", []string{"rev-parse", "--is-inside-work-tree"}).Return("true", nil)
 	g := &git{
 		env: env,
@@ -36,8 +36,7 @@ func TestGetGitOutputForCommand(t *testing.T) {
 	env := new(MockedEnvironment)
 	env.On("runCommand", "git", append(args, commandArgs...)).Return(want, nil)
 	g := &git{
-		env:         env,
-		commandPath: "git",
+		env: env,
 	}
 	got := g.getGitCommandOutput(commandArgs...)
 	assert.Equal(t, want, got)
@@ -86,7 +85,6 @@ func setupHEADContextEnv(context *detachedContext) *git {
 		repo: &gitRepo{
 			root: "",
 		},
-		commandPath: "git",
 	}
 	return g
 }
@@ -213,8 +211,7 @@ func TestGetStashContextZeroEntries(t *testing.T) {
 	env := new(MockedEnvironment)
 	env.On("runCommand", "git", []string{"-c", "core.quotepath=false", "-c", "color.status=false", "rev-list", "--walk-reflogs", "--count", "refs/stash"}).Return("", nil)
 	g := &git{
-		env:         env,
-		commandPath: "git",
+		env: env,
 	}
 	got := g.getStashContext()
 	assert.Equal(t, want, got)
@@ -225,8 +222,7 @@ func TestGetStashContextMultipleEntries(t *testing.T) {
 	env := new(MockedEnvironment)
 	env.On("runCommand", "git", []string{"-c", "core.quotepath=false", "-c", "color.status=false", "rev-list", "--walk-reflogs", "--count", "refs/stash"}).Return("2", nil)
 	g := &git{
-		env:         env,
-		commandPath: "git",
+		env: env,
 	}
 	got := g.getStashContext()
 	assert.Equal(t, want, got)
@@ -394,8 +390,7 @@ func bootstrapUpstreamTest(upstream string) *git {
 		repo: &gitRepo{
 			upstream: "origin/main",
 		},
-		props:       props,
-		commandPath: "git",
+		props: props,
 	}
 	return g
 }
diff --git a/src/segment_kubectl.go b/src/segment_kubectl.go
index 8d2c8411..52e88974 100644
--- a/src/segment_kubectl.go
+++ b/src/segment_kubectl.go
@@ -16,10 +16,10 @@ func (k *kubectl) init(props *properties, env environmentInfo) {
 }
 
 func (k *kubectl) enabled() bool {
-	commandPath, commandExists := k.env.hasCommand("kubectl")
-	if !commandExists {
+	cmd := "kubectl"
+	if !k.env.hasCommand(cmd) {
 		return false
 	}
-	k.contextName, _ = k.env.runCommand(commandPath, "config", "current-context")
+	k.contextName, _ = k.env.runCommand(cmd, "config", "current-context")
 	return k.contextName != ""
 }
diff --git a/src/segment_kubectl_test.go b/src/segment_kubectl_test.go
index 3fb2f02c..070e8759 100644
--- a/src/segment_kubectl_test.go
+++ b/src/segment_kubectl_test.go
@@ -13,7 +13,7 @@ type kubectlArgs struct {
 
 func bootStrapKubectlTest(args *kubectlArgs) *kubectl {
 	env := new(MockedEnvironment)
-	env.On("hasCommand", "kubectl").Return("kubectl", args.enabled)
+	env.On("hasCommand", "kubectl").Return(args.enabled)
 	env.On("runCommand", "kubectl", []string{"config", "current-context"}).Return(args.contextName, nil)
 	k := &kubectl{
 		env:   env,
diff --git a/src/segment_language.go b/src/segment_language.go
index 2ad4b9d9..b1d5c253 100644
--- a/src/segment_language.go
+++ b/src/segment_language.go
@@ -87,9 +87,8 @@ func (l *language) getVersion() bool {
 // hasCommand checks if one of the commands exists and sets it as executable
 func (l *language) hasCommand() bool {
 	for i, command := range l.commands {
-		commandPath, commandExists := l.env.hasCommand(command)
-		if commandExists {
-			l.executable = commandPath
+		if l.env.hasCommand(command) {
+			l.executable = command
 			break
 		}
 		if i == len(l.commands)-1 {
diff --git a/src/segment_language_test.go b/src/segment_language_test.go
index fb0589ed..33b9afb6 100644
--- a/src/segment_language_test.go
+++ b/src/segment_language_test.go
@@ -37,7 +37,7 @@ func (l *languageArgs) hasvalue(value string, list []string) bool {
 func bootStrapLanguageTest(args *languageArgs) *language {
 	env := new(MockedEnvironment)
 	for _, command := range args.commands {
-		env.On("hasCommand", command).Return(command, args.hasvalue(command, args.enabledCommands))
+		env.On("hasCommand", command).Return(args.hasvalue(command, args.enabledCommands))
 		env.On("runCommand", command, []string{args.versionParam}).Return(args.version, nil)
 	}
 	for _, extension := range args.extensions {
diff --git a/src/segment_path_test.go b/src/segment_path_test.go
index 5cfc3ffb..21479259 100644
--- a/src/segment_path_test.go
+++ b/src/segment_path_test.go
@@ -72,9 +72,9 @@ func (env *MockedEnvironment) getPlatform() string {
 	return args.String(0)
 }
 
-func (env *MockedEnvironment) hasCommand(command string) (string, bool) {
+func (env *MockedEnvironment) hasCommand(command string) bool {
 	args := env.Called(command)
-	return args.String(0), args.Bool(1)
+	return args.Bool(0)
 }
 
 func (env *MockedEnvironment) runCommand(command string, args ...string) (string, error) {
diff --git a/src/segment_python_test.go b/src/segment_python_test.go
index e99d92fc..984ce0f8 100644
--- a/src/segment_python_test.go
+++ b/src/segment_python_test.go
@@ -16,7 +16,7 @@ type pythonArgs struct {
 
 func bootStrapPythonTest(args *pythonArgs) *python {
 	env := new(MockedEnvironment)
-	env.On("hasCommand", "python").Return("python", true)
+	env.On("hasCommand", "python").Return(true)
 	env.On("runCommand", "python", []string{"--version"}).Return("Python 3.8.4", nil)
 	env.On("hasFiles", "*.py").Return(true)
 	env.On("getenv", "VIRTUAL_ENV").Return(args.virtualEnvName)
diff --git a/src/segment_terraform.go b/src/segment_terraform.go
index 26b52855..0188ad75 100644
--- a/src/segment_terraform.go
+++ b/src/segment_terraform.go
@@ -16,10 +16,10 @@ func (tf *terraform) init(props *properties, env environmentInfo) {
 }
 
 func (tf *terraform) enabled() bool {
-	commandPath, commandExists := tf.env.hasCommand("terraform")
-	if !commandExists || !tf.env.hasFolder(".terraform") {
+	cmd := "terraform"
+	if !tf.env.hasCommand(cmd) || !tf.env.hasFolder(".terraform") {
 		return false
 	}
-	tf.workspaceName, _ = tf.env.runCommand(commandPath, "workspace", "show")
+	tf.workspaceName, _ = tf.env.runCommand(cmd, "workspace", "show")
 	return true
 }
diff --git a/src/segment_terraform_test.go b/src/segment_terraform_test.go
index 3970f2c8..dc98ec99 100644
--- a/src/segment_terraform_test.go
+++ b/src/segment_terraform_test.go
@@ -14,7 +14,7 @@ type terraformArgs struct {
 
 func bootStrapTerraformTest(args *terraformArgs) *terraform {
 	env := new(MockedEnvironment)
-	env.On("hasCommand", "terraform").Return("terraform", args.hasTfCommand)
+	env.On("hasCommand", "terraform").Return(args.hasTfCommand)
 	env.On("hasFolder", ".terraform").Return(args.hasTfFolder)
 	env.On("runCommand", "terraform", []string{"workspace", "show"}).Return(args.workspaceName, nil)
 	k := &terraform{