From 9e8b10215b2d86313524f29c5881adc5c259949e Mon Sep 17 00:00:00 2001
From: ethan-l-geotab <ethanliong@geotab.com>
Date: Mon, 4 May 2026 17:17:34 -0400
Subject: [PATCH] fix: address second round of review feedback for valve
 encryption

- Use .iterator() instead of .select() in migrate/rollback for memory efficiency
- Fix unsafe rollback: use _fernet directly so decryption failures raise and
  leave the encrypted value intact (instead of silently overwriting with {})
- Validate decrypted result is a dict before returning
- Remove misleading log message claiming encrypted value is preserved after failure

Co-Authored-By: Claude Opus 4 (1M context) <noreply@anthropic.com>
---
 .../migrations/019_encrypt_user_valves.py         | 15 ++++++++-------
 backend/open_webui/utils/valve_encryption.py      | 14 ++++++++++----
 2 files changed, 18 insertions(+), 11 deletions(-)

diff --git a/backend/open_webui/internal/migrations/019_encrypt_user_valves.py b/backend/open_webui/internal/migrations/019_encrypt_user_valves.py
index 009a3cc2a9..7c0361cc22 100644
--- a/backend/open_webui/internal/migrations/019_encrypt_user_valves.py
+++ b/backend/open_webui/internal/migrations/019_encrypt_user_valves.py
@@ -46,7 +46,7 @@ def migrate(migrator: Migrator, database: pw.Database, *, fake=False):
     """Encrypt existing plaintext user valves in user.settings."""
     User = migrator.orm["user"]
 
-    for user in User.select().where(User.settings.is_null(False)):
+    for user in User.select().where(User.settings.is_null(False)).iterator():
         settings_raw = user.settings
         if not settings_raw:
             continue
@@ -88,11 +88,11 @@ def migrate(migrator: Migrator, database: pw.Database, *, fake=False):
 
 def rollback(migrator: Migrator, database: pw.Database, *, fake=False):
     """Decrypt user valves back to plaintext."""
-    from open_webui.utils.valve_encryption import decrypt_user_valves
+    from open_webui.utils.valve_encryption import _fernet
 
     User = migrator.orm["user"]
 
-    for user in User.select().where(User.settings.is_null(False)):
+    for user in User.select().where(User.settings.is_null(False)).iterator():
         settings_raw = user.settings
         if not settings_raw:
             continue
@@ -122,13 +122,14 @@ def rollback(migrator: Migrator, database: pw.Database, *, fake=False):
             for valve_id, valve_data in valves.items():
                 if isinstance(valve_data, str):
                     try:
-                        settings[category]["valves"][valve_id] = decrypt_user_valves(
-                            valve_data
-                        )
+                        decrypted = json.loads(_fernet.decrypt(valve_data.encode()).decode())
+                        if not isinstance(decrypted, dict):
+                            raise ValueError(f"Expected dict, got {type(decrypted).__name__}")
+                        settings[category]["valves"][valve_id] = decrypted
                         changed = True
                     except Exception:
                         log.warning(
-                            "Rollback failed to decrypt valve %s/%s for user %s",
+                            "Rollback failed to decrypt valve %s/%s for user %s; skipping",
                             category, valve_id, user.id,
                         )
 
diff --git a/backend/open_webui/utils/valve_encryption.py b/backend/open_webui/utils/valve_encryption.py
index 45389d5881..92c24428ea 100644
--- a/backend/open_webui/utils/valve_encryption.py
+++ b/backend/open_webui/utils/valve_encryption.py
@@ -45,17 +45,23 @@ def decrypt_user_valves(stored) -> dict:
     if isinstance(stored, str):
         try:
             decrypted = _fernet.decrypt(stored.encode()).decode()
-            return json.loads(decrypted)
+            result = json.loads(decrypted)
+            if not isinstance(result, dict):
+                log.error(
+                    "Decrypted user valves produced unexpected type %s; returning empty valves.",
+                    type(result).__name__,
+                )
+                return {}
+            return result
         except InvalidToken:
             log.error(
                 "Failed to decrypt user valves: key mismatch or corrupted data. "
-                "Returning empty valves. The original encrypted value is preserved in the database."
+                "Returning empty valves."
             )
             return {}
         except json.JSONDecodeError:
             log.error(
-                "Decrypted user valves but got malformed JSON. "
-                "Returning empty valves. The original encrypted value is preserved in the database."
+                "Decrypted user valves but got malformed JSON. Returning empty valves."
             )
             return {}
         except Exception as e: