add unit test to verify fix

coder · Emyrk · Aug 8, 2023 · Aug 7, 2023 · Aug 7, 2023 · Aug 7, 2023
commit 6dfc5466a67c00111f600f19a7f6c8f652e059cb
diff --git a/coderd/coderdtest/coderdtest.go b/coderd/coderdtest/coderdtest.go
@@ -1022,9 +1022,31 @@ func NewAWSInstanceIdentity(t *testing.T, instanceID string) (awsidentity.Certif
 type OIDCConfig struct {
 	key    *rsa.PrivateKey
 	issuer string
+	// These are optional
+	refreshToken     string
+	oidcTokenExpires func() time.Time
+	tokenSource      func() (*oauth2.Token, error)
 }
 
-func NewOIDCConfig(t *testing.T, issuer string) *OIDCConfig {
+func WithRefreshToken(token string) func(cfg *OIDCConfig) {
+	return func(cfg *OIDCConfig) {
+		cfg.refreshToken = token
+	}
+}
+
+func WithTokenExpires(expFunc func() time.Time) func(cfg *OIDCConfig) {
+	return func(cfg *OIDCConfig) {
+		cfg.oidcTokenExpires = expFunc
+	}
+}
+
+func WithTokenSource(src func() (*oauth2.Token, error)) func(cfg *OIDCConfig) {
+	return func(cfg *OIDCConfig) {
+		cfg.tokenSource = src
+	}
+}
+
+func NewOIDCConfig(t *testing.T, issuer string, opts ...func(cfg *OIDCConfig)) *OIDCConfig {
 	t.Helper()
 
 	block, _ := pem.Decode([]byte(testRSAPrivateKey))
@@ -1035,27 +1057,52 @@ func NewOIDCConfig(t *testing.T, issuer string) *OIDCConfig {
 		issuer = "https://coder.com"
 	}
 
-	return &OIDCConfig{
+	cfg := &OIDCConfig{
 		key:    pkey,
 		issuer: issuer,
 	}
+	for _, opt := range opts {
+		opt(cfg)
+	}
+	return cfg
 }
 
 func (*OIDCConfig) AuthCodeURL(state string, _ ...oauth2.AuthCodeOption) string {
 	return "/?state=" + url.QueryEscape(state)
 }
 
-func (*OIDCConfig) TokenSource(context.Context, *oauth2.Token) oauth2.TokenSource {
-	return nil
+type tokenSource struct {
+	src func() (*oauth2.Token, error)
+}
+
+func (s tokenSource) Token() (*oauth2.Token, error) {
+	return s.src()
 }
 
-func (*OIDCConfig) Exchange(_ context.Context, code string, _ ...oauth2.AuthCodeOption) (*oauth2.Token, error) {
+func (cfg *OIDCConfig) TokenSource(context.Context, *oauth2.Token) oauth2.TokenSource {
+	if cfg.tokenSource == nil {
+		return nil
+	}
+	return tokenSource{
+		src: cfg.tokenSource,
+	}
+}
+
+func (cfg *OIDCConfig) Exchange(_ context.Context, code string, _ ...oauth2.AuthCodeOption) (*oauth2.Token, error) {
 	token, err := base64.StdEncoding.DecodeString(code)
 	if err != nil {
 		return nil, xerrors.Errorf("decode code: %w", err)
 	}
+
+	var exp time.Time
+	if cfg.oidcTokenExpires != nil {
+		exp = cfg.oidcTokenExpires()
+	}
+
 	return (&oauth2.Token{
-		AccessToken: "token",
+		AccessToken:  "token",
+		RefreshToken: cfg.refreshToken,
+		Expiry:       exp,
 	}).WithExtra(map[string]interface{}{
 		"id_token": string(token),
 	}), nil

diff --git a/coderd/userauth_test.go b/coderd/userauth_test.go
@@ -9,6 +9,7 @@ import (
 	"net/http/cookiejar"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/coreos/go-oidc/v3/oidc"
 	"github.com/golang-jwt/jwt"
@@ -24,12 +25,94 @@ import (
 	"github.com/coder/coder/coderd/audit"
 	"github.com/coder/coder/coderd/coderdtest"
 	"github.com/coder/coder/coderd/database"
+	"github.com/coder/coder/coderd/database/dbauthz"
 	"github.com/coder/coder/coderd/database/dbgen"
 	"github.com/coder/coder/coderd/database/dbtestutil"
 	"github.com/coder/coder/codersdk"
 	"github.com/coder/coder/testutil"
 )
 
+// This test specifically tests logging in with OIDC when an expired
+// OIDC session token exists.
+// The token refreshing should not happen since we are reauthenticating.
+func TestOIDCOauthLoginWithExisting(t *testing.T) {
+	conf := coderdtest.NewOIDCConfig(t, "",
+		// Provide a refresh token so we use the refresh token flow
+		coderdtest.WithRefreshToken("refresh_token"),
+		// We need to set the expire in the future for the first api calls.
+		coderdtest.WithTokenExpires(func() time.Time {
+			return time.Now().Add(time.Hour).UTC()
+		}),
+		// No refresh should actually happen in this test.
+		coderdtest.WithTokenSource(func() (*oauth2.Token, error) {
+			return nil, xerrors.New("token should not require refresh")
+		}),
+	)
+	logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
+	auditor := audit.NewMock()
+	const username = "alice"
+	claims := jwt.MapClaims{
+		"email":              "alice@coder.com",
+		"email_verified":     true,
+		"preferred_username": username,
+	}
+	config := conf.OIDCConfig(t, claims)
+
+	config.AllowSignups = true
+	config.IgnoreUserInfo = true
+	client, _, api := coderdtest.NewWithAPI(t, &coderdtest.Options{
+		Auditor:    auditor,
+		OIDCConfig: config,
+		Logger:     &logger,
+	})
+
+	// Signup alice
+	resp := oidcCallback(t, client, conf.EncodeClaims(t, claims))
+	// Set the client to use this OIDC context
+	authCookie := authCookieValue(resp.Cookies())
+	client.SetSessionToken(authCookie)
+	_ = resp.Body.Close()
+
+	ctx := testutil.Context(t, testutil.WaitLong*1000)
+	// Verify the user and oauth link
+	user, err := client.User(ctx, "me")
+	require.NoError(t, err)
+	require.Equal(t, username, user.Username)
+
+	// nolint:gocritic
+	link, err := api.Database.GetUserLinkByUserIDLoginType(dbauthz.AsSystemRestricted(ctx), database.GetUserLinkByUserIDLoginTypeParams{
+		UserID:    user.ID,
+		LoginType: database.LoginType(user.LoginType),
+	})
+	require.NoError(t, err, "failed to get user link")
+
+	// Expire the link
+	// nolint:gocritic
+	_, err = api.Database.UpdateUserLink(dbauthz.AsSystemRestricted(ctx), database.UpdateUserLinkParams{
+		OAuthAccessToken:  link.OAuthAccessToken,
+		OAuthRefreshToken: link.OAuthRefreshToken,
+		OAuthExpiry:       time.Now().Add(time.Hour * -1).UTC(),
+		UserID:            link.UserID,
+		LoginType:         link.LoginType,
+	})
+	require.NoError(t, err, "failed to update user link")
+
+	// Log in again with OIDC
+	loginAgain := oidcCallbackWithState(t, client, conf.EncodeClaims(t, claims), "seconds_login", func(req *http.Request) {
+		req.AddCookie(&http.Cookie{
+			Name:  codersdk.SessionTokenCookie,
+			Value: authCookie,
+			Path:  "/",
+		})
+	})
+	require.Equal(t, http.StatusTemporaryRedirect, loginAgain.StatusCode)
+
+	// Try to use new login
+	client.SetSessionToken(authCookieValue(resp.Cookies()))
+	_, err = client.User(ctx, "me")
+	require.NoError(t, err, "use new session")
+}
+
 func TestUserLogin(t *testing.T) {
 	t.Parallel()
 	t.Run("OK", func(t *testing.T) {
@@ -819,7 +902,7 @@ func TestUserOIDC(t *testing.T) {
 		})
 		require.NoError(t, err)
 
-		resp := oidcCallbackWithState(t, user, code, convertResponse.StateString)
+		resp := oidcCallbackWithState(t, user, code, convertResponse.StateString, nil)
 		require.Equal(t, http.StatusTemporaryRedirect, resp.StatusCode)
 	})
 
@@ -1045,10 +1128,10 @@ func oauth2Callback(t *testing.T, client *codersdk.Client) *http.Response {
 }
 
 func oidcCallback(t *testing.T, client *codersdk.Client, code string) *http.Response {
-	return oidcCallbackWithState(t, client, code, "somestate")
+	return oidcCallbackWithState(t, client, code, "somestate", nil)
 }
 
-func oidcCallbackWithState(t *testing.T, client *codersdk.Client, code, state string) *http.Response {
+func oidcCallbackWithState(t *testing.T, client *codersdk.Client, code, state string, modify func(r *http.Request)) *http.Response {
 	t.Helper()
 
 	client.HTTPClient.CheckRedirect = func(req *http.Request, via []*http.Request) error {
@@ -1062,6 +1145,9 @@ func oidcCallbackWithState(t *testing.T, client *codersdk.Client, code, state st
 		Name:  codersdk.OAuth2StateCookie,
 		Value: state,
 	})
+	if modify != nil {
+		modify(req)
+	}
 	res, err := client.HTTPClient.Do(req)
 	require.NoError(t, err)
 	defer res.Body.Close()