Fix checking for credentials performance (#2452)

Co-authored-by: Jesse Duffield <jessedduffield@gmail.com>

2 files changed, 133 insertions, 11 deletions

diff --git a/pkg/commands/oscommands/cmd_obj_runner.go b/pkg/commands/oscommands/cmd_obj_runner.go
index 498ef4bab..3df1e2916 100644
--- a/pkg/commands/oscommands/cmd_obj_runner.go
+++ b/pkg/commands/oscommands/cmd_obj_runner.go
@@ -323,6 +323,23 @@ func (self *cmdObjRunner) processOutput(reader io.Reader, writer io.Writer, prom
 // having a function that returns a function because we need to maintain some state inbetween calls hence the closure
 func (self *cmdObjRunner) getCheckForCredentialRequestFunc() func([]byte) (CredentialType, bool) {
 	var ttyText strings.Builder
+	prompts := map[string]CredentialType{
+		`Password:`:                              Password,
+		`.+'s password:`:                         Password,
+		`Password\s*for\s*'.+':`:                 Password,
+		`Username\s*for\s*'.+':`:                 Username,
+		`Enter\s*passphrase\s*for\s*key\s*'.+':`: Passphrase,
+		`Enter\s*PIN\s*for\s*.+\s*key\s*.+:`:     PIN,
+	}
+
+	compiledPrompts := map[*regexp.Regexp]CredentialType{}
+	for pattern, askFor := range prompts {
+		compiledPattern := regexp.MustCompile(pattern)
+		compiledPrompts[compiledPattern] = askFor
+	}
+
+	newlineRegex := regexp.MustCompile("\n")
+
 	// this function takes each word of output from the command and builds up a string to see if we're being asked for a password
 	return func(newBytes []byte) (CredentialType, bool) {
 		_, err := ttyText.Write(newBytes)
@@ -330,22 +347,18 @@ func (self *cmdObjRunner) getCheckForCredentialRequestFunc() func([]byte) (Crede
 			self.log.Error(err)
 		}
 
-		prompts := map[string]CredentialType{
-			`Password:`:                              Password,
-			`.+'s password:`:                         Password,
-			`Password\s*for\s*'.+':`:                 Password,
-			`Username\s*for\s*'.+':`:                 Username,
-			`Enter\s*passphrase\s*for\s*key\s*'.+':`: Passphrase,
-			`Enter\s*PIN\s*for\s*.+\s*key\s*.+:`:     PIN,
-		}
-
-		for pattern, askFor := range prompts {
-			if match, _ := regexp.MatchString(pattern, ttyText.String()); match {
+		for pattern, askFor := range compiledPrompts {
+			if match := pattern.Match([]byte(ttyText.String())); match {
 				ttyText.Reset()
 				return askFor, true
 			}
 		}
 
+		if indices := newlineRegex.FindIndex([]byte(ttyText.String())); indices != nil {
+			newText := []byte(ttyText.String()[indices[1]:])
+			ttyText.Reset()
+			ttyText.Write(newText)
+		}
 		return 0, false
 	}
 }
diff --git a/pkg/commands/oscommands/cmd_obj_runner_test.go b/pkg/commands/oscommands/cmd_obj_runner_test.go
new file mode 100644
index 000000000..ab26b9827
--- /dev/null
+++ b/pkg/commands/oscommands/cmd_obj_runner_test.go
@@ -0,0 +1,109 @@
+package oscommands
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/jesseduffield/lazygit/pkg/utils"
+)
+
+func getRunner() *cmdObjRunner {
+	log := utils.NewDummyLog()
+	return &cmdObjRunner{
+		log:   log,
+		guiIO: NewNullGuiIO(log),
+	}
+}
+
+func TestProcessOutput(t *testing.T) {
+	defaultPromptUserForCredential := func(ct CredentialType) string {
+		switch ct {
+		case Password:
+			return "password"
+		case Username:
+			return "username"
+		case Passphrase:
+			return "passphrase"
+		case PIN:
+			return "pin"
+		default:
+			panic("unexpected credential type")
+		}
+	}
+
+	scenarios := []struct {
+		name                    string
+		promptUserForCredential func(CredentialType) string
+		output                  string
+		expectedToWrite         string
+	}{
+		{
+			name:                    "no output",
+			promptUserForCredential: defaultPromptUserForCredential,
+			output:                  "",
+			expectedToWrite:         "",
+		},
+		{
+			name:                    "password prompt",
+			promptUserForCredential: defaultPromptUserForCredential,
+			output:                  "Password:",
+			expectedToWrite:         "password",
+		},
+		{
+			name:                    "password prompt 2",
+			promptUserForCredential: defaultPromptUserForCredential,
+			output:                  "Bill's password:",
+			expectedToWrite:         "password",
+		},
+		{
+			name:                    "password prompt 3",
+			promptUserForCredential: defaultPromptUserForCredential,
+			output:                  "Password for 'Bill':",
+			expectedToWrite:         "password",
+		},
+		{
+			name:                    "username prompt",
+			promptUserForCredential: defaultPromptUserForCredential,
+			output:                  "Username for 'Bill':",
+			expectedToWrite:         "username",
+		},
+		{
+			name:                    "passphrase prompt",
+			promptUserForCredential: defaultPromptUserForCredential,
+			output:                  "Enter passphrase for key '123':",
+			expectedToWrite:         "passphrase",
+		},
+		{
+			name:                    "pin prompt",
+			promptUserForCredential: defaultPromptUserForCredential,
+			output:                  "Enter PIN for key '123':",
+			expectedToWrite:         "pin",
+		},
+		{
+			name:                    "username and password prompt",
+			promptUserForCredential: defaultPromptUserForCredential,
+			output:                  "Password:\nUsername for 'Alice':\n",
+			expectedToWrite:         "passwordusername",
+		},
+		{
+			name:                    "user submits empty credential",
+			promptUserForCredential: func(ct CredentialType) string { return "" },
+			output:                  "Password:\n",
+			expectedToWrite:         "",
+		},
+	}
+
+	for _, scenario := range scenarios {
+		t.Run(scenario.name, func(t *testing.T) {
+			runner := getRunner()
+			reader := strings.NewReader(scenario.output)
+			writer := &strings.Builder{}
+
+			runner.processOutput(reader, writer, scenario.promptUserForCredential)
+
+			if writer.String() != scenario.expectedToWrite {
+				t.Errorf("expected to write '%s' but got '%s'", scenario.expectedToWrite, writer.String())
+			}
+		})
+	}
+}