coder · Emyrk · Aug 25, 2023 · Aug 23, 2023 · Aug 24, 2023 · Aug 24, 2023
diff --git a/coderd/coderdtest/coderdtest.go b/coderd/coderdtest/coderdtest.go
@@ -31,15 +31,13 @@ import (
 	"time"
 
 	"cloud.google.com/go/compute/metadata"
-	"github.com/coreos/go-oidc/v3/oidc"
 	"github.com/fullsailor/pkcs7"
-	"github.com/golang-jwt/jwt"
+	"github.com/golang-jwt/jwt/v4"
 	"github.com/google/uuid"
 	"github.com/moby/moby/pkg/namesgenerator"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
-	"golang.org/x/oauth2"
 	"golang.org/x/xerrors"
 	"google.golang.org/api/idtoken"
 	"google.golang.org/api/option"
@@ -1020,152 +1018,6 @@ 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 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))
-	pkey, err := x509.ParsePKCS1PrivateKey(block.Bytes)
-	require.NoError(t, err)
-
-	if issuer == "" {
-		issuer = "https://coder.com"
-	}
-
-	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)
-}
-
-type tokenSource struct {
-	src func() (*oauth2.Token, error)
-}
-
-func (s tokenSource) Token() (*oauth2.Token, error) {
-	return s.src()
-}
-
-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",
-		RefreshToken: cfg.refreshToken,
-		Expiry:       exp,
-	}).WithExtra(map[string]interface{}{
-		"id_token": string(token),
-	}), nil
-}
-
-func (cfg *OIDCConfig) EncodeClaims(t *testing.T, claims jwt.MapClaims) string {
-	t.Helper()
-
-	if _, ok := claims["exp"]; !ok {
-		claims["exp"] = time.Now().Add(time.Hour).UnixMilli()
-	}
-
-	if _, ok := claims["iss"]; !ok {
-		claims["iss"] = cfg.issuer
-	}
-
-	if _, ok := claims["sub"]; !ok {
-		claims["sub"] = "testme"
-	}
-
-	signed, err := jwt.NewWithClaims(jwt.SigningMethodRS256, claims).SignedString(cfg.key)
-	require.NoError(t, err)
-
-	return base64.StdEncoding.EncodeToString([]byte(signed))
-}
-
-func (cfg *OIDCConfig) OIDCConfig(t *testing.T, userInfoClaims jwt.MapClaims, opts ...func(cfg *coderd.OIDCConfig)) *coderd.OIDCConfig {
-	// By default, the provider can be empty.
-	// This means it won't support any endpoints!
-	provider := &oidc.Provider{}
-	if userInfoClaims != nil {
-		resp, err := json.Marshal(userInfoClaims)
-		require.NoError(t, err)
-		srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			w.WriteHeader(http.StatusOK)
-			_, _ = w.Write(resp)
-		}))
-		t.Cleanup(srv.Close)
-		cfg := &oidc.ProviderConfig{
-			UserInfoURL: srv.URL,
-		}
-		provider = cfg.NewProvider(context.Background())
-	}
-	newCFG := &coderd.OIDCConfig{
-		OAuth2Config: cfg,
-		Verifier: oidc.NewVerifier(cfg.issuer, &oidc.StaticKeySet{
-			PublicKeys: []crypto.PublicKey{cfg.key.Public()},
-		}, &oidc.Config{
-			SkipClientIDCheck: true,
-		}),
-		Provider:      provider,
-		UsernameField: "preferred_username",
-		EmailField:    "email",
-		AuthURLParams: map[string]string{"access_type": "offline"},
-		GroupField:    "groups",
-	}
-	for _, opt := range opts {
-		opt(newCFG)
-	}
-	return newCFG
-}
-
 // NewAzureInstanceIdentity returns a metadata client and ID token validator for faking
 // instance authentication for Azure.
 func NewAzureInstanceIdentity(t *testing.T, instanceID string) (x509.VerifyOptions, *http.Client) {
@@ -1254,22 +1106,6 @@ func SDKError(t *testing.T, err error) *codersdk.Error {
 	return cerr
 }
 
-const testRSAPrivateKey = `-----BEGIN RSA PRIVATE KEY-----
-MIICXQIBAAKBgQDLets8+7M+iAQAqN/5BVyCIjhTQ4cmXulL+gm3v0oGMWzLupUS
-v8KPA+Tp7dgC/DZPfMLaNH1obBBhJ9DhS6RdS3AS3kzeFrdu8zFHLWF53DUBhS92
-5dCAEuJpDnNizdEhxTfoHrhuCmz8l2nt1pe5eUK2XWgd08Uc93h5ij098wIDAQAB
-AoGAHLaZeWGLSaen6O/rqxg2laZ+jEFbMO7zvOTruiIkL/uJfrY1kw+8RLIn+1q0
-wLcWcuEIHgKKL9IP/aXAtAoYh1FBvRPLkovF1NZB0Je/+CSGka6wvc3TGdvppZJe
-rKNcUvuOYLxkmLy4g9zuY5qrxFyhtIn2qZzXEtLaVOHzPQECQQDvN0mSajpU7dTB
-w4jwx7IRXGSSx65c+AsHSc1Rj++9qtPC6WsFgAfFN2CEmqhMbEUVGPv/aPjdyWk9
-pyLE9xR/AkEA2cGwyIunijE5v2rlZAD7C4vRgdcMyCf3uuPcgzFtsR6ZhyQSgLZ8
-YRPuvwm4cdPJMmO3YwBfxT6XGuSc2k8MjQJBAI0+b8prvpV2+DCQa8L/pjxp+VhR
-Xrq2GozrHrgR7NRokTB88hwFRJFF6U9iogy9wOx8HA7qxEbwLZuhm/4AhbECQC2a
-d8h4Ht09E+f3nhTEc87mODkl7WJZpHL6V2sORfeq/eIkds+H6CJ4hy5w/bSw8tjf
-sz9Di8sGIaUbLZI2rd0CQQCzlVwEtRtoNCyMJTTrkgUuNufLP19RZ5FpyXxBO5/u
-QastnN77KfUwdj3SJt44U/uh1jAIv4oSLBr8HYUkbnI8
------END RSA PRIVATE KEY-----`
-
 func DeploymentValues(t testing.TB) *codersdk.DeploymentValues {
 	var cfg codersdk.DeploymentValues
 	opts := cfg.Options()

diff --git a/coderd/coderdtest/oidctest/helper.go b/coderd/coderdtest/oidctest/helper.go
@@ -0,0 +1,103 @@
+package oidctest
+
+import (
+	"net/http"
+	"testing"
+	"time"
+
+	"github.com/golang-jwt/jwt/v4"
+	"github.com/stretchr/testify/require"
+
+	"github.com/coder/coder/v2/coderd/database"
+	"github.com/coder/coder/v2/coderd/database/dbauthz"
+	"github.com/coder/coder/v2/coderd/httpmw"
+	"github.com/coder/coder/v2/codersdk"
+	"github.com/coder/coder/v2/testutil"
+)
+
+// LoginHelper helps with logging in a user and refreshing their oauth tokens.
+// It is mainly because refreshing oauth tokens is a bit tricky and requires
+// some database manipulation.
+type LoginHelper struct {
+	fake   *FakeIDP
+	client *codersdk.Client
+}
+
+func NewLoginHelper(client *codersdk.Client, fake *FakeIDP) *LoginHelper {
+	if client == nil {
+		panic("client must not be nil")
+	}
+	if fake == nil {
+		panic("fake must not be nil")
+	}
+	return &LoginHelper{
+		fake:   fake,
+		client: client,
+	}
+}
+
+// Login just helps by making an unauthenticated client and logging in with
+// the given claims. All Logins should be unauthenticated, so this is a
+// convenience method.
+func (h *LoginHelper) Login(t *testing.T, idTokenClaims jwt.MapClaims) (*codersdk.Client, *http.Response) {
+	t.Helper()
+	unauthenticatedClient := codersdk.New(h.client.URL)
+
+	return h.fake.Login(t, unauthenticatedClient, idTokenClaims)
+}
+
+// ExpireOauthToken expires the oauth token for the given user.
+func (*LoginHelper) ExpireOauthToken(t *testing.T, db database.Store, user *codersdk.Client) database.UserLink {
+	t.Helper()
+
+	//nolint:gocritic // Testing
+	ctx := dbauthz.AsSystemRestricted(testutil.Context(t, testutil.WaitMedium))
+
+	id, _, err := httpmw.SplitAPIToken(user.SessionToken())
+	require.NoError(t, err)
+
+	// We need to get the OIDC link and update it in the database to force
+	// it to be expired.
+	key, err := db.GetAPIKeyByID(ctx, id)
+	require.NoError(t, err, "get api key")
+
+	link, err := db.GetUserLinkByUserIDLoginType(ctx, database.GetUserLinkByUserIDLoginTypeParams{
+		UserID:    key.UserID,
+		LoginType: database.LoginTypeOIDC,
+	})
+	require.NoError(t, err, "get user link")
+
+	// Expire the oauth link for the given user.
+	updated, err := db.UpdateUserLink(ctx, database.UpdateUserLinkParams{
+		OAuthAccessToken:  link.OAuthAccessToken,
+		OAuthRefreshToken: link.OAuthRefreshToken,
+		OAuthExpiry:       time.Now().Add(time.Hour * -1),
+		UserID:            link.UserID,
+		LoginType:         link.LoginType,
+	})
+	require.NoError(t, err, "expire user link")
+
+	return updated
+}
+
+// ForceRefresh forces the client to refresh its oauth token. It does this by
+// expiring the oauth token, then doing an authenticated call. This will force
+// the API Key middleware to refresh the oauth token.
+//
+// A unit test assertion makes sure the refresh token is used.
+func (h *LoginHelper) ForceRefresh(t *testing.T, db database.Store, user *codersdk.Client, idToken jwt.MapClaims) {
+	t.Helper()
+
+	link := h.ExpireOauthToken(t, db, user)
+	// Updates the claims that the IDP will return. By default, it always
+	// uses the original claims for the original oauth token.
+	h.fake.UpdateRefreshClaims(link.OAuthRefreshToken, idToken)
+
+	t.Cleanup(func() {
+		require.True(t, h.fake.RefreshUsed(link.OAuthRefreshToken), "refresh token must be used, but has not. Did you forget to call the returned function from this call?")
+	})
+
+	// Do any authenticated call to force the refresh
+	_, err := user.User(testutil.Context(t, testutil.WaitShort), "me")
+	require.NoError(t, err, "user must be able to be fetched")
+}