fixing tests

coder · sreya · Oct 25, 2024 · Oct 5, 2024 · Oct 5, 2024 · Oct 5, 2024
commit 2ea31abd1d14d81d00b9a00b6468544185437ee3
diff --git a/tailnet/resume.go b/tailnet/resume.go
@@ -3,11 +3,8 @@ package tailnet
 import (
 	"context"
 	"crypto/rand"
-	"database/sql"
-	"encoding/hex"
 	"time"
 
-	"github.com/go-jose/go-jose/v4"
 	"github.com/go-jose/go-jose/v4/jwt"
 	"github.com/google/uuid"
 	"golang.org/x/xerrors"
@@ -53,47 +50,6 @@ func GenerateResumeTokenSigningKey() (ResumeTokenSigningKey, error) {
 	return key, nil
 }
 
-type ResumeTokenSigningKeyDatabaseStore interface {
-	GetCoordinatorResumeTokenSigningKey(ctx context.Context) (string, error)
-	UpsertCoordinatorResumeTokenSigningKey(ctx context.Context, key string) error
-}
-
-// ResumeTokenSigningKeyFromDatabase retrieves the coordinator resume token
-// signing key from the database. If the key is not found, a new key is
-// generated and inserted into the database.
-func ResumeTokenSigningKeyFromDatabase(ctx context.Context, db ResumeTokenSigningKeyDatabaseStore) (ResumeTokenSigningKey, error) {
-	var resumeTokenKey ResumeTokenSigningKey
-	resumeTokenKeyStr, err := db.GetCoordinatorResumeTokenSigningKey(ctx)
-	if err != nil && !xerrors.Is(err, sql.ErrNoRows) {
-		return resumeTokenKey, xerrors.Errorf("get coordinator resume token key: %w", err)
-	}
-	if decoded, err := hex.DecodeString(resumeTokenKeyStr); err != nil || len(decoded) != len(resumeTokenKey) {
-		newKey, err := GenerateResumeTokenSigningKey()
-		if err != nil {
-			return resumeTokenKey, xerrors.Errorf("generate fresh coordinator resume token key: %w", err)
-		}
-
-		resumeTokenKeyStr = hex.EncodeToString(newKey[:])
-		err = db.UpsertCoordinatorResumeTokenSigningKey(ctx, resumeTokenKeyStr)
-		if err != nil {
-			return resumeTokenKey, xerrors.Errorf("insert freshly generated coordinator resume token key to database: %w", err)
-		}
-	}
-
-	resumeTokenKeyBytes, err := hex.DecodeString(resumeTokenKeyStr)
-	if err != nil {
-		return resumeTokenKey, xerrors.Errorf("decode coordinator resume token key from database: %w", err)
-	}
-	if len(resumeTokenKeyBytes) != len(resumeTokenKey) {
-		return resumeTokenKey, xerrors.Errorf("coordinator resume token key in database is not the correct length, expect %d got %d", len(resumeTokenKey), len(resumeTokenKeyBytes))
-	}
-	copy(resumeTokenKey[:], resumeTokenKeyBytes)
-	if resumeTokenKey == [64]byte{} {
-		return resumeTokenKey, xerrors.Errorf("coordinator resume token key in database is empty")
-	}
-	return resumeTokenKey, nil
-}
-
 type ResumeTokenKeyProvider struct {
 	key    jwtutils.SigningKeyManager
 	clock  quartz.Clock
@@ -111,19 +67,11 @@ func NewResumeTokenKeyProvider(key jwtutils.SigningKeyManager, clock quartz.Cloc
 	}
 }
 
-type resumeTokenPayload struct {
-	jwt.Claims
-	PeerID uuid.UUID `json:"sub"`
-	Expiry int64     `json:"exp"`
-}
-
 func (p ResumeTokenKeyProvider) GenerateResumeToken(ctx context.Context, peerID uuid.UUID) (*proto.RefreshResumeTokenResponse, error) {
 	exp := p.clock.Now().Add(p.expiry)
-	payload := resumeTokenPayload{
-		PeerID: peerID,
-		Claims: jwt.Claims{
-			Expiry: jwt.NewNumericDate(exp),
-		},
+	payload := jwt.Claims{
+		Subject: peerID.String(),
+		Expiry:  jwt.NewNumericDate(exp),
 	}
 
 	token, err := jwtutils.Sign(ctx, p.key, payload)
@@ -142,12 +90,16 @@ func (p ResumeTokenKeyProvider) GenerateResumeToken(ctx context.Context, peerID
 // returns the payload. If the token is invalid or expired, an error is
 // returned.
 func (p ResumeTokenKeyProvider) VerifyResumeToken(ctx context.Context, str string) (uuid.UUID, error) {
-	var tok resumeTokenPayload
+	var tok jwt.Claims
 	err := jwtutils.Verify(ctx, p.key, str, &tok, jwtutils.WithVerifyExpected(jwt.Expected{
 		Time: p.clock.Now(),
 	}))
 	if err != nil {
 		return uuid.Nil, xerrors.Errorf("verify payload: %w", err)
 	}
-	return tok.PeerID, nil
+	parsed, err := uuid.Parse(tok.Subject)
+	if err != nil {
+		return uuid.Nil, xerrors.Errorf("parse peerID from token: %w", err)
+	}
+	return parsed, nil
 }
diff --git a/tailnet/resume_test.go b/tailnet/resume_test.go
@@ -1,118 +1,20 @@
 package tailnet_test
 
 import (
-	"context"
-	"encoding/hex"
 	"testing"
 	"time"
 
+	"github.com/go-jose/go-jose/v4"
+	"github.com/go-jose/go-jose/v4/jwt"
 	"github.com/google/uuid"
-	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
-	"go.uber.org/mock/gomock"
 
-	"github.com/coder/coder/v2/coderd/database/dbmock"
-	"github.com/coder/coder/v2/coderd/database/dbtestutil"
 	"github.com/coder/coder/v2/coderd/jwtutils"
 	"github.com/coder/coder/v2/tailnet"
 	"github.com/coder/coder/v2/testutil"
 	"github.com/coder/quartz"
 )
 
-func TestResumeTokenSigningKeyFromDatabase(t *testing.T) {
-	t.Parallel()
-
-	assertRandomKey := func(t *testing.T, key tailnet.ResumeTokenSigningKey) {
-		t.Helper()
-		assert.NotEqual(t, tailnet.ResumeTokenSigningKey{}, key, "key should not be empty")
-		assert.NotEqualValues(t, [64]byte{1}, key, "key should not be all 1s")
-	}
-
-	t.Run("GenerateRetrieve", func(t *testing.T) {
-		t.Parallel()
-
-		db, _ := dbtestutil.NewDB(t)
-		ctx := testutil.Context(t, testutil.WaitShort)
-		key1, err := tailnet.ResumeTokenSigningKeyFromDatabase(ctx, db)
-		require.NoError(t, err)
-		assertRandomKey(t, key1)
-
-		key2, err := tailnet.ResumeTokenSigningKeyFromDatabase(ctx, db)
-		require.NoError(t, err)
-		require.Equal(t, key1, key2, "keys should not be different")
-	})
-
-	t.Run("GetError", func(t *testing.T) {
-		t.Parallel()
-
-		db := dbmock.NewMockStore(gomock.NewController(t))
-		db.EXPECT().GetCoordinatorResumeTokenSigningKey(gomock.Any()).Return("", assert.AnError)
-
-		ctx := testutil.Context(t, testutil.WaitShort)
-		_, err := tailnet.ResumeTokenSigningKeyFromDatabase(ctx, db)
-		require.ErrorIs(t, err, assert.AnError)
-	})
-
-	t.Run("UpsertError", func(t *testing.T) {
-		t.Parallel()
-
-		db := dbmock.NewMockStore(gomock.NewController(t))
-		db.EXPECT().GetCoordinatorResumeTokenSigningKey(gomock.Any()).Return("", nil)
-		db.EXPECT().UpsertCoordinatorResumeTokenSigningKey(gomock.Any(), gomock.Any()).Return(assert.AnError)
-
-		ctx := testutil.Context(t, testutil.WaitShort)
-		_, err := tailnet.ResumeTokenSigningKeyFromDatabase(ctx, db)
-		require.ErrorIs(t, err, assert.AnError)
-	})
-
-	t.Run("DecodeErrorShouldRegenerate", func(t *testing.T) {
-		t.Parallel()
-
-		db := dbmock.NewMockStore(gomock.NewController(t))
-		db.EXPECT().GetCoordinatorResumeTokenSigningKey(gomock.Any()).Return("invalid", nil)
-
-		var storedKey tailnet.ResumeTokenSigningKey
-		db.EXPECT().UpsertCoordinatorResumeTokenSigningKey(gomock.Any(), gomock.Any()).Do(func(_ context.Context, value string) error {
-			keyBytes, err := hex.DecodeString(value)
-			require.NoError(t, err)
-			require.Len(t, keyBytes, len(storedKey))
-			copy(storedKey[:], keyBytes)
-			return nil
-		})
-
-		ctx := testutil.Context(t, testutil.WaitShort)
-		key, err := tailnet.ResumeTokenSigningKeyFromDatabase(ctx, db)
-		require.NoError(t, err)
-		assertRandomKey(t, key)
-		require.Equal(t, storedKey, key, "key should match stored value")
-	})
-
-	t.Run("LengthErrorShouldRegenerate", func(t *testing.T) {
-		t.Parallel()
-
-		db := dbmock.NewMockStore(gomock.NewController(t))
-		db.EXPECT().GetCoordinatorResumeTokenSigningKey(gomock.Any()).Return("deadbeef", nil)
-		db.EXPECT().UpsertCoordinatorResumeTokenSigningKey(gomock.Any(), gomock.Any()).Return(nil)
-
-		ctx := testutil.Context(t, testutil.WaitShort)
-		key, err := tailnet.ResumeTokenSigningKeyFromDatabase(ctx, db)
-		require.NoError(t, err)
-		assertRandomKey(t, key)
-	})
-
-	t.Run("EmptyError", func(t *testing.T) {
-		t.Parallel()
-
-		db := dbmock.NewMockStore(gomock.NewController(t))
-		emptyKey := hex.EncodeToString(make([]byte, 64))
-		db.EXPECT().GetCoordinatorResumeTokenSigningKey(gomock.Any()).Return(emptyKey, nil)
-
-		ctx := testutil.Context(t, testutil.WaitShort)
-		_, err := tailnet.ResumeTokenSigningKeyFromDatabase(ctx, db)
-		require.ErrorContains(t, err, "is empty")
-	})
-}
-
 func TestResumeTokenKeyProvider(t *testing.T) {
 	t.Parallel()
 
@@ -156,7 +58,7 @@ func TestResumeTokenKeyProvider(t *testing.T) {
 		_ = clock.Advance(tailnet.DefaultResumeTokenExpiry + time.Second)
 
 		_, err = provider.VerifyResumeToken(ctx, token.Token)
-		require.ErrorContains(t, err, "expired")
+		require.ErrorIs(t, err, jwt.ErrExpired)
 	})
 
 	t.Run("InvalidToken", func(t *testing.T) {
@@ -175,17 +77,20 @@ func TestResumeTokenKeyProvider(t *testing.T) {
 		// Generate a resume token with a different key
 		otherKey, err := tailnet.GenerateResumeTokenSigningKey()
 		require.NoError(t, err)
-		otherProvider := tailnet.NewResumeTokenKeyProvider(newKeySigner(otherKey), quartz.NewMock(t), tailnet.DefaultResumeTokenExpiry)
+		otherSigner := newKeySigner(otherKey)
+		otherProvider := tailnet.NewResumeTokenKeyProvider(otherSigner, quartz.NewMock(t), tailnet.DefaultResumeTokenExpiry)
 		token, err := otherProvider.GenerateResumeToken(ctx, uuid.New())
 		require.NoError(t, err)
 
-		provider := tailnet.NewResumeTokenKeyProvider(newKeySigner(key), quartz.NewMock(t), tailnet.DefaultResumeTokenExpiry)
+		signer := newKeySigner(key)
+		signer.ID = otherSigner.ID
+		provider := tailnet.NewResumeTokenKeyProvider(signer, quartz.NewMock(t), tailnet.DefaultResumeTokenExpiry)
 		_, err = provider.VerifyResumeToken(ctx, token.Token)
-		require.ErrorContains(t, err, "verify JWS")
+		require.ErrorIs(t, err, jose.ErrCryptoFailure)
 	})
 }
 
-func newKeySigner(key tailnet.ResumeTokenSigningKey) jwtutils.SigningKeyManager {
+func newKeySigner(key tailnet.ResumeTokenSigningKey) jwtutils.StaticKeyManager {
 	return jwtutils.StaticKeyManager{
 		ID:  uuid.New().String(),
 		Key: key[:],