feat(otp): add one_time_tokens query helpers

internal/api/pkce.go

@@ -10,7 +10,7 @@ import (
 )
 
 const (
-	PKCEPrefix                    = "pkce_"
+	PKCEPrefix                    = models.PKCEPrefix
 	MinCodeChallengeLength        = 43
 	MaxCodeChallengeLength        = 128
 	InvalidPKCEParamsErrorMessage = "PKCE flow requires code_challenge_method and code_challenge"

internal/models/one_time_token.go

@@ -92,6 +92,8 @@ func (t *OneTimeTokenType) Scan(src interface{}) error {
 	return nil
 }
 
+const PKCEPrefix = "pkce_"
+
 type OneTimeTokenNotFoundError struct {
 }
 
@@ -163,21 +165,48 @@ func CreateOneTimeToken(
 }
 
 func FindOneTimeToken(tx *storage.Connection, tokenHash string, tokenTypes ...OneTimeTokenType) (*OneTimeToken, error) {
+	return findOneTimeToken(tx, tokenHash, false, tokenTypes...)
+}
+
+// FindOneTimeTokenWithPKCEFallback finds the one time token of the given
+// types whose hash is either tokenHash or tokenHash with the "pkce_" prefix,
+// in a single query. An exact match is preferred over a prefixed one.
+// It returns OneTimeTokenNotFoundError when no row exists.
+func FindOneTimeTokenWithPKCEFallback(tx *storage.Connection, tokenHash string, tokenTypes ...OneTimeTokenType) (*OneTimeToken, error) {
+	return findOneTimeToken(tx, tokenHash, true, tokenTypes...)
+}
+
+// findOneTimeToken finds the one time token of the given types by tokenHash.
+// With pkceFallback it also accepts PKCEPrefix+tokenHash and prefers the
+// exact match. It returns OneTimeTokenNotFoundError when no row exists.
+func findOneTimeToken(tx *storage.Connection, tokenHash string, pkceFallback bool, tokenTypes ...OneTimeTokenType) (*OneTimeToken, error) {
 	oneTimeToken := &OneTimeToken{}
 
 	query := tx.Eager().Q()
 
+	hashClause, hashArgs := "token_hash = ?", []interface{}{tokenHash}
+	if pkceFallback {
+		hashClause, hashArgs = "token_hash in (?, ?)", []interface{}{tokenHash, PKCEPrefix + tokenHash}
+	}
+
 	switch len(tokenTypes) {
 	case 2:
-		query = query.Where("(token_type = ? or token_type = ?) and token_hash = ?", tokenTypes[0], tokenTypes[1], tokenHash) // #nosec G602
+		args := append([]interface{}{tokenTypes[0], tokenTypes[1]}, hashArgs...) // #nosec G602
+		query = query.Where("(token_type = ? or token_type = ?) and "+hashClause, args...)
 
 	case 1:
-		query = query.Where("token_type = ? and token_hash = ?", tokenTypes[0], tokenHash)
+		args := append([]interface{}{tokenTypes[0]}, hashArgs...)
+		query = query.Where("token_type = ? and "+hashClause, args...)
 
 	default:
 		panic("at most 2 token types are accepted")
 	}
 
+	if pkceFallback {
+		// true sorts before false in descending order, so this allows us to prefer an exact match
+		query = query.Order("token_hash = ? desc", tokenHash)
+	}
+
 	if err := query.First(oneTimeToken); err != nil {
 		if errors.Cause(err) == sql.ErrNoRows {
 			return nil, OneTimeTokenNotFoundError{}
@@ -189,6 +218,30 @@ func FindOneTimeToken(tx *storage.Connection, tokenHash string, tokenTypes ...On
 	return oneTimeToken, nil
 }
 
+// FindOneTimeTokenByRelatesTo finds the newest one time token of the given
+// token type by the relatesTo field.
+//
+// relates_to is not unique across users. For PhoneChangeToken in particular,
+// two users can hold rows for the same phone number, so the returned row does
+// not identify a user on its own. Callers must check the user against the
+// request before they trust the result.
+//
+// It returns OneTimeTokenNotFoundError when no row exists.
+func FindOneTimeTokenByRelatesTo(tx *storage.Connection, relatesTo string, tokenType OneTimeTokenType) (*OneTimeToken, error) {
+	oneTimeToken := &OneTimeToken{}
+
+	err := tx.Eager().Q().
+		Where("token_type = ? and relates_to = ?", tokenType, strings.ToLower(relatesTo)).
+		Order("created_at desc").
+		First(oneTimeToken)
+	if errors.Cause(err) == sql.ErrNoRows {
+		return nil, OneTimeTokenNotFoundError{}
+	} else if err != nil {
+		return nil, errors.Wrap(err, "error finding one time token")
+	}
+	return oneTimeToken, nil
+}
+
 // FindUserByOneTimeToken finds the user holding the one-time token matching
 // tokenHash for any of the given token types.
 func FindUserByOneTimeToken(tx *storage.Connection, tokenHash string, tokenTypes ...OneTimeTokenType) (*User, error) {
@@ -208,7 +261,7 @@ func FindUserByEmailChangeCurrentAndAudience(tx *storage.Connection, email, toke
 	}
 
 	if ott == nil {
-		ott, err = FindOneTimeToken(tx, "pkce_"+token, EmailChangeTokenCurrent)
+		ott, err = FindOneTimeToken(tx, PKCEPrefix+token, EmailChangeTokenCurrent)
 		if err != nil {
 			return nil, err
 		}
@@ -237,7 +290,7 @@ func FindUserByEmailChangeNewAndAudience(tx *storage.Connection, email, token, a
 	}
 
 	if ott == nil {
-		ott, err = FindOneTimeToken(tx, "pkce_"+token, EmailChangeTokenNew)
+		ott, err = FindOneTimeToken(tx, PKCEPrefix+token, EmailChangeTokenNew)
 		if err != nil && !IsNotFoundError(err) {
 			return nil, err
 		}

internal/models/one_time_token_test.go

@@ -48,6 +48,14 @@ func (ts *OneTimeTokenTestSuite) createUser() *User {
 	return u
 }
 
+// seedToken starts from an empty table and returns the user who owns the row.
+func (ts *OneTimeTokenTestSuite) seedToken(hash string, tokenType OneTimeTokenType) *User {
+	TruncateAll(ts.db)
+	u := ts.createUser()
+	require.NoError(ts.T(), CreateOneTimeToken(ts.db, u.ID, u.GetEmail(), hash, tokenType, time.Minute))
+	return u
+}
+
 func (ts *OneTimeTokenTestSuite) TestCreateOneTimeToken() {
 	cases := map[string]time.Duration{
 		"future window": 15 * time.Minute,
@@ -97,3 +105,142 @@ func (ts *OneTimeTokenTestSuite) TestCreateOneTimeTokenResendReplacesWindow() {
 	require.True(ts.T(), second.ExpiresAt.After(*first.ExpiresAt),
 		"resend must move expires_at forward, first=%s second=%s", first.ExpiresAt, second.ExpiresAt)
 }
+
+func (ts *OneTimeTokenTestSuite) TestFindOneTimeToken() {
+	ts.Run("matches the exact hash only, not the pkce_ prefixed form", func() {
+		ts.seedToken("pkce_hash", ConfirmationToken)
+
+		ott, err := FindOneTimeToken(ts.db, "hash", ConfirmationToken)
+		require.True(ts.T(), IsNotFoundError(err), "expected not found error, got %v", err)
+		require.Nil(ts.T(), ott)
+	})
+
+	ts.Run("does not return a row of another token type", func() {
+		ts.seedToken("hash", RecoveryToken)
+
+		ott, err := FindOneTimeToken(ts.db, "hash", ConfirmationToken)
+		require.True(ts.T(), IsNotFoundError(err), "expected not found error, got %v", err)
+		require.Nil(ts.T(), ott)
+	})
+
+	ts.Run("matches either of two token types", func() {
+		u := ts.seedToken("hash", RecoveryToken)
+
+		// The row has the second type, so this also checks the argument order.
+		ott, err := FindOneTimeToken(ts.db, "hash", ConfirmationToken, RecoveryToken)
+		require.NoError(ts.T(), err)
+		require.Equal(ts.T(), RecoveryToken, ott.TokenType)
+		require.Equal(ts.T(), u.ID, ott.UserID)
+	})
+}
+
+func (ts *OneTimeTokenTestSuite) TestFindOneTimeTokenWithPKCEFallback() {
+	ts.Run("exact hash match", func() {
+		u := ts.seedToken("hash", ConfirmationToken)
+
+		ott, err := FindOneTimeTokenWithPKCEFallback(ts.db, "hash", ConfirmationToken)
+		require.NoError(ts.T(), err)
+		require.Equal(ts.T(), "hash", ott.TokenHash)
+		require.Equal(ts.T(), u.ID, ott.UserID)
+	})
+
+	ts.Run("falls back to pkce_ prefixed hash", func() {
+		u := ts.seedToken("pkce_hash", ConfirmationToken)
+
+		ott, err := FindOneTimeTokenWithPKCEFallback(ts.db, "hash", ConfirmationToken)
+		require.NoError(ts.T(), err)
+		require.Equal(ts.T(), "pkce_hash", ott.TokenHash)
+		require.Equal(ts.T(), u.ID, ott.UserID)
+	})
+
+	ts.Run("prefers exact match over pkce_ prefixed hash", func() {
+		// (user_id, token_type) is unique, so the two candidates have to be
+		// different types. Both types are passed so both are eligible.
+		u := ts.seedToken("hash", ConfirmationToken)
+		require.NoError(ts.T(), CreateOneTimeToken(ts.db, u.ID, u.GetEmail(), "pkce_hash", RecoveryToken, time.Minute))
+
+		ott, err := FindOneTimeTokenWithPKCEFallback(ts.db, "hash", ConfirmationToken, RecoveryToken)
+		require.NoError(ts.T(), err)
+		require.Equal(ts.T(), "hash", ott.TokenHash)
+		require.Equal(ts.T(), ConfirmationToken, ott.TokenType)
+	})
+
+	ts.Run("not found when neither hash exists", func() {
+		TruncateAll(ts.db)
+		ts.createUser()
+
+		ott, err := FindOneTimeTokenWithPKCEFallback(ts.db, "missing", ConfirmationToken)
+		require.True(ts.T(), IsNotFoundError(err), "expected not found error, got %v", err)
+		require.Nil(ts.T(), ott)
+	})
+
+	ts.Run("token type filter applies to the pkce_ fallback", func() {
+		ts.seedToken("pkce_hash", RecoveryToken)
+
+		ott, err := FindOneTimeTokenWithPKCEFallback(ts.db, "hash", ConfirmationToken)
+		require.True(ts.T(), IsNotFoundError(err), "expected not found error, got %v", err)
+		require.Nil(ts.T(), ott)
+	})
+}
+
+func (ts *OneTimeTokenTestSuite) TestFindOneTimeTokenByRelatesTo() {
+	ts.Run("returns the row matching relates_to and token type", func() {
+		TruncateAll(ts.db)
+		u := ts.createUser()
+		require.NoError(ts.T(), CreateOneTimeToken(ts.db, u.ID, "+15551234567", "hash", PhoneChangeToken, time.Minute))
+
+		ott, err := FindOneTimeTokenByRelatesTo(ts.db, "+15551234567", PhoneChangeToken)
+		require.NoError(ts.T(), err)
+		require.Equal(ts.T(), "hash", ott.TokenHash)
+		require.Equal(ts.T(), u.ID, ott.UserID)
+		require.Equal(ts.T(), PhoneChangeToken, ott.TokenType)
+	})
+
+	ts.Run("lowercases relates_to before matching", func() {
+		TruncateAll(ts.db)
+		u := ts.createUser()
+		require.NoError(ts.T(), CreateOneTimeToken(ts.db, u.ID, "User@Example.com", "hash", ConfirmationToken, time.Minute))
+
+		ott, err := FindOneTimeTokenByRelatesTo(ts.db, "USER@EXAMPLE.COM", ConfirmationToken)
+		require.NoError(ts.T(), err)
+		require.Equal(ts.T(), "user@example.com", ott.RelatesTo)
+		require.Equal(ts.T(), u.ID, ott.UserID)
+	})
+
+	ts.Run("does not leak across token types", func() {
+		TruncateAll(ts.db)
+		u := ts.createUser()
+		require.NoError(ts.T(), CreateOneTimeToken(ts.db, u.ID, "+15551234567", "hash", PhoneChangeToken, time.Minute))
+
+		ott, err := FindOneTimeTokenByRelatesTo(ts.db, "+15551234567", ConfirmationToken)
+		require.True(ts.T(), IsNotFoundError(err), "expected not found error, got %v", err)
+		require.Nil(ts.T(), ott)
+	})
+
+	ts.Run("not found when no row has the relates_to value", func() {
+		TruncateAll(ts.db)
+		u := ts.createUser()
+		require.NoError(ts.T(), CreateOneTimeToken(ts.db, u.ID, "+15551234567", "hash", PhoneChangeToken, time.Minute))
+
+		ott, err := FindOneTimeTokenByRelatesTo(ts.db, "+15559999999", PhoneChangeToken)
+		require.True(ts.T(), IsNotFoundError(err), "expected not found error, got %v", err)
+		require.Nil(ts.T(), ott)
+	})
+
+	ts.Run("returns the newest row when two users share the value", func() {
+		TruncateAll(ts.db)
+		first := ts.createUser()
+
+		second, err := NewUser("", "other@example.com", "password", ts.config.JWT.Aud, nil)
+		require.NoError(ts.T(), err)
+		require.NoError(ts.T(), ts.db.Create(second))
+
+		require.NoError(ts.T(), CreateOneTimeToken(ts.db, first.ID, "+15551234567", "first-hash", PhoneChangeToken, time.Minute))
+		require.NoError(ts.T(), CreateOneTimeToken(ts.db, second.ID, "+15551234567", "second-hash", PhoneChangeToken, time.Minute))
+
+		ott, err := FindOneTimeTokenByRelatesTo(ts.db, "+15551234567", PhoneChangeToken)
+		require.NoError(ts.T(), err)
+		require.Equal(ts.T(), "second-hash", ott.TokenHash)
+		require.Equal(ts.T(), second.ID, ott.UserID)
+	})
+}