Simplify 2FA Push auth processing

Use a single channel to handle both Push Auth and OTP results

4 files changed, 228 insertions, 405 deletions

diff --git a/internal/command/twofactorverify/twofactorverify.go b/internal/command/twofactorverify/twofactorverify.go
index 1043bfb..4041de0 100644
--- a/internal/command/twofactorverify/twofactorverify.go
+++ b/internal/command/twofactorverify/twofactorverify.go
@@ -14,105 +14,63 @@ import (
 	"gitlab.com/gitlab-org/gitlab-shell/v14/internal/gitlabnet/twofactorverify"
 )
 
+const (
+	timeout = 30 * time.Second
+	prompt  = "OTP: "
+)
+
 type Command struct {
 	Config     *config.Config
-	Client     *twofactorverify.Client
 	Args       *commandargs.Shell
 	ReadWriter *readwriter.ReadWriter
 }
 
-type Result struct {
-	Error   error
-	Status  string
-	Success bool
-}
-
 func (c *Command) Execute(ctx context.Context) error {
-	ctxlog := log.ContextLogger(ctx)
-
-	// config.GetHTTPClient isn't thread-safe so save Client in struct for concurrency
-	// workaround until #518 is fixed
-	var err error
-	c.Client, err = twofactorverify.NewClient(c.Config)
-
+	client, err := twofactorverify.NewClient(c.Config)
 	if err != nil {
-		ctxlog.WithError(err).Error("twofactorverify: execute: OTP verification failed")
 		return err
 	}
 
-	// Create timeout context
-	// TODO: make timeout configurable
-	const ctxTimeout = 30
-	timeoutCtx, cancelTimeout := context.WithTimeout(ctx, ctxTimeout*time.Second)
-	defer cancelTimeout()
+	ctx, cancel := context.WithTimeout(ctx, timeout)
+	defer cancel()
+
+	fmt.Fprint(c.ReadWriter.Out, prompt)
 
-	// Background push notification with timeout
-	pushauth := make(chan Result)
+	resultCh := make(chan string)
 	go func() {
-		defer close(pushauth)
-		ctxlog.Info("twofactorverify: execute: waiting for push auth")
-		status, success, err := c.pushAuth(timeoutCtx)
-		ctxlog.WithError(err).Info("twofactorverify: execute: push auth verified")
-
-		select {
-		case <-timeoutCtx.Done(): // push cancelled by manual OTP
-			// skip writing to channel
-			ctxlog.Info("twofactorverify: execute: push auth cancelled")
-		default:
-			pushauth <- Result{Error: err, Status: status, Success: success}
+		err := client.PushAuth(ctx, c.Args)
+		if err == nil {
+			resultCh <- "OTP has been validated by Push Authentication. Git operations are now allowed."
 		}
 	}()
 
-	// Also allow manual OTP entry while waiting for push, with same timeout as push
-	verify := make(chan Result)
 	go func() {
-		defer close(verify)
-		ctxlog.Info("twofactorverify: execute: waiting for user input")
-		answer := ""
-		answer = c.getOTP(timeoutCtx)
-		ctxlog.Info("twofactorverify: execute: user input received")
-
-		select {
-		case <-timeoutCtx.Done(): // manual OTP cancelled by push
-			// skip writing to channel
-			ctxlog.Info("twofactorverify: execute: verify cancelled")
-		default:
-			ctxlog.Info("twofactorverify: execute: verifying entered OTP")
-			status, success, err := c.verifyOTP(timeoutCtx, answer)
-			ctxlog.WithError(err).Info("twofactorverify: execute: OTP verified")
-			verify <- Result{Error: err, Status: status, Success: success}
+		answer, err := c.getOTP(ctx)
+		if err != nil {
+			resultCh <- formatErr(err)
 		}
-	}()
 
-	for {
-		select {
-		case res := <-verify:
-			if res.Status == "" {
-				// channel closed; don't print anything
-			} else {
-				fmt.Fprint(c.ReadWriter.Out, res.Status)
-				return nil
-			}
-		case res := <-pushauth:
-			if res.Status == "" {
-				// channel closed; don't print anything
-			} else {
-				fmt.Fprint(c.ReadWriter.Out, res.Status)
-				return nil
-			}
-		case <-timeoutCtx.Done(): // push timed out
-			fmt.Fprint(c.ReadWriter.Out, "\nOTP verification timed out\n")
-			return nil
+		if err := client.VerifyOTP(ctx, c.Args, answer); err != nil {
+			resultCh <- formatErr(err)
+		} else {
+			resultCh <- "OTP validation successful. Git operations are now allowed."
 		}
+	}()
+
+	var message string
+	select {
+	case message = <-resultCh:
+	case <-ctx.Done():
+		message = formatErr(ctx.Err())
 	}
 
+	log.WithContextFields(ctx, log.Fields{"message": message}).Info("Two factor verify command finished")
+	fmt.Fprintf(c.ReadWriter.Out, "\n%v\n", message)
+
 	return nil
 }
 
-func (c *Command) getOTP(ctx context.Context) string {
-	prompt := "OTP: "
-	fmt.Fprint(c.ReadWriter.Out, prompt)
-
+func (c *Command) getOTP(ctx context.Context) (string, error) {
 	var answer string
 	otpLength := int64(64)
 	reader := io.LimitReader(c.ReadWriter.In, otpLength)
@@ -120,41 +78,13 @@ func (c *Command) getOTP(ctx context.Context) string {
 		log.ContextLogger(ctx).WithError(err).Debug("twofactorverify: getOTP: Failed to get user input")
 	}
 
-	return answer
-}
-
-func (c *Command) verifyOTP(ctx context.Context, otp string) (status string, success bool, err error) {
-	reason := ""
-
-	success, reason, err = c.Client.VerifyOTP(ctx, c.Args, otp)
-	if success {
-		status = fmt.Sprintf("\nOTP validation successful. Git operations are now allowed.\n")
-	} else {
-		if err != nil {
-			status = fmt.Sprintf("\nOTP validation failed.\n%v\n", err)
-		} else {
-			status = fmt.Sprintf("\nOTP validation failed.\n%v\n", reason)
-		}
+	if answer == "" {
+		return "", fmt.Errorf("OTP cannot be blank.")
 	}
 
-	err = nil
-
-	return
+	return answer, nil
 }
 
-func (c *Command) pushAuth(ctx context.Context) (status string, success bool, err error) {
-	reason := ""
-
-	success, reason, err = c.Client.PushAuth(ctx, c.Args)
-	if success {
-		status = fmt.Sprintf("\nPush OTP validation successful. Git operations are now allowed.\n")
-	} else {
-		if err != nil {
-			status = fmt.Sprintf("\nPush OTP validation failed.\n%v\n", err)
-		} else {
-			status = fmt.Sprintf("\nPush OTP validation failed.\n%v\n", reason)
-		}
-	}
-
-	return
+func formatErr(err error) string {
+	return fmt.Sprintf("OTP validation failed: %v", err)
 }
diff --git a/internal/command/twofactorverify/twofactorverify_test.go b/internal/command/twofactorverify/twofactorverify_test.go
new file mode 100644
index 0000000..fcc095f
--- /dev/null
+++ b/internal/command/twofactorverify/twofactorverify_test.go
@@ -0,0 +1,191 @@
+package twofactorverify
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"io"
+	"net/http"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"gitlab.com/gitlab-org/gitlab-shell/v14/client/testserver"
+	"gitlab.com/gitlab-org/gitlab-shell/v14/internal/command/commandargs"
+	"gitlab.com/gitlab-org/gitlab-shell/v14/internal/command/readwriter"
+	"gitlab.com/gitlab-org/gitlab-shell/v14/internal/config"
+	"gitlab.com/gitlab-org/gitlab-shell/v14/internal/gitlabnet/twofactorverify"
+)
+
+type blockingReader struct{}
+
+func (*blockingReader) Read([]byte) (int, error) {
+	waitInfinitely := make(chan struct{})
+	<-waitInfinitely
+
+	return 0, nil
+}
+
+func setup(t *testing.T) []testserver.TestRequestHandler {
+	waitInfinitely := make(chan struct{})
+	requests := []testserver.TestRequestHandler{
+		{
+			Path: "/api/v4/internal/two_factor_manual_otp_check",
+			Handler: func(w http.ResponseWriter, r *http.Request) {
+				b, err := io.ReadAll(r.Body)
+				defer r.Body.Close()
+
+				require.NoError(t, err)
+
+				var requestBody *twofactorverify.RequestBody
+				require.NoError(t, json.Unmarshal(b, &requestBody))
+
+				switch requestBody.KeyId {
+				case "verify_via_otp", "verify_via_otp_with_push_error":
+					body := map[string]interface{}{
+						"success": true,
+					}
+					json.NewEncoder(w).Encode(body)
+				case "wait_infinitely":
+					<-waitInfinitely
+				case "error":
+					body := map[string]interface{}{
+						"success": false,
+						"message": "error message",
+					}
+					require.NoError(t, json.NewEncoder(w).Encode(body))
+				case "broken":
+					w.WriteHeader(http.StatusInternalServerError)
+				}
+			},
+		},
+		{
+			Path: "/api/v4/internal/two_factor_push_otp_check",
+			Handler: func(w http.ResponseWriter, r *http.Request) {
+				b, err := io.ReadAll(r.Body)
+				defer r.Body.Close()
+
+				require.NoError(t, err)
+
+				var requestBody *twofactorverify.RequestBody
+				require.NoError(t, json.Unmarshal(b, &requestBody))
+
+				switch requestBody.KeyId {
+				case "verify_via_push":
+					body := map[string]interface{}{
+						"success": true,
+					}
+					json.NewEncoder(w).Encode(body)
+				case "verify_via_otp_with_push_error":
+					w.WriteHeader(http.StatusInternalServerError)
+				default:
+					<-waitInfinitely
+				}
+			},
+		},
+	}
+
+	return requests
+}
+
+const errorHeader = "OTP validation failed: "
+
+func TestExecute(t *testing.T) {
+	requests := setup(t)
+
+	url := testserver.StartSocketHttpServer(t, requests)
+
+	testCases := []struct {
+		desc           string
+		arguments      *commandargs.Shell
+		input          io.Reader
+		expectedOutput string
+	}{
+		{
+			desc:           "Verify via OTP",
+			arguments:      &commandargs.Shell{GitlabKeyId: "verify_via_otp"},
+			expectedOutput: "OTP validation successful. Git operations are now allowed.\n",
+		},
+		{
+			desc:           "Verify via OTP",
+			arguments:      &commandargs.Shell{GitlabKeyId: "verify_via_otp_with_push_error"},
+			expectedOutput: "OTP validation successful. Git operations are now allowed.\n",
+		},
+		{
+			desc:           "Verify via push authentication",
+			arguments:      &commandargs.Shell{GitlabKeyId: "verify_via_push"},
+			input:          &blockingReader{},
+			expectedOutput: "OTP has been validated by Push Authentication. Git operations are now allowed.\n",
+		},
+		{
+			desc:           "With an empty OTP",
+			arguments:      &commandargs.Shell{GitlabKeyId: "verify_via_otp"},
+			input:          bytes.NewBufferString("\n"),
+			expectedOutput: errorHeader + "OTP cannot be blank.\n",
+		},
+		{
+			desc:           "With bad response",
+			arguments:      &commandargs.Shell{GitlabKeyId: "-1"},
+			expectedOutput: errorHeader + "Parsing failed\n",
+		},
+		{
+			desc:           "With API returns an error",
+			arguments:      &commandargs.Shell{GitlabKeyId: "error"},
+			expectedOutput: errorHeader + "error message\n",
+		},
+		{
+			desc:           "With API fails",
+			arguments:      &commandargs.Shell{GitlabKeyId: "broken"},
+			expectedOutput: errorHeader + "Internal API error (500)\n",
+		},
+		{
+			desc:           "With missing arguments",
+			arguments:      &commandargs.Shell{},
+			expectedOutput: errorHeader + "who='' is invalid\n",
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.desc, func(t *testing.T) {
+			output := &bytes.Buffer{}
+
+			input := tc.input
+			if input == nil {
+				input = bytes.NewBufferString("123456\n")
+			}
+
+			cmd := &Command{
+				Config:     &config.Config{GitlabUrl: url},
+				Args:       tc.arguments,
+				ReadWriter: &readwriter.ReadWriter{Out: output, In: input},
+			}
+
+			err := cmd.Execute(context.Background())
+
+			require.NoError(t, err)
+			require.Equal(t, prompt+"\n"+tc.expectedOutput, output.String())
+		})
+	}
+}
+
+func TestCanceledContext(t *testing.T) {
+	requests := setup(t)
+
+	output := &bytes.Buffer{}
+
+	url := testserver.StartSocketHttpServer(t, requests)
+	cmd := &Command{
+		Config:     &config.Config{GitlabUrl: url},
+		Args:       &commandargs.Shell{GitlabKeyId: "wait_infinitely"},
+		ReadWriter: &readwriter.ReadWriter{Out: output, In: &bytes.Buffer{}},
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+
+	errCh := make(chan error)
+	go func() { errCh <- cmd.Execute(ctx) }()
+	cancel()
+
+	require.NoError(t, <-errCh)
+	require.Equal(t, prompt+"\n"+errorHeader+"context canceled\n", output.String())
+}
diff --git a/internal/command/twofactorverify/twofactorverifymanual_test.go b/internal/command/twofactorverify/twofactorverifymanual_test.go
deleted file mode 100644
index 59db281..0000000
--- a/internal/command/twofactorverify/twofactorverifymanual_test.go
+++ /dev/null
@@ -1,154 +0,0 @@
-package twofactorverify
-
-import (
-	"bytes"
-	"context"
-	"encoding/json"
-	"io"
-	"net/http"
-	"testing"
-	"time"
-
-	"github.com/stretchr/testify/require"
-
-	"gitlab.com/gitlab-org/gitlab-shell/v14/client/testserver"
-	"gitlab.com/gitlab-org/gitlab-shell/v14/internal/command/commandargs"
-	"gitlab.com/gitlab-org/gitlab-shell/v14/internal/command/readwriter"
-	"gitlab.com/gitlab-org/gitlab-shell/v14/internal/config"
-	"gitlab.com/gitlab-org/gitlab-shell/v14/internal/gitlabnet/twofactorverify"
-)
-
-func setupManual(t *testing.T) []testserver.TestRequestHandler {
-	requests := []testserver.TestRequestHandler{
-		{
-			Path: "/api/v4/internal/two_factor_manual_otp_check",
-			Handler: func(w http.ResponseWriter, r *http.Request) {
-				b, err := io.ReadAll(r.Body)
-				defer r.Body.Close()
-
-				require.NoError(t, err)
-
-				var requestBody *twofactorverify.RequestBody
-				require.NoError(t, json.Unmarshal(b, &requestBody))
-
-				var body map[string]interface{}
-				switch requestBody.KeyId {
-				case "1":
-					body = map[string]interface{}{
-						"success": true,
-					}
-					json.NewEncoder(w).Encode(body)
-				case "error":
-					body = map[string]interface{}{
-						"success": false,
-						"message": "error message",
-					}
-					require.NoError(t, json.NewEncoder(w).Encode(body))
-				case "broken":
-					w.WriteHeader(http.StatusInternalServerError)
-				}
-			},
-		},
-		{
-			Path: "/api/v4/internal/two_factor_push_otp_check",
-			Handler: func(w http.ResponseWriter, r *http.Request) {
-				b, err := io.ReadAll(r.Body)
-				defer r.Body.Close()
-
-				require.NoError(t, err)
-
-				var requestBody *twofactorverify.RequestBody
-				require.NoError(t, json.Unmarshal(b, &requestBody))
-
-				time.Sleep(5 * time.Second)
-
-				var body map[string]interface{}
-				switch requestBody.KeyId {
-				case "1":
-					body = map[string]interface{}{
-						"success": true,
-					}
-					json.NewEncoder(w).Encode(body)
-				case "error":
-					body = map[string]interface{}{
-						"success": false,
-						"message": "error message",
-					}
-					require.NoError(t, json.NewEncoder(w).Encode(body))
-				case "broken":
-					w.WriteHeader(http.StatusInternalServerError)
-				default:
-					body = map[string]interface{}{
-						"success": true,
-						"message": "default message",
-					}
-					json.NewEncoder(w).Encode(body)
-				}
-			},
-		},
-	}
-
-	return requests
-}
-
-const (
-	manualQuestion    = "OTP: \n"
-	manualErrorHeader = "OTP validation failed.\n"
-)
-
-func TestExecuteManual(t *testing.T) {
-	requests := setupManual(t)
-
-	url := testserver.StartSocketHttpServer(t, requests)
-
-	testCases := []struct {
-		desc           string
-		arguments      *commandargs.Shell
-		answer         string
-		expectedOutput string
-	}{
-		{
-			desc:           "With a known key id",
-			arguments:      &commandargs.Shell{GitlabKeyId: "1"},
-			answer:         "123456\n",
-			expectedOutput: manualQuestion + "OTP validation successful. Git operations are now allowed.\n",
-		},
-		{
-			desc:           "With bad response",
-			arguments:      &commandargs.Shell{GitlabKeyId: "-1"},
-			answer:         "123456\n",
-			expectedOutput: manualQuestion + manualErrorHeader + "Parsing failed\n",
-		},
-		{
-			desc:           "With API returns an error",
-			arguments:      &commandargs.Shell{GitlabKeyId: "error"},
-			answer:         "yes\n",
-			expectedOutput: manualQuestion + manualErrorHeader + "error message\n",
-		},
-		{
-			desc:           "With API fails",
-			arguments:      &commandargs.Shell{GitlabKeyId: "broken"},
-			answer:         "yes\n",
-			expectedOutput: manualQuestion + manualErrorHeader + "Internal API error (500)\n",
-		},
-	}
-
-	for _, tc := range testCases {
-		t.Run(tc.desc, func(t *testing.T) {
-			output := &bytes.Buffer{}
-
-			input := bytes.NewBufferString(tc.answer)
-
-			cmd := &Command{
-				Config:     &config.Config{GitlabUrl: url},
-				Args:       tc.arguments,
-				ReadWriter: &readwriter.ReadWriter{Out: output, In: input},
-			}
-
-			err := cmd.Execute(context.Background())
-
-			require.NoError(t, err)
-			require.Equal(t, tc.expectedOutput, output.String())
-		})
-	}
-}
diff --git a/internal/command/twofactorverify/twofactorverifypush_test.go b/internal/command/twofactorverify/twofactorverifypush_test.go
deleted file mode 100644
index 4b58cc5..0000000
--- a/internal/command/twofactorverify/twofactorverifypush_test.go
+++ /dev/null
@@ -1,144 +0,0 @@
-package twofactorverify
-
-import (
-	"bytes"
-	"context"
-	"encoding/json"
-	"io"
-	"net/http"
-	"testing"
-	"time"
-
-	"github.com/stretchr/testify/require"
-
-	"gitlab.com/gitlab-org/gitlab-shell/v14/client/testserver"
-	"gitlab.com/gitlab-org/gitlab-shell/v14/internal/command/commandargs"
-	"gitlab.com/gitlab-org/gitlab-shell/v14/internal/command/readwriter"
-	"gitlab.com/gitlab-org/gitlab-shell/v14/internal/config"
-	"gitlab.com/gitlab-org/gitlab-shell/v14/internal/gitlabnet/twofactorverify"
-)
-
-func setupPush(t *testing.T) []testserver.TestRequestHandler {
-	requests := []testserver.TestRequestHandler{
-		{
-			Path: "/api/v4/internal/two_factor_push_otp_check",
-			Handler: func(w http.ResponseWriter, r *http.Request) {
-				b, err := io.ReadAll(r.Body)
-				defer r.Body.Close()
-
-				require.NoError(t, err)
-
-				var requestBody *twofactorverify.RequestBody
-				require.NoError(t, json.Unmarshal(b, &requestBody))
-
-				var body map[string]interface{}
-				switch requestBody.KeyId {
-				case "1":
-					body = map[string]interface{}{
-						"success": true,
-					}
-					json.NewEncoder(w).Encode(body)
-				case "error":
-					body = map[string]interface{}{
-						"success": false,
-						"message": "error message",
-					}
-					require.NoError(t, json.NewEncoder(w).Encode(body))
-				case "broken":
-					w.WriteHeader(http.StatusInternalServerError)
-				}
-			},
-		},
-		{
-			Path: "/api/v4/internal/two_factor_manual_otp_check",
-			Handler: func(w http.ResponseWriter, r *http.Request) {
-				b, err := io.ReadAll(r.Body)
-				defer r.Body.Close()
-
-				require.NoError(t, err)
-
-				var requestBody *twofactorverify.RequestBody
-				require.NoError(t, json.Unmarshal(b, &requestBody))
-
-				time.Sleep(20 * time.Second)
-
-				var body map[string]interface{}
-				switch requestBody.KeyId {
-				case "1":
-					body = map[string]interface{}{
-						"success": true,
-					}
-					json.NewEncoder(w).Encode(body)
-				case "error":
-					body = map[string]interface{}{
-						"success": false,
-						"message": "error message",
-					}
-					require.NoError(t, json.NewEncoder(w).Encode(body))
-				case "broken":
-					w.WriteHeader(http.StatusInternalServerError)
-				}
-			},
-		},
-	}
-
-	return requests
-}
-
-const (
-	pushQuestion    = "OTP: \n"
-	pushErrorHeader = "Push OTP validation failed.\n"
-)
-
-func TestExecutePush(t *testing.T) {
-	requests := setupPush(t)
-
-	url := testserver.StartSocketHttpServer(t, requests)
-
-	testCases := []struct {
-		desc           string
-		arguments      *commandargs.Shell
-		answer         string
-		expectedOutput string
-	}{
-		{
-			desc:           "When push is provided",
-			arguments:      &commandargs.Shell{GitlabKeyId: "1"},
-			answer:         "",
-			expectedOutput: pushQuestion + "Push OTP validation successful. Git operations are now allowed.\n",
-		},
-		{
-			desc:           "With API returns an error",
-			arguments:      &commandargs.Shell{GitlabKeyId: "error"},
-			answer:         "",
-			expectedOutput: pushQuestion + pushErrorHeader + "error message\n",
-		},
-		{
-			desc:           "With API fails",
-			arguments:      &commandargs.Shell{GitlabKeyId: "broken"},
-			answer:         "",
-			expectedOutput: pushQuestion + pushErrorHeader + "Internal API error (500)\n",
-		},
-	}
-
-	for _, tc := range testCases {
-		t.Run(tc.desc, func(t *testing.T) {
-			output := &bytes.Buffer{}
-
-			var input io.Reader
-			// make input wait for push auth tests
-			input, _ = io.Pipe()
-
-			cmd := &Command{
-				Config:     &config.Config{GitlabUrl: url},
-				Args:       tc.arguments,
-				ReadWriter: &readwriter.ReadWriter{Out: output, In: input},
-			}
-
-			err := cmd.Execute(context.Background())
-
-			require.NoError(t, err)
-			require.Equal(t, tc.expectedOutput, output.String())
-		})
-	}
-}