From 60d04502877d5745e9fab969d541e0020a5d7ff7 Mon Sep 17 00:00:00 2001
From: Laura Brehm <laurabrehm@hey.com>
Date: Wed, 28 Aug 2024 13:09:42 +0100
Subject: [PATCH] oauth/api: drain timer channel on each iteration

Previously, if while polling for oauth device-code login results a user
suspended the process (such as with CTRL-Z) and then restored it with
`fg`, an error might occur in the form of:

```
failed waiting for authentication: You are polling faster than the specified interval of 5 seconds.
```

This is due to our use of a `time.Ticker` here - if no receiver drains
the ticker channel (and timers/tickers use a buffered channel behind the
scenes), more than one tick will pile up, causing the program to "tick"
twice, in fast succession, after it is resumed.

The new implementation replaces the `time.Ticker` with a `time.Timer`
(`time.Ticker` is just a nice wrapper) and introduces a helper function
`resetTimer` to ensure that before every `select`, the timer is stopped
and it's channel is drained.

Signed-off-by: Laura Brehm <laurabrehm@hey.com>
---
 cli/internal/oauth/api/api.go      | 41 +++++++++++++++++++++++++++---
 cli/internal/oauth/api/api_test.go |  2 +-
 2 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/cli/internal/oauth/api/api.go b/cli/internal/oauth/api/api.go
index ddcba749fdeb..cf6c1be961e5 100644
--- a/cli/internal/oauth/api/api.go
+++ b/cli/internal/oauth/api/api.go
@@ -96,18 +96,32 @@ func tryDecodeOAuthError(resp *http.Response) error {
 // authenticated or we have reached the time limit for authenticating (based on
 // the response from GetDeviceCode).
 func (a API) WaitForDeviceToken(ctx context.Context, state State) (TokenResponse, error) {
-	ticker := time.NewTicker(state.IntervalDuration())
+	// Ticker for polling tenant for login – based on the interval
+	// specified by the tenant response.
+	ticker := time.NewTimer(state.IntervalDuration())
 	defer ticker.Stop()
-	timeout := time.After(state.ExpiryDuration())
+	// The tenant tells us for as long as we can poll it for credentials
+	// while the user logs in through their browser. Timeout if we don't get
+	// credentials within this period.
+	timeout := time.NewTimer(state.ExpiryDuration())
+	defer timeout.Stop()
 
 	for {
+		resetTimer(ticker, state.IntervalDuration())
 		select {
 		case <-ctx.Done():
+			// user canceled login
 			return TokenResponse{}, ctx.Err()
 		case <-ticker.C:
+			// tick, check for user login
 			res, err := a.getDeviceToken(ctx, state)
 			if err != nil {
-				return res, err
+				if errors.Is(err, context.Canceled) {
+					// if the caller canceled the context, continue
+					// and let the select hit the ctx.Done() branch
+					continue
+				}
+				return TokenResponse{}, err
 			}
 
 			if res.Error != nil {
@@ -119,14 +133,33 @@ func (a API) WaitForDeviceToken(ctx context.Context, state State) (TokenResponse
 			}
 
 			return res, nil
-		case <-timeout:
+		case <-timeout.C:
+			// login timed out
 			return TokenResponse{}, ErrTimeout
 		}
 	}
 }
 
+// resetTimer is a helper function thatstops, drains and resets the timer.
+// This is necessary in go versions <1.23, since the timer isn't stopped +
+// the timer's channel isn't drained on timer.Reset.
+// See: https://go-review.googlesource.com/c/go/+/568341
+// FIXME: remove/simplify this after we update to go1.23
+func resetTimer(t *time.Timer, d time.Duration) {
+	if !t.Stop() {
+		select {
+		case <-t.C:
+		default:
+		}
+	}
+	t.Reset(d)
+}
+
 // getToken calls the token endpoint of Auth0 and returns the response.
 func (a API) getDeviceToken(ctx context.Context, state State) (TokenResponse, error) {
+	ctx, cancel := context.WithTimeout(ctx, 1*time.Minute)
+	defer cancel()
+
 	data := url.Values{
 		"client_id":   {a.ClientID},
 		"grant_type":  {"urn:ietf:params:oauth:grant-type:device_code"},
diff --git a/cli/internal/oauth/api/api_test.go b/cli/internal/oauth/api/api_test.go
index 34b9600f5cb3..ec8ea876ff77 100644
--- a/cli/internal/oauth/api/api_test.go
+++ b/cli/internal/oauth/api/api_test.go
@@ -198,7 +198,7 @@ func TestWaitForDeviceToken(t *testing.T) {
 		state := State{
 			DeviceCode: "aDeviceCode",
 			UserCode:   "aUserCode",
-			Interval:   1,
+			Interval:   5,
 			ExpiresIn:  1,
 		}