Handle unbuffered channels per review comment

1 files changed, 37 insertions, 36 deletions

diff --git a/internal/command/twofactorverify/twofactorverify.go b/internal/command/twofactorverify/twofactorverify.go
index 9f11895..7877745 100644
--- a/internal/command/twofactorverify/twofactorverify.go
+++ b/internal/command/twofactorverify/twofactorverify.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"io"
+	"time"
 
 	"gitlab.com/gitlab-org/gitlab-shell/internal/command/commandargs"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/command/readwriter"
@@ -33,41 +34,41 @@ func (c *Command) Execute(ctx context.Context) error {
 		return err
 	}
 
-	verify := make(chan Result)
-	pushauth := make(chan Result)
+	// Create timeout context 
+	// TODO: make timeout configurable
+	const ctxTimeout = 30
+	timeoutCtx, cancel := context.WithTimeout(ctx, ctxTimeout * time.Second)
+	defer cancel()
 
+	// Background push notification with timeout
+	pushauth := make(chan Result)
 	go func() {
-		status, success, err := c.verifyOTP(ctx, c.getOTP())
-		verify <- Result{Error: err, Status: status, Success: success}
+		defer close(pushauth)
+		status, success, err := c.pushAuth(timeoutCtx)
+		pushauth <- Result{Error: err, Status: status, Success: success}
 	}()
 
+	// Also allow manual OTP entry while waiting for push, with same timeout as push
+	verify := make(chan Result)
 	go func() {
-		status, success, err := c.pushAuth(ctx)
-		pushauth <- Result{Error: err, Status: status, Success: success}
+		defer close(verify)
+		status, success, err := c.verifyOTP(timeoutCtx, c.getOTP(timeoutCtx))
+		verify <- Result{Error: err, Status: status, Success: success}
 	}()
 
-L:
-	for {
-		select {
-		case res := <-verify:
-			if res.Error != nil {
-				return res.Error
-			}
-			fmt.Fprint(c.ReadWriter.Out, res.Status)
-			break L
-		case res := <-pushauth:
-			if res.Success {
-				fmt.Fprint(c.ReadWriter.Out, res.Status)
-				break L
-			} else {
-				// ignore reject from remote, need to wait for user input in this case
-			}
-		}
+	select {
+	case res := <-verify: // manual OTP
+		fmt.Fprint(c.ReadWriter.Out, res.Status)
+		case res := <-pushauth: // push
+		fmt.Fprint(c.ReadWriter.Out, res.Status)
+	case <-timeoutCtx.Done(): // push timed out
+		fmt.Fprint(c.ReadWriter.Out, "OTP verification timed out")
 	}
+
 	return nil
 }
 
-func (c *Command) getOTP() string {
+func (c *Command) getOTP(ctx context.Context) string {
 	prompt := "OTP: "
 	fmt.Fprint(c.ReadWriter.Out, prompt)
 
@@ -79,38 +80,38 @@ func (c *Command) getOTP() string {
 	return answer
 }
 
-func (c *Command) pushAuth(ctx context.Context) (status string, success bool, err error) {
+func (c *Command) verifyOTP(ctx context.Context, otp string) (status string, success bool, err error) {
 	reason := ""
 
-	success, reason, err = c.Client.PushAuth(ctx, c.Args)
+	success, reason, err = c.Client.VerifyOTP(ctx, c.Args, otp)
 	if success {
-		status = fmt.Sprintf("\nPush OTP validation successful. Git operations are now allowed.\n")
+		status = fmt.Sprintf("\nOTP validation successful. Git operations are now allowed.\n")
 	} else {
 		if err != nil {
-			status = fmt.Sprintf("\nPush OTP validation failed.\n%v\n", err)
+			status = fmt.Sprintf("\nOTP validation failed.\n%v\n", err)
 		} else {
-			status = fmt.Sprintf("\nPush OTP validation failed.\n%v\n", reason)
+			status = fmt.Sprintf("\nOTP validation failed.\n%v\n", reason)
 		}
 	}
 
+	err = nil
+
 	return
 }
 
-func (c *Command) verifyOTP(ctx context.Context, otp string) (status string, success bool, err error) {
+func (c *Command) pushAuth(ctx context.Context) (status string, success bool, err error) {
 	reason := ""
 
-	success, reason, err = c.Client.VerifyOTP(ctx, c.Args, otp)
+	success, reason, err = c.Client.PushAuth(ctx, c.Args)
 	if success {
-		status = fmt.Sprintf("\nOTP validation successful. Git operations are now allowed.\n")
+		status = fmt.Sprintf("\nPush OTP validation successful. Git operations are now allowed.\n")
 	} else {
 		if err != nil {
-			status = fmt.Sprintf("\nOTP validation failed.\n%v\n", err)
+			status = fmt.Sprintf("\nPush OTP validation failed.\n%v\n", err)
 		} else {
-			status = fmt.Sprintf("\nOTP validation failed.\n%v\n", reason)
+			status = fmt.Sprintf("\nPush OTP validation failed.\n%v\n", reason)
 		}
 	}
 
-	err = nil
-
 	return
 }