fix(sql): replace non-atomic check-then-insert with upsert in WalletStore.StoreIdentity

The previous implementation performed a non-atomic IdentityExists check
followed by a separate INSERT, creating a TOCTOU race under concurrent
writers. Replace with a single atomic INSERT ... ON CONFLICT DO NOTHING,
which is natively supported by both PostgreSQL and SQLite.

Signed-off-by: Rama542 <Rama542@users.noreply.github.com>

Author: Rama542

token/services/storage/db/sql/common/wallet.go

@@ -87,14 +87,10 @@ func (db *WalletStore) GetWalletIDs(ctx context.Context, roleID int) ([]driver.W
 }
 
 func (db *WalletStore) StoreIdentity(ctx context.Context, identity token.Identity, eID string, wID driver.WalletID, roleID int, meta []byte) error {
-	// TODO AF Use upsert
-	if db.IdentityExists(ctx, identity, wID, roleID) {
-		return nil
-	}
-
 	query, args := q.InsertInto(db.table.Wallets).
 		Fields("identity_hash", "meta", "wallet_id", "role_id", "created_at", "enrollment_id").
 		Row(identity.UniqueID(), meta, wID, roleID, time.Now().UTC(), eID).
+		OnConflictDoNothing().
 		Format()
 	logging.Debug(logger, query)
 

token/services/storage/db/sql/common/wallet_test_utils.go

@@ -107,14 +107,9 @@ func TestStoreIdentity(t *testing.T, store walletStoreConstructor) {
 	walletID := driver.WalletID("my wallet")
 	roleID := 5
 
-	mockDB.
-		ExpectQuery("SELECT wallet_id FROM WALLETS WHERE \\(identity_hash = \\$1\\) AND \\(wallet_id = \\$2\\) AND \\(role_id = \\$3\\)").
-		WithArgs(tokenID.UniqueID(), walletID, roleID).
-		WillReturnRows(mockDB.NewRows([]string{"wallet_id"}))
-
 	mockDB.ExpectExec("INSERT INTO WALLETS "+
 		"\\(identity_hash, meta, wallet_id, role_id, created_at, enrollment_id\\) "+
-		"VALUES \\(\\$1, \\$2, \\$3, \\$4, \\$5, \\$6\\)").
+		"VALUES \\(\\$1, \\$2, \\$3, \\$4, \\$5, \\$6\\) ON CONFLICT DO NOTHING").
 		WithArgs(tokenID.UniqueID(), []uint8(nil), walletID, roleID, sqlmock.AnyArg(), eID).
 		WillReturnResult(sqlmock.NewResult(1, 1))
 
@@ -123,3 +118,37 @@ func TestStoreIdentity(t *testing.T, store walletStoreConstructor) {
 	gomega.Expect(mockDB.ExpectationsWereMet()).To(gomega.Succeed())
 	gomega.Expect(err).ToNot(gomega.HaveOccurred())
 }
+
+func TestStoreIdentityIdempotent(t *testing.T, store walletStoreConstructor) {
+	gomega.RegisterTestingT(t)
+	db, mockDB, err := sqlmock.New()
+	gomega.Expect(err).ToNot(gomega.HaveOccurred())
+
+	tokenID := token.Identity([]byte("1234"))
+	eID := "5678"
+	walletID := driver.WalletID("my wallet")
+	roleID := 5
+
+	insertQuery := "INSERT INTO WALLETS " +
+		"\\(identity_hash, meta, wallet_id, role_id, created_at, enrollment_id\\) " +
+		"VALUES \\(\\$1, \\$2, \\$3, \\$4, \\$5, \\$6\\) ON CONFLICT DO NOTHING"
+
+	// First call: row inserted (1 row affected)
+	mockDB.ExpectExec(insertQuery).
+		WithArgs(tokenID.UniqueID(), []uint8(nil), walletID, roleID, sqlmock.AnyArg(), eID).
+		WillReturnResult(sqlmock.NewResult(1, 1))
+
+	// Second call: conflict, 0 rows affected — must still return nil
+	mockDB.ExpectExec(insertQuery).
+		WithArgs(tokenID.UniqueID(), []uint8(nil), walletID, roleID, sqlmock.AnyArg(), eID).
+		WillReturnResult(sqlmock.NewResult(0, 0))
+
+	s := store(db)
+	err = s.StoreIdentity(t.Context(), tokenID, eID, walletID, roleID, nil)
+	gomega.Expect(err).ToNot(gomega.HaveOccurred())
+
+	err = s.StoreIdentity(t.Context(), tokenID, eID, walletID, roleID, nil)
+	gomega.Expect(err).ToNot(gomega.HaveOccurred())
+
+	gomega.Expect(mockDB.ExpectationsWereMet()).To(gomega.Succeed())
+}

token/services/storage/db/sql/postgres/wallet_test.go

@@ -40,3 +40,7 @@ func TestIdentityExists(t *testing.T) {
 func TestStoreIdentity(t *testing.T) {
 	common2.TestStoreIdentity(t, mockWalletStore)
 }
+
+func TestStoreIdentityIdempotent(t *testing.T) {
+	common2.TestStoreIdentityIdempotent(t, mockWalletStore)
+}

token/services/storage/db/sql/sqlite/wallet_test.go

@@ -40,3 +40,7 @@ func TestIdentityExists(t *testing.T) {
 func TestStoreIdentity(t *testing.T) {
 	common2.TestStoreIdentity(t, mockWalletStore)
 }
+
+func TestStoreIdentityIdempotent(t *testing.T) {
+	common2.TestStoreIdentityIdempotent(t, mockWalletStore)
+}