add majority of tests

Author: hf

internal/api/oauthserver/handlers.go

@@ -437,7 +437,7 @@ func (s *Server) handleAuthorizationCodeGrant(ctx context.Context, w http.Respon
 
 		// Issue the refresh token and access token
 		var terr error
-		tokenResponse, terr = tokenService.IssueRefreshToken(r, tx, user, authMethod, grantParams)
+		tokenResponse, terr = tokenService.IssueRefreshToken(r, w.Header(), tx, user, authMethod, grantParams)
 		if terr != nil {
 			return terr
 		}
@@ -488,7 +488,7 @@ func (s *Server) handleRefreshTokenGrant(ctx context.Context, w http.ResponseWri
 	}
 
 	db := s.db.WithContext(ctx)
-	tokenResponse, err := tokenService.RefreshTokenGrant(ctx, db, r, tokens.RefreshTokenGrantParams{
+	tokenResponse, err := tokenService.RefreshTokenGrant(ctx, db, r, w.Header(), tokens.RefreshTokenGrantParams{
 		RefreshToken: params.RefreshToken,
 		ClientID:     clientID,
 	})

internal/api/token.go

@@ -190,7 +190,7 @@ func (a *API) ResourceOwnerPasswordGrant(ctx context.Context, w http.ResponseWri
 		}); terr != nil {
 			return terr
 		}
-		token, terr = a.tokenService.IssueRefreshToken(r, tx, user, models.PasswordGrant, grantParams)
+		token, terr = a.tokenService.IssueRefreshToken(r, w.Header(), tx, user, models.PasswordGrant, grantParams)
 		if terr != nil {
 			return terr
 		}
@@ -260,7 +260,7 @@ func (a *API) PKCE(ctx context.Context, w http.ResponseWriter, r *http.Request)
 		}); terr != nil {
 			return terr
 		}
-		token, terr = a.tokenService.IssueRefreshToken(r, tx, user, authMethod, grantParams)
+		token, terr = a.tokenService.IssueRefreshToken(r, w.Header(), tx, user, authMethod, grantParams)
 		if terr != nil {
 			// error type is already handled in issueRefreshToken
 			return terr
@@ -295,7 +295,7 @@ func (a *API) generateAccessToken(r *http.Request, tx *storage.Connection, user
 }
 
 func (a *API) issueRefreshToken(r *http.Request, conn *storage.Connection, user *models.User, authenticationMethod models.AuthenticationMethod, grantParams models.GrantParams) (*tokens.AccessTokenResponse, error) {
-	return a.tokenService.IssueRefreshToken(r, conn, user, authenticationMethod, grantParams)
+	return a.tokenService.IssueRefreshToken(r, make(http.Header), conn, user, authenticationMethod, grantParams)
 }
 
 func (a *API) updateMFASessionAndClaims(r *http.Request, tx *storage.Connection, user *models.User, authenticationMethod models.AuthenticationMethod, grantParams models.GrantParams) (*tokens.AccessTokenResponse, error) {

internal/api/token_refresh.go

@@ -3,7 +3,10 @@ package api
 import (
 	"context"
 	"net/http"
+	"regexp"
 
+	"github.com/supabase/auth/internal/api/apierrors"
+	"github.com/supabase/auth/internal/crypto"
 	"github.com/supabase/auth/internal/tokens"
 )
 
@@ -12,15 +15,42 @@ type RefreshTokenGrantParams struct {
 	RefreshToken string `json:"refresh_token"`
 }
 
+var legacyRefreshTokenPattern = regexp.MustCompile("^[a-z0-9]{12}$")
+
+func (p *RefreshTokenGrantParams) Validate() error {
+	if len(p.RefreshToken) < 12 {
+		return apierrors.NewBadRequestError(apierrors.ErrorCodeValidationFailed, "Refresh token is not valid")
+	}
+
+	if len(p.RefreshToken) == 12 {
+		if !legacyRefreshTokenPattern.MatchString(p.RefreshToken) {
+			return apierrors.NewBadRequestError(apierrors.ErrorCodeValidationFailed, "Refresh token is not valid")
+		}
+
+		return nil
+	}
+
+	_, err := crypto.ParseRefreshToken(p.RefreshToken)
+	if err != nil {
+		return apierrors.NewBadRequestError(apierrors.ErrorCodeValidationFailed, "Refresh token is not valid").WithInternalError(err)
+	}
+
+	return nil
+}
+
 // RefreshTokenGrant implements the refresh_token grant type flow
 func (a *API) RefreshTokenGrant(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
 	params := &RefreshTokenGrantParams{}
 	if err := retrieveRequestParams(r, params); err != nil {
 		return err
 	}
 
+	if err := params.Validate(); err != nil {
+		return err
+	}
+
 	db := a.db.WithContext(ctx)
-	tokenResponse, err := a.tokenService.RefreshTokenGrant(ctx, db, r, tokens.RefreshTokenGrantParams{
+	tokenResponse, err := a.tokenService.RefreshTokenGrant(ctx, db, r, w.Header(), tokens.RefreshTokenGrantParams{
 		RefreshToken: params.RefreshToken,
 	})
 	if err != nil {

internal/crypto/crypto_test.go

@@ -105,4 +105,5 @@ func TestEncryptedStringDecryptNegative(t *testing.T) {
 
 func TestSecureToken(t *testing.T) {
 	assert.Equal(t, len(SecureAlphanumeric(22)), 22)
+	assert.Equal(t, len(SecureAlphanumeric(7)), 8)
 }

internal/crypto/refresh_tokens.go

@@ -61,7 +61,7 @@ func (r *RefreshToken) Encode(hmacSha256Key []byte) string {
 
 	result = append(result, 0)
 	result = append(result, r.SessionID.Bytes()...)
-	result = binary.AppendUvarint(result, uint64(r.Counter))
+	result = binary.AppendUvarint(result, safeUint64(r.Counter))
 
 	// Note on truncating the HMAC-SHA-256 output:
 	// This does not impact security as the brute-force space is 2^128 and
@@ -90,6 +90,22 @@ var (
 	ErrRefreshTokenCounterInvalid  = errors.New("crypto: refresh token's counter is not valid")
 )
 
+func safeInt64(v uint64) int64 {
+	if v > math.MaxInt64 {
+		return math.MaxInt64
+	}
+
+	return int64(v)
+}
+
+func safeUint64(v int64) uint64 {
+	if v < 0 {
+		return 0
+	}
+
+	return uint64(v)
+}
+
 func ParseRefreshToken(token string) (*RefreshToken, error) {
 	bytes, err := base64.RawURLEncoding.DecodeString(token)
 	if err != nil {
@@ -111,10 +127,7 @@ func ParseRefreshToken(token string) (*RefreshToken, error) {
 		return nil, ErrRefreshTokenChecksumInvalid
 	}
 
-	sessionID, err := uuid.FromBytes(parseFrom[0:16])
-	if err != nil {
-		return nil, err
-	}
+	sessionID := uuid.FromBytesOrNil(parseFrom[0:16])
 
 	parseFrom = parseFrom[16:]
 
@@ -123,10 +136,6 @@ func ParseRefreshToken(token string) (*RefreshToken, error) {
 		return nil, ErrRefreshTokenCounterInvalid
 	}
 
-	if counter > math.MaxInt64 {
-		return nil, ErrRefreshTokenCounterInvalid
-	}
-
 	parseFrom = parseFrom[counterBytes:]
 
 	if len(parseFrom) != 16 {
@@ -140,7 +149,7 @@ func ParseRefreshToken(token string) (*RefreshToken, error) {
 
 		Version:   0,
 		SessionID: sessionID,
-		Counter:   int64(counter),
+		Counter:   safeInt64(counter),
 		Signature: signature,
 	}, nil
 }

internal/crypto/refresh_tokens_test.go

@@ -4,13 +4,22 @@ import (
 	"crypto/sha256"
 	"encoding/base64"
 	"fmt"
+	"math"
 	"strings"
 	"testing"
 
 	"github.com/gofrs/uuid"
 	"github.com/stretchr/testify/require"
 )
 
+func TestSafeIntegers(t *testing.T) {
+	require.Equal(t, int64(math.MaxInt64), safeInt64(math.MaxUint64))
+	require.Equal(t, int64(math.MaxInt64), safeInt64(math.MaxInt64))
+
+	require.Equal(t, int64(0), safeUint64(-1))
+	require.Equal(t, int64(math.MaxInt64), safeInt64(math.MaxInt64))
+}
+
 func TestRefreshTokenParse(t *testing.T) {
 	negativeExamples := []struct {
 		value []byte
@@ -84,3 +93,9 @@ func TestRefreshTokenParse(t *testing.T) {
 	require.Equal(t, original.Raw, parsed.Raw)
 	require.Equal(t, original.Signature, parsed.Signature)
 }
+
+func TestRefreshTokenTableName(t *testing.T) {
+	require.Panics(t, func() {
+		RefreshToken{}.TableName()
+	})
+}

internal/models/sessions.go

@@ -167,6 +167,23 @@ func (s *Session) UpdateOnlyRefreshToken(tx *storage.Connection) error {
 	return tx.UpdateOnly(s, "refresh_token_counter")
 }
 
+func (s *Session) ReEncryptRefreshTokenHmacKey(tx *storage.Connection, dbEncryption conf.DatabaseEncryptionConfiguration) error {
+	key, _, err := s.GetRefreshTokenHmacKey(dbEncryption)
+	if err != nil {
+		return err
+	}
+
+	es, err := crypto.NewEncryptedString(s.ID.String(), []byte(base64.RawURLEncoding.EncodeToString(key)), dbEncryption.EncryptionKeyID, dbEncryption.EncryptionKey)
+	if err != nil {
+		return err
+	}
+
+	encryptedValue := es.String()
+	s.RefreshTokenHmacKey = &encryptedValue
+
+	return tx.UpdateOnly(s, "refresh_token_hmac_key")
+}
+
 type SessionValidityReason = int
 
 const (

internal/models/user.go

@@ -636,8 +636,13 @@ func FindUserByID(tx *storage.Connection, id uuid.UUID) (*User, error) {
 // lock will only be acquired if there's no other lock. In case there is a
 // lock, a IsNotFound(err) error will be returned.
 //
-// Second value returned is either *models.RefreshToken or *models.TODO.
+// Second value returned is either *models.RefreshToken or *crypto.RefreshToken.
 func FindUserWithRefreshToken(tx *storage.Connection, dbEncryption conf.DatabaseEncryptionConfiguration, token string, forUpdate bool) (*User, any, *Session, error) {
+	if len(token) < 12 {
+		// not a valid refresh token so don't bother looking it up in the database
+		return nil, nil, nil, SessionNotFoundError{}
+	}
+
 	if len(token) == 12 {
 		return findUserWithLegacyRefreshToken(tx, token, forUpdate)
 	}
@@ -648,11 +653,11 @@ func FindUserWithRefreshToken(tx *storage.Connection, dbEncryption conf.Database
 func findUserWithRefreshToken(tx *storage.Connection, dbEncryption conf.DatabaseEncryptionConfiguration, token string, forUpdate bool) (*User, *crypto.RefreshToken, *Session, error) {
 	refreshToken, err := crypto.ParseRefreshToken(token)
 	if err != nil {
-		return nil, nil, nil, err
+		// refresh token is not valid
+		return nil, nil, nil, SessionNotFoundError{}
 	}
 
-	// first find the session to check the token's signature
-	session, err := FindSessionByID(tx, refreshToken.SessionID, false)
+	session, err := FindSessionByID(tx, refreshToken.SessionID, forUpdate)
 	if err != nil {
 		return nil, nil, nil, err
 	}
@@ -663,28 +668,28 @@ func findUserWithRefreshToken(tx *storage.Connection, dbEncryption conf.Database
 		return nil, nil, nil, SessionNotFoundError{}
 	}
 
-	key, _, err := session.GetRefreshTokenHmacKey(dbEncryption)
+	key, shouldReEncrypt, err := session.GetRefreshTokenHmacKey(dbEncryption)
 	if err != nil {
 		return nil, nil, nil, err
 	}
 
 	if !refreshToken.CheckSignature(key) {
-		// TODO: return SessionNotFound, log informational
-		return nil, nil, nil, fmt.Errorf("refresh token for session %s with counter %v has invalid signature", session.ID.String(), refreshToken.Counter)
-	}
-
-	user, err := FindUserByID(tx, session.UserID)
-	if err != nil {
-		return nil, nil, nil, err
+		// refresh token signature is not valid for this session
+		return nil, nil, nil, SessionNotFoundError{}
 	}
 
-	if forUpdate {
-		session, err = FindSessionByID(tx, refreshToken.SessionID, forUpdate)
+	if shouldReEncrypt && forUpdate {
+		err := session.ReEncryptRefreshTokenHmacKey(tx, dbEncryption)
 		if err != nil {
 			return nil, nil, nil, err
 		}
 	}
 
+	user, err := FindUserByID(tx, session.UserID)
+	if err != nil {
+		return nil, nil, nil, err
+	}
+
 	return user, refreshToken, session, nil
 }