coder · Emyrk · Jan 25, 2024 · Jan 22, 2024 · Jan 22, 2024 · Jan 23, 2024
diff --git a/coderd/database/modelmethods.go b/coderd/database/modelmethods.go
@@ -6,6 +6,7 @@ import (
 	"time"
 
 	"golang.org/x/exp/maps"
+	"golang.org/x/oauth2"
 
 	"github.com/coder/coder/v2/coderd/database/dbtime"
 	"github.com/coder/coder/v2/coderd/rbac"
@@ -268,6 +269,14 @@ func (u ExternalAuthLink) RBACObject() rbac.Object {
 	return rbac.ResourceUserData.WithID(u.UserID).WithOwner(u.UserID.String())
 }
 
+func (u ExternalAuthLink) OAuthToken() *oauth2.Token {
+	return &oauth2.Token{
+		AccessToken:  u.OAuthAccessToken,
+		RefreshToken: u.OAuthRefreshToken,
+		Expiry:       u.OAuthExpiry,
+	}
+}
+
 func (u UserLink) RBACObject() rbac.Object {
 	// I assume UserData is ok?
 	return rbac.ResourceUserData.WithOwner(u.UserID.String()).WithID(u.UserID)

diff --git a/coderd/externalauth.go b/coderd/externalauth.go
@@ -57,7 +57,7 @@ func (api *API) externalAuthByID(w http.ResponseWriter, r *http.Request) {
 	}
 	var eg errgroup.Group
 	eg.Go(func() (err error) {
-		res.Authenticated, res.User, err = config.ValidateToken(ctx, link.OAuthAccessToken)
+		res.Authenticated, res.User, err = config.ValidateToken(ctx, link.OAuthToken())
 		return err
 	})
 	eg.Go(func() (err error) {

diff --git a/coderd/externalauth/externalauth.go b/coderd/externalauth/externalauth.go
@@ -138,7 +138,7 @@ func (c *Config) RefreshToken(ctx context.Context, db database.Store, externalAu
 	retryCtx, retryCtxCancel := context.WithTimeout(ctx, time.Second)
 	defer retryCtxCancel()
 validate:
-	valid, _, err := c.ValidateToken(ctx, token.AccessToken)
+	valid, _, err := c.ValidateToken(ctx, token)
 	if err != nil {
 		return externalAuthLink, false, xerrors.Errorf("validate external auth token: %w", err)
 	}
@@ -179,7 +179,14 @@ validate:
 
 // ValidateToken ensures the Git token provided is valid!
 // The user is optionally returned if the provider supports it.
-func (c *Config) ValidateToken(ctx context.Context, token string) (bool, *codersdk.ExternalAuthUser, error) {
+func (c *Config) ValidateToken(ctx context.Context, link *oauth2.Token) (bool, *codersdk.ExternalAuthUser, error) {
+	if link == nil {
+		return false, nil, xerrors.New("validate external auth token: token is nil")
+	}
+	if !link.Expiry.IsZero() && link.Expiry.Before(dbtime.Now()) {
+		return false, nil, nil
+	}
+
 	if c.ValidateURL == "" {
 		// Default that the token is valid if no validation URL is provided.
 		return true, nil, nil
@@ -189,7 +196,7 @@ func (c *Config) ValidateToken(ctx context.Context, token string) (bool, *coders
 		return false, nil, err
 	}
 
-	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token))
+	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", link.AccessToken))
 	res, err := c.InstrumentedOAuth2Config.Do(ctx, promoauth.SourceValidateToken, req)
 	if err != nil {
 		return false, nil, err
@@ -396,10 +403,15 @@ func (c *DeviceAuth) ExchangeDeviceCode(ctx context.Context, deviceCode string)
 	if body.Error != "" {
 		return nil, xerrors.New(body.Error)
 	}
+	// If expiresIn is 0, then the token never expires.
+	expires := dbtime.Now().Add(time.Duration(body.ExpiresIn) * time.Second)
+	if body.ExpiresIn == 0 {
+		expires = time.Time{}
+	}
 	return &oauth2.Token{
 		AccessToken:  body.AccessToken,
 		RefreshToken: body.RefreshToken,
-		Expiry:       dbtime.Now().Add(time.Duration(body.ExpiresIn) * time.Second),
+		Expiry:       expires,
 	}, nil
 }
 

diff --git a/coderd/promoauth/oauth2_test.go b/coderd/promoauth/oauth2_test.go
@@ -75,7 +75,7 @@ func TestInstrument(t *testing.T) {
 	require.Equal(t, count("TokenSource"), 1)
 
 	// Try a validate
-	valid, _, err := cfg.ValidateToken(ctx, refreshed.AccessToken)
+	valid, _, err := cfg.ValidateToken(ctx, refreshed)
 	require.NoError(t, err)
 	require.True(t, valid)
 	require.Equal(t, count("ValidateToken"), 1)

diff --git a/coderd/workspaceagents.go b/coderd/workspaceagents.go
@@ -2031,78 +2031,26 @@ func (api *API) workspaceAgentsExternalAuth(rw http.ResponseWriter, r *http.Requ
 		return
 	}
 
-	if listen {
-		// Since we're ticking frequently and this sign-in operation is rare,
-		// we are OK with polling to avoid the complexity of pubsub.
-		ticker, done := api.NewTicker(time.Second)
-		defer done()
-		var previousToken database.ExternalAuthLink
-		for {
-			select {
-			case <-ctx.Done():
-				return
-			case <-ticker:
-			}
-			externalAuthLink, err := api.Database.GetExternalAuthLink(ctx, database.GetExternalAuthLinkParams{
-				ProviderID: externalAuthConfig.ID,
-				UserID:     workspace.OwnerID,
-			})
-			if err != nil {
-				if errors.Is(err, sql.ErrNoRows) {
-					continue
-				}
-				httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
-					Message: "Failed to get external auth link.",
-					Detail:  err.Error(),
-				})
-				return
-			}
-
-			// Expiry may be unset if the application doesn't configure tokens
-			// to expire.
-			// See
-			// https://docs.github.com/en/apps/creating-github-apps/authenticating-with-a-github-app/generating-a-user-access-token-for-a-github-app.
-			if externalAuthLink.OAuthExpiry.Before(dbtime.Now()) && !externalAuthLink.OAuthExpiry.IsZero() {
-				continue
-			}
-
-			// Only attempt to revalidate an oauth token if it has actually changed.
-			// No point in trying to validate the same token over and over again.
-			if previousToken.OAuthAccessToken == externalAuthLink.OAuthAccessToken &&
-				previousToken.OAuthRefreshToken == externalAuthLink.OAuthRefreshToken &&
-				previousToken.OAuthExpiry == externalAuthLink.OAuthExpiry {
-				continue
-			}
-
-			valid, _, err := externalAuthConfig.ValidateToken(ctx, externalAuthLink.OAuthAccessToken)
-			if err != nil {
-				api.Logger.Warn(ctx, "failed to validate external auth token",
-					slog.F("workspace_owner_id", workspace.OwnerID.String()),
-					slog.F("validate_url", externalAuthConfig.ValidateURL),
-					slog.Error(err),
-				)
-			}
-			previousToken = externalAuthLink
-			if !valid {
-				continue
-			}
-			resp, err := createExternalAuthResponse(externalAuthConfig.Type, externalAuthLink.OAuthAccessToken, externalAuthLink.OAuthExtra)
-			if err != nil {
-				httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
-					Message: "Failed to create external auth response.",
-					Detail:  err.Error(),
-				})
-				return
-			}
-			httpapi.Write(ctx, rw, http.StatusOK, resp)
+	var previousToken *database.ExternalAuthLink
+	// handleRetrying will attempt to continually check for a new token
+	// if listen is true. This is useful if an error is encountered in the
+	// original single flow.
+	//
+	// By default, if no errors are encountered, then the single flow response
+	// is returned.
+	handleRetrying := func(code int, response any) {
+		if !listen {
+			httpapi.Write(ctx, rw, code, response)
 			return
 		}
+
+		api.workspaceAgentsExternalAuthListen(ctx, rw, previousToken, externalAuthConfig, workspace)
 	}
 
 	// This is the URL that will redirect the user with a state token.
 	redirectURL, err := api.AccessURL.Parse(fmt.Sprintf("/external-auth/%s", externalAuthConfig.ID))
 	if err != nil {
-		httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
+		handleRetrying(http.StatusInternalServerError, codersdk.Response{
 			Message: "Failed to parse access URL.",
 			Detail:  err.Error(),
 		})
@@ -2115,36 +2063,40 @@ func (api *API) workspaceAgentsExternalAuth(rw http.ResponseWriter, r *http.Requ
 	})
 	if err != nil {
 		if !errors.Is(err, sql.ErrNoRows) {
-			httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
+			handleRetrying(http.StatusInternalServerError, codersdk.Response{
 				Message: "Failed to get external auth link.",
 				Detail:  err.Error(),
 			})
 			return
 		}
 
-		httpapi.Write(ctx, rw, http.StatusOK, agentsdk.ExternalAuthResponse{
+		handleRetrying(http.StatusOK, agentsdk.ExternalAuthResponse{
 			URL: redirectURL.String(),
 		})
 		return
 	}
 
-	externalAuthLink, updated, err := externalAuthConfig.RefreshToken(ctx, api.Database, externalAuthLink)
+	externalAuthLink, valid, err := externalAuthConfig.RefreshToken(ctx, api.Database, externalAuthLink)
 	if err != nil {
-		httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
+		handleRetrying(http.StatusInternalServerError, codersdk.Response{
 			Message: "Failed to refresh external auth token.",
 			Detail:  err.Error(),
 		})
 		return
 	}
-	if !updated {
-		httpapi.Write(ctx, rw, http.StatusOK, agentsdk.ExternalAuthResponse{
+	if !valid {
+		// Set the previous token so the retry logic will skip validating the
+		// same token again. This should only be set if the token is invalid and there
+		// was no error. If it is invalid because of an error, then we should recheck.
+		previousToken = &externalAuthLink
+		handleRetrying(http.StatusOK, agentsdk.ExternalAuthResponse{
 			URL: redirectURL.String(),
 		})
 		return
 	}
 	resp, err := createExternalAuthResponse(externalAuthConfig.Type, externalAuthLink.OAuthAccessToken, externalAuthLink.OAuthExtra)
 	if err != nil {
-		httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
+		handleRetrying(http.StatusInternalServerError, codersdk.Response{
 			Message: "Failed to create external auth response.",
 			Detail:  err.Error(),
 		})
@@ -2153,6 +2105,81 @@ func (api *API) workspaceAgentsExternalAuth(rw http.ResponseWriter, r *http.Requ
 	httpapi.Write(ctx, rw, http.StatusOK, resp)
 }
 
+func (api *API) workspaceAgentsExternalAuthListen(ctx context.Context, rw http.ResponseWriter, previous *database.ExternalAuthLink, externalAuthConfig *externalauth.Config, workspace database.Workspace) {
+	// Since we're ticking frequently and this sign-in operation is rare,
+	// we are OK with polling to avoid the complexity of pubsub.
+	ticker, done := api.NewTicker(time.Second)
+	defer done()
+	// If we have a previous token that is invalid, we should not check this again.
+	// This serves to prevent doing excessive unauthorized requests to the external
+	// auth provider. For github, this limit is 60 per hour, so saving a call
+	// per invalid token can be significant.
+	var previousToken database.ExternalAuthLink
+	if previous != nil {
+		previousToken = *previous
+	}
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-ticker:
+		}
+		externalAuthLink, err := api.Database.GetExternalAuthLink(ctx, database.GetExternalAuthLinkParams{
+			ProviderID: externalAuthConfig.ID,
+			UserID:     workspace.OwnerID,
+		})
+		if err != nil {
+			if errors.Is(err, sql.ErrNoRows) {
+				continue
+			}
+			httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
+				Message: "Failed to get external auth link.",
+				Detail:  err.Error(),
+			})
+			return
+		}
+
+		// Expiry may be unset if the application doesn't configure tokens
+		// to expire.
+		// See
+		// https://docs.github.com/en/apps/creating-github-apps/authenticating-with-a-github-app/generating-a-user-access-token-for-a-github-app.
+		if externalAuthLink.OAuthExpiry.Before(dbtime.Now()) && !externalAuthLink.OAuthExpiry.IsZero() {
+			continue
+		}
+
+		// Only attempt to revalidate an oauth token if it has actually changed.
+		// No point in trying to validate the same token over and over again.
+		if previousToken.OAuthAccessToken == externalAuthLink.OAuthAccessToken &&
+			previousToken.OAuthRefreshToken == externalAuthLink.OAuthRefreshToken &&
+			previousToken.OAuthExpiry == externalAuthLink.OAuthExpiry {
+			continue
+		}
+
+		valid, _, err := externalAuthConfig.ValidateToken(ctx, externalAuthLink.OAuthToken())
+		if err != nil {
+			api.Logger.Warn(ctx, "failed to validate external auth token",
+				slog.F("workspace_owner_id", workspace.OwnerID.String()),
+				slog.F("validate_url", externalAuthConfig.ValidateURL),
+				slog.Error(err),
+			)
+		}
+		previousToken = externalAuthLink
+		if !valid {
+			continue
+		}
+		resp, err := createExternalAuthResponse(externalAuthConfig.Type, externalAuthLink.OAuthAccessToken, externalAuthLink.OAuthExtra)
+		if err != nil {
+			httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
+				Message: "Failed to create external auth response.",
+				Detail:  err.Error(),
+			})
+			return
+		}
+		httpapi.Write(ctx, rw, http.StatusOK, resp)
+		return
+	}
+}
+
 // createExternalAuthResponse creates an ExternalAuthResponse based on the
 // provider type. This is to support legacy `/workspaceagents/me/gitauth`
 // which uses `Username` and `Password`.

diff --git a/coderd/workspaceagents_test.go b/coderd/workspaceagents_test.go
@@ -1577,7 +1577,7 @@ func TestWorkspaceAgentExternalAuthListen(t *testing.T) {
 			},
 			ExternalAuthConfigs: []*externalauth.Config{
 				fake.ExternalAuthConfig(t, providerID, nil, func(cfg *externalauth.Config) {
-					cfg.Type = codersdk.EnhancedExternalAuthProviderGitHub.String()
+					cfg.Type = codersdk.EnhancedExternalAuthProviderGitLab.String()
 				}),
 			},
 		})
@@ -1623,7 +1623,8 @@ func TestWorkspaceAgentExternalAuthListen(t *testing.T) {
 			ticks <- time.Now()
 		}
 		cancel()
-		// We expect only 1
+		// We expect only 1. One from the initial "Refresh" attempt, and the
+		// other should be skipped.
 		// In a failed test, you will likely see 9, as the last one
 		// gets canceled.
 		require.Equal(t, 1, validateCalls, "validate calls duplicated on same token")