From a79e4dc5925e99130ec6d0b837f85ab14d46248a Mon Sep 17 00:00:00 2001
From: James Elliott <james-d-elliott@users.noreply.github.com>
Date: Thu, 2 Dec 2021 15:16:45 +1100
Subject: [PATCH] fix(storage): duo/u2f upsert failure on postgresql (#2658)

This replaces the standard duo_devices upsert with a PostgreSQL specific one and ensures the u2f_devices upsert uses the new unique key for the ON CONFLICT check.
---
 .../handlers/handler_register_u2f_step2.go     |  7 ++++---
 internal/storage/sql_provider.go               |  4 ++--
 .../storage/sql_provider_backend_postgres.go   |  1 +
 internal/storage/sql_provider_queries.go       | 18 ++++++++++++------
 4 files changed, 19 insertions(+), 11 deletions(-)

diff --git a/internal/handlers/handler_register_u2f_step2.go b/internal/handlers/handler_register_u2f_step2.go
index dd2b70ff9..251242112 100644
--- a/internal/handlers/handler_register_u2f_step2.go
+++ b/internal/handlers/handler_register_u2f_step2.go
@@ -48,9 +48,10 @@ func SecondFactorU2FRegister(ctx *middlewares.AutheliaCtx) {
 	publicKey := elliptic.Marshal(elliptic.P256(), registration.PubKey.X, registration.PubKey.Y)
 
 	err = ctx.Providers.StorageProvider.SaveU2FDevice(ctx, models.U2FDevice{
-		Username:  userSession.Username,
-		KeyHandle: registration.KeyHandle,
-		PublicKey: publicKey},
+		Username:    userSession.Username,
+		Description: "Primary",
+		KeyHandle:   registration.KeyHandle,
+		PublicKey:   publicKey},
 	)
 
 	if err != nil {
diff --git a/internal/storage/sql_provider.go b/internal/storage/sql_provider.go
index 6a983bd79..816f39f18 100644
--- a/internal/storage/sql_provider.go
+++ b/internal/storage/sql_provider.go
@@ -339,8 +339,8 @@ func (p *SQLProvider) UpdateTOTPConfigurationSecret(ctx context.Context, config
 
 // SaveU2FDevice saves a registered U2F device.
 func (p *SQLProvider) SaveU2FDevice(ctx context.Context, device models.U2FDevice) (err error) {
-	if _, err = p.db.ExecContext(ctx, p.sqlUpsertU2FDevice, device.Username, device.KeyHandle, device.PublicKey); err != nil {
-		return fmt.Errorf("error upserting U2F device secret: %v", err)
+	if _, err = p.db.ExecContext(ctx, p.sqlUpsertU2FDevice, device.Username, device.Description, device.KeyHandle, device.PublicKey); err != nil {
+		return fmt.Errorf("error upserting U2F device: %v", err)
 	}
 
 	return nil
diff --git a/internal/storage/sql_provider_backend_postgres.go b/internal/storage/sql_provider_backend_postgres.go
index b5e475680..5eadf8539 100644
--- a/internal/storage/sql_provider_backend_postgres.go
+++ b/internal/storage/sql_provider_backend_postgres.go
@@ -27,6 +27,7 @@ func NewPostgreSQLProvider(config *schema.Configuration) (provider *PostgreSQLPr
 	// Specific alterations to this provider.
 	// PostgreSQL doesn't have a UPSERT statement but has an ON CONFLICT operation instead.
 	provider.sqlUpsertU2FDevice = fmt.Sprintf(queryFmtPostgresUpsertU2FDevice, tableU2FDevices)
+	provider.sqlUpsertDuoDevice = fmt.Sprintf(queryFmtPostgresUpsertDuoDevice, tableDuoDevices)
 	provider.sqlUpsertTOTPConfig = fmt.Sprintf(queryFmtPostgresUpsertTOTPConfiguration, tableTOTPConfigurations)
 	provider.sqlUpsertPreferred2FAMethod = fmt.Sprintf(queryFmtPostgresUpsertPreferred2FAMethod, tableUserPreferences)
 	provider.sqlUpsertEncryptionValue = fmt.Sprintf(queryFmtPostgresUpsertEncryptionValue, tableEncryption)
diff --git a/internal/storage/sql_provider_queries.go b/internal/storage/sql_provider_queries.go
index 1bf1a0b30..9b6e66d4d 100644
--- a/internal/storage/sql_provider_queries.go
+++ b/internal/storage/sql_provider_queries.go
@@ -119,14 +119,14 @@ const (
 		WHERE username = ?;`
 
 	queryFmtUpsertU2FDevice = `
-		REPLACE INTO %s (username, key_handle, public_key)
-		VALUES (?, ?, ?);`
+		REPLACE INTO %s (username, description, key_handle, public_key)
+		VALUES (?, ?, ?, ?);`
 
 	queryFmtPostgresUpsertU2FDevice = `
-		INSERT INTO %s (username, key_handle, public_key)
-		VALUES ($1, $2, $3)
-			ON CONFLICT (username)
-			DO UPDATE SET key_handle=$2, public_key=$3;`
+		INSERT INTO %s (username, description, key_handle, public_key)
+		VALUES ($1, $2, $3, $4)
+			ON CONFLICT (username, description)
+			DO UPDATE SET key_handle=$3, public_key=$4;`
 )
 
 const (
@@ -134,6 +134,12 @@ const (
 		REPLACE INTO %s (username, device, method)
 		VALUES (?, ?, ?);`
 
+	queryFmtPostgresUpsertDuoDevice = `
+		INSERT INTO %s (username, device, method)
+		VALUES ($1, $2, $3)
+			ON CONFLICT (username)
+			DO UPDATE SET device=$2, method=$3;`
+
 	queryFmtDeleteDuoDevice = `
 		DELETE
 		FROM %s