chore(kratos): rewrite recovery-address lookup as region-pinned CTE

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
GitOrigin-RevId: 4c9def2adf32a1528bddef30ac6a04d5e7143941

Author: 2 people

identity/pool.go

@@ -116,8 +116,8 @@ type (
 		// FindRecoveryAddressByValue returns a matching address or sql.ErrNoRows if no address could be found.
 		FindRecoveryAddressByValue(ctx context.Context, via, address string) (*RecoveryAddress, error)
 
-		// FindAllRecoveryAddressesForIdentityByRecoveryAddressValue finds all recovery addresses for an identity if at least one of its recovery addresses matches the provided value.
-		FindAllRecoveryAddressesForIdentityByRecoveryAddressValue(ctx context.Context, anyRecoveryAddress string) ([]RecoveryAddress, error)
+		// FindAllRecoveryAddressValuesForIdentityByRecoveryAddressValue finds the values of all recovery addresses for an identity if at least one of its recovery addresses matches the provided value.
+		FindAllRecoveryAddressValuesForIdentityByRecoveryAddressValue(ctx context.Context, anyRecoveryAddress string) ([]string, error)
 	}
 
 	PoolProvider interface {

identity/test/pool.go

@@ -9,7 +9,6 @@ import (
 	"encoding/json"
 	"fmt"
 	"net/http"
-	"slices"
 	"strings"
 	"testing"
 	"time"
@@ -2021,12 +2020,6 @@ func TestPool(ctx context.Context, p persistence.Persister, m *identity.Manager,
 		})
 
 		t.Run("suite=recovery-address", func(t *testing.T) {
-			sortAddresses := func(addresses []identity.RecoveryAddress) {
-				slices.SortFunc(addresses, func(a, b identity.RecoveryAddress) int {
-					return strings.Compare(a.Value, b.Value)
-				})
-			}
-
 			createIdentityWithAddresses := func(t *testing.T, email string) *identity.Identity {
 				var i identity.Identity
 				require.NoError(t, faker.FakeData(&i))
@@ -2045,7 +2038,7 @@ func TestPool(ctx context.Context, p persistence.Persister, m *identity.Manager,
 				_, err := p.FindRecoveryAddressByValue(ctx, identity.AddressTypeEmail, "does-not-exist")
 				require.ErrorIs(t, err, sqlcon.ErrNoRows())
 
-				allAddresses, err := p.FindAllRecoveryAddressesForIdentityByRecoveryAddressValue(ctx, "does-not-exist")
+				allAddresses, err := p.FindAllRecoveryAddressValuesForIdentityByRecoveryAddressValue(ctx, "does-not-exist")
 				require.NoError(t, err)
 				require.Len(t, allAddresses, 0)
 			})
@@ -2080,19 +2073,18 @@ func TestPool(ctx context.Context, p persistence.Persister, m *identity.Manager,
 						})
 					})
 
-					t.Run("method=FindAllRecoveryAddressesForIdentityByRecoveryAddressValue", func(t *testing.T) {
+					t.Run("method=FindAllRecoveryAddressValuesForIdentityByRecoveryAddressValue", func(t *testing.T) {
 						t.Run(fmt.Sprintf("case=%d", k), func(t *testing.T) {
-							allAddresses, err := p.FindAllRecoveryAddressesForIdentityByRecoveryAddressValue(ctx, expected.Value)
+							allAddresses, err := p.FindAllRecoveryAddressValuesForIdentityByRecoveryAddressValue(ctx, expected.Value)
 							require.NoError(t, err)
 							require.Len(t, allAddresses, 2)
-							sortAddresses(allAddresses)
-							require.Equal(t, expected.Value, allAddresses[0].Value)
-							require.Equal(t, expected.Value+"_other", allAddresses[1].Value)
+							require.Equal(t, expected.Value, allAddresses[0])
+							require.Equal(t, expected.Value+"_other", allAddresses[1])
 						})
 
 						t.Run("not if on another network", func(t *testing.T) {
 							_, p := testhelpers.NewNetwork(t, ctx, p)
-							allAddresses, err := p.FindAllRecoveryAddressesForIdentityByRecoveryAddressValue(ctx, expected.Value)
+							allAddresses, err := p.FindAllRecoveryAddressValuesForIdentityByRecoveryAddressValue(ctx, expected.Value)
 							require.NoError(t, err)
 							require.Len(t, allAddresses, 0)
 						})
@@ -2108,19 +2100,18 @@ func TestPool(ctx context.Context, p persistence.Persister, m *identity.Manager,
 				_, err := p.FindRecoveryAddressByValue(ctx, identity.AddressTypeEmail, email)
 				require.NoError(t, err)
 
-				allAddresses, err := p.FindAllRecoveryAddressesForIdentityByRecoveryAddressValue(ctx, emailLower)
+				allAddresses, err := p.FindAllRecoveryAddressValuesForIdentityByRecoveryAddressValue(ctx, emailLower)
 				require.NoError(t, err)
 				require.Len(t, allAddresses, 2)
-				sortAddresses(allAddresses)
-				require.Equal(t, allAddresses[0].Value, emailLower)
-				require.Equal(t, allAddresses[1].Value, emailLower+"_other")
+				require.Equal(t, allAddresses[0], emailLower)
+				require.Equal(t, allAddresses[1], emailLower+"_other")
 
 				t.Run("can not find if on another network", func(t *testing.T) {
 					_, p := testhelpers.NewNetwork(t, ctx, p)
 					_, err := p.FindRecoveryAddressByValue(ctx, identity.AddressTypeEmail, email)
 					require.ErrorIs(t, err, sqlcon.ErrNoRows())
 
-					allAddresses, err := p.FindAllRecoveryAddressesForIdentityByRecoveryAddressValue(ctx, emailLower)
+					allAddresses, err := p.FindAllRecoveryAddressValuesForIdentityByRecoveryAddressValue(ctx, emailLower)
 					require.NoError(t, err)
 					require.Len(t, allAddresses, 0)
 				})
@@ -2132,7 +2123,7 @@ func TestPool(ctx context.Context, p persistence.Persister, m *identity.Manager,
 				_, err = p.FindRecoveryAddressByValue(ctx, identity.AddressTypeEmail, email)
 				require.EqualError(t, err, sqlcon.ErrNoRows().Error())
 
-				allAddresses, err = p.FindAllRecoveryAddressesForIdentityByRecoveryAddressValue(ctx, emailLower)
+				allAddresses, err = p.FindAllRecoveryAddressValuesForIdentityByRecoveryAddressValue(ctx, emailLower)
 				require.NoError(t, err)
 				require.Len(t, allAddresses, 0)
 
@@ -2141,7 +2132,7 @@ func TestPool(ctx context.Context, p persistence.Persister, m *identity.Manager,
 					_, err := p.FindRecoveryAddressByValue(ctx, identity.AddressTypeEmail, email)
 					require.ErrorIs(t, err, sqlcon.ErrNoRows())
 
-					allAddresses, err := p.FindAllRecoveryAddressesForIdentityByRecoveryAddressValue(ctx, emailLower)
+					allAddresses, err := p.FindAllRecoveryAddressValuesForIdentityByRecoveryAddressValue(ctx, emailLower)
 					require.NoError(t, err)
 					require.Len(t, allAddresses, 0)
 				})
@@ -2152,21 +2143,18 @@ func TestPool(ctx context.Context, p persistence.Persister, m *identity.Manager,
 				assert.Equal(t, identity.AddressTypeEmail, actual.Via)
 				assert.Equal(t, emailNextLower, actual.Value)
 
-				allAddresses, err = p.FindAllRecoveryAddressesForIdentityByRecoveryAddressValue(ctx, emailNextLower)
+				allAddresses, err = p.FindAllRecoveryAddressValuesForIdentityByRecoveryAddressValue(ctx, emailNextLower)
 				require.NoError(t, err)
 				require.Len(t, allAddresses, 2)
-				sortAddresses(allAddresses)
-				assert.Equal(t, identity.AddressTypeEmail, allAddresses[0].Via)
-				assert.Equal(t, emailNextLower, allAddresses[0].Value)
-				assert.Equal(t, identity.AddressTypeEmail, allAddresses[1].Via)
-				assert.Equal(t, emailNextLower+"_other", allAddresses[1].Value)
+				assert.Equal(t, emailNextLower, allAddresses[0])
+				assert.Equal(t, emailNextLower+"_other", allAddresses[1])
 
 				t.Run("can not find if on another network", func(t *testing.T) {
 					_, p := testhelpers.NewNetwork(t, ctx, p)
 					_, err := p.FindRecoveryAddressByValue(ctx, identity.AddressTypeEmail, emailNext)
 					require.ErrorIs(t, err, sqlcon.ErrNoRows())
 
-					allAddresses, err := p.FindAllRecoveryAddressesForIdentityByRecoveryAddressValue(ctx, emailNextLower)
+					allAddresses, err := p.FindAllRecoveryAddressValuesForIdentityByRecoveryAddressValue(ctx, emailNextLower)
 					require.NoError(t, err)
 					require.Len(t, allAddresses, 0)
 				})
@@ -2367,30 +2355,25 @@ func TestPool(ctx context.Context, p persistence.Persister, m *identity.Manager,
 					createdIDs = append(createdIDs, i.ID)
 
 					// Test 1: Find all addresses using normalized phone1
-					allAddresses, err := p.FindAllRecoveryAddressesForIdentityByRecoveryAddressValue(ctx, normalizedPhone1)
+					allAddresses, err := p.FindAllRecoveryAddressValuesForIdentityByRecoveryAddressValue(ctx, normalizedPhone1)
 					require.NoError(t, err, "should find all recovery addresses with normalized phone")
 					require.Len(t, allAddresses, 2, "should return all recovery addresses for the identity")
-					sortAddresses(allAddresses)
-					assert.Equal(t, normalizedPhone1, allAddresses[0].Value)
-					assert.Equal(t, normalizedPhone2, allAddresses[1].Value)
-					assert.Equal(t, identity.AddressTypeSMS, allAddresses[0].Via)
-					assert.Equal(t, identity.AddressTypeSMS, allAddresses[1].Via)
+					assert.Equal(t, normalizedPhone1, allAddresses[0])
+					assert.Equal(t, normalizedPhone2, allAddresses[1])
 
 					// Test 2: Find all addresses using non-normalized phone1 (gets normalized, then matches)
-					allAddresses2, err := p.FindAllRecoveryAddressesForIdentityByRecoveryAddressValue(ctx, nonNormalizedPhone1)
+					allAddresses2, err := p.FindAllRecoveryAddressValuesForIdentityByRecoveryAddressValue(ctx, nonNormalizedPhone1)
 					require.NoError(t, err, "should find all recovery addresses with non-normalized phone")
 					require.Len(t, allAddresses2, 2, "should return all recovery addresses when querying with non-normalized phone")
-					sortAddresses(allAddresses2)
-					assert.Equal(t, normalizedPhone1, allAddresses2[0].Value)
-					assert.Equal(t, normalizedPhone2, allAddresses2[1].Value)
+					assert.Equal(t, normalizedPhone1, allAddresses2[0])
+					assert.Equal(t, normalizedPhone2, allAddresses2[1])
 
 					// Test 3: Find all addresses using normalized phone2
-					allAddresses3, err := p.FindAllRecoveryAddressesForIdentityByRecoveryAddressValue(ctx, normalizedPhone2)
+					allAddresses3, err := p.FindAllRecoveryAddressValuesForIdentityByRecoveryAddressValue(ctx, normalizedPhone2)
 					require.NoError(t, err)
 					require.Len(t, allAddresses3, 2, "should return all recovery addresses when querying with second phone")
-					sortAddresses(allAddresses3)
-					assert.Equal(t, normalizedPhone1, allAddresses3[0].Value)
-					assert.Equal(t, normalizedPhone2, allAddresses3[1].Value)
+					assert.Equal(t, normalizedPhone1, allAddresses3[0])
+					assert.Equal(t, normalizedPhone2, allAddresses3[1])
 				})
 
 				t.Run("create with multiple SMS addresses and find all with non-normalized phone", func(t *testing.T) {
@@ -2413,22 +2396,18 @@ func TestPool(ctx context.Context, p persistence.Persister, m *identity.Manager,
 					createdIDs = append(createdIDs, i.ID)
 
 					// Test 1: Find all addresses using non-normalized phone1 (gets normalized for search)
-					allAddresses, err := p.FindAllRecoveryAddressesForIdentityByRecoveryAddressValue(ctx, nonNormPhone1)
+					allAddresses, err := p.FindAllRecoveryAddressValuesForIdentityByRecoveryAddressValue(ctx, nonNormPhone1)
 					require.NoError(t, err, "should find all recovery addresses with non-normalized phone")
 					require.Len(t, allAddresses, 2, "should return all recovery addresses for the identity")
-					sortAddresses(allAddresses)
-					assert.Equal(t, normPhone1, allAddresses[0].Value, "stored values should be normalized")
-					assert.Equal(t, normPhone2, allAddresses[1].Value)
-					assert.Equal(t, identity.AddressTypeSMS, allAddresses[0].Via)
-					assert.Equal(t, identity.AddressTypeSMS, allAddresses[1].Via)
+					assert.Equal(t, normPhone1, allAddresses[0], "stored values should be normalized")
+					assert.Equal(t, normPhone2, allAddresses[1])
 
 					// Test 2: Find all addresses using normalized phone1 (direct match)
-					allAddresses2, err := p.FindAllRecoveryAddressesForIdentityByRecoveryAddressValue(ctx, normPhone1)
+					allAddresses2, err := p.FindAllRecoveryAddressValuesForIdentityByRecoveryAddressValue(ctx, normPhone1)
 					require.NoError(t, err)
 					require.Len(t, allAddresses2, 2, "should return all recovery addresses when querying with normalized phone")
-					sortAddresses(allAddresses2)
-					assert.Equal(t, normPhone1, allAddresses2[0].Value)
-					assert.Equal(t, normPhone2, allAddresses2[1].Value)
+					assert.Equal(t, normPhone1, allAddresses2[0])
+					assert.Equal(t, normPhone2, allAddresses2[1])
 				})
 
 				t.Run("backward compatibility with legacy non-normalized phone data", func(t *testing.T) {
@@ -2470,14 +2449,11 @@ func TestPool(ctx context.Context, p persistence.Persister, m *identity.Manager,
 
 					// Test: Find all addresses with exact non-normalized phone (what user would enter)
 					// Should find all recovery addresses for the identity
-					allAddresses, err := p.FindAllRecoveryAddressesForIdentityByRecoveryAddressValue(ctx, legacyPhone1)
+					allAddresses, err := p.FindAllRecoveryAddressValuesForIdentityByRecoveryAddressValue(ctx, legacyPhone1)
 					require.NoError(t, err, "should find all recovery addresses with non-normalized search")
 					require.Len(t, allAddresses, 2, "should return all recovery addresses for the identity")
-					sortAddresses(allAddresses)
-					assert.Equal(t, legacyPhone1, allAddresses[0].Value)
-					assert.Equal(t, legacyPhone2, allAddresses[1].Value)
-					assert.Equal(t, identity.AddressTypeSMS, allAddresses[0].Via)
-					assert.Equal(t, identity.AddressTypeSMS, allAddresses[1].Via)
+					assert.Equal(t, legacyPhone1, allAddresses[0])
+					assert.Equal(t, legacyPhone2, allAddresses[1])
 				})
 			})
 		})

persistence/sql/identity/persister_identity.go

@@ -1541,11 +1541,11 @@ func (p *IdentityPersister) FindRecoveryAddressByValue(ctx context.Context, via,
 	return &addr, nil
 }
 
-// FindAllRecoveryAddressesForIdentityByRecoveryAddressValue returns all
-// recovery addresses for an identity if at least one of those addresses matches
-// the provided value.
-func (p *IdentityPersister) FindAllRecoveryAddressesForIdentityByRecoveryAddressValue(ctx context.Context, anyRecoveryAddress string) (recoveryAddresses []identity.RecoveryAddress, err error) {
-	ctx, span := p.r.Tracer(ctx).Tracer().Start(ctx, "persistence.sql.FindAllRecoveryAddressesForIdentityByRecoveryAddressValue",
+// FindAllRecoveryAddressValuesForIdentityByRecoveryAddressValue returns the
+// values of all recovery addresses for an identity if at least one of those
+// addresses matches the provided value.
+func (p *IdentityPersister) FindAllRecoveryAddressValuesForIdentityByRecoveryAddressValue(ctx context.Context, anyRecoveryAddress string) (recoveryAddresses []string, err error) {
+	ctx, span := p.r.Tracer(ctx).Tracer().Start(ctx, "persistence.sql.FindAllRecoveryAddressValuesForIdentityByRecoveryAddressValue",
 		trace.WithAttributes(
 			attribute.Stringer("network.id", p.NetworkID(ctx))))
 	defer otelx.End(span, &err)
@@ -1560,19 +1560,28 @@ func (p *IdentityPersister) FindAllRecoveryAddressesForIdentityByRecoveryAddress
 	//
 	// This is all done in one query with a self-join.
 	// We also bound the results for safety.
+	via := identity.AddressTypeSMS
+	if strings.ContainsRune(anyRecoveryAddress, '@') {
+		via = identity.AddressTypeEmail
+	}
+
+	nid := p.NetworkID(ctx)
 	err = p.GetConnection(ctx).RawQuery(`
-SELECT A.id, A.via, A.value, A.identity_id, A.created_at, A.updated_at, A.nid, A.break_glass_for_organization
+SELECT A.value
 FROM identity_recovery_addresses A
 JOIN identity_recovery_addresses B
-ON A.identity_id = B.identity_id
-AND A.nid = B.nid
-WHERE B.value IN (?,?)
-AND A.nid = ?
+  ON A.identity_id = B.identity_id
+  AND A.nid = B.nid
+WHERE A.nid = ?
+  AND B.via = ?
+  AND B.value IN (?, ?)
+ORDER BY A.value
 LIMIT 10
 		`,
+		nid,
+		via,
 		x.GracefulNormalization(anyRecoveryAddress),
 		anyRecoveryAddress,
-		p.NetworkID(ctx),
 	).
 		All(&recoveryAddresses)
 	if err != nil {

selfservice/strategy/code/strategy_recovery.go

@@ -481,7 +481,7 @@ func (s *Strategy) recoveryV2HandleStateAwaitingAddress(r *http.Request, f *reco
 	}
 
 	// Need to retrieve all possible recovery addresses and present a choice.
-	recoveryAddresses, err := s.deps.IdentityPool().FindAllRecoveryAddressesForIdentityByRecoveryAddressValue(r.Context(), body.RecoveryAddress)
+	recoveryAddresses, err := s.deps.IdentityPool().FindAllRecoveryAddressValuesForIdentityByRecoveryAddressValue(r.Context(), body.RecoveryAddress)
 	// Real error.
 	if err != nil && !errors.Is(err, sqlcon.ErrNoRows()) {
 		return err
@@ -491,12 +491,12 @@ func (s *Strategy) recoveryV2HandleStateAwaitingAddress(r *http.Request, f *reco
 	if len(recoveryAddresses) == 0 {
 		// To avoid an attacker from using this case to probe for existing addresses, we pretend it exists.
 		// This is the same behavior as in Recovery V1.
-		recoveryAddresses = append(recoveryAddresses, identity.RecoveryAddress{Value: body.RecoveryAddress})
+		recoveryAddresses = append(recoveryAddresses, body.RecoveryAddress)
 	}
 
 	f.State = flow.StateRecoveryAwaitingAddressChoice
 
-	if len(recoveryAddresses) == 1 && recoveryAddresses[0].Value == body.RecoveryAddress {
+	if len(recoveryAddresses) == 1 && recoveryAddresses[0] == body.RecoveryAddress {
 		// Skip two states for convenience:
 		// - No need to present a choice with only one option
 		// - No need to ask for the full address if there is only one and it was just provided in full
@@ -522,20 +522,18 @@ func (s *Strategy) recoveryV2HandleStateAwaitingAddress(r *http.Request, f *reco
 	f.State = flow.StateRecoveryAwaitingAddressChoice
 	f.UI.Messages.Set(text.NewRecoveryAskToChooseAddress())
 
-	slices.SortFunc(recoveryAddresses, func(a, b identity.RecoveryAddress) int {
-		return strings.Compare(a.Value, b.Value)
-	})
+	slices.Sort(recoveryAddresses)
 
 	for _, a := range recoveryAddresses {
 		// NOTE: Only send the masked value and the hash, to avoid information exfiltration.
 		// Why the hash? So that we can recognize later, when the user chooses the masked address in the list,
 		// that the chosen masked address is the `recovery_address` provided in the beginning,
 		// and then we do not ask again the user to provide it in full.
-		hashBase64 := AddressToHashBase64(a.Value)
+		hashBase64 := AddressToHashBase64(a)
 		f.UI.GetNodes().Append(node.NewInputField("recovery_select_address", hashBase64, node.CodeGroup, node.InputAttributeTypeSubmit).
 			WithMetaLabel(&text.Message{
 				ID:   text.InfoNodeLabel,
-				Text: MaskAddress(a.Value),
+				Text: MaskAddress(a),
 				Type: text.Info,
 			}))
 	}
@@ -592,11 +590,11 @@ func (s *Strategy) recoveryV2HandleStateAwaitingAddressChoice(r *http.Request, f
 
 	// Retrieve the selected recovery address in plaintext to determine the input label and type.
 	var plaintextRecoveryAddress string
-	recoveryAddresses, err := s.deps.IdentityPool().FindAllRecoveryAddressesForIdentityByRecoveryAddressValue(r.Context(), body.RecoveryAddress)
+	recoveryAddresses, err := s.deps.IdentityPool().FindAllRecoveryAddressValuesForIdentityByRecoveryAddressValue(r.Context(), body.RecoveryAddress)
 	if err == nil {
 		for _, a := range recoveryAddresses {
-			if subtle.ConstantTimeCompare([]byte(AddressToHashBase64(a.Value)), []byte(body.RecoverySelectAddress)) == 1 {
-				plaintextRecoveryAddress = a.Value
+			if subtle.ConstantTimeCompare([]byte(AddressToHashBase64(a)), []byte(body.RecoverySelectAddress)) == 1 {
+				plaintextRecoveryAddress = a
 				break
 			}
 		}
@@ -728,7 +726,7 @@ func (s *Strategy) recoveryV2HandleStateAwaitingCode(w http.ResponseWriter, r *h
 
 func (s *Strategy) recoveryV2HandleGoBack(r *http.Request, f *recovery.Flow, body *recoverySubmitPayload) error {
 	// If no address choice needs to take place, just go to the first screen.
-	recoveryAddresses, _ := s.deps.IdentityPool().FindAllRecoveryAddressesForIdentityByRecoveryAddressValue(r.Context(), body.RecoveryAddress)
+	recoveryAddresses, _ := s.deps.IdentityPool().FindAllRecoveryAddressValuesForIdentityByRecoveryAddressValue(r.Context(), body.RecoveryAddress)
 	if len(recoveryAddresses) <= 1 {
 		f.State = flow.StateRecoveryAwaitingAddress
 		err := s.PopulateRecoveryMethod(r, f)