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>
This commit is contained in:
ethan-l-geotab 2026-05-04 17:17:34 -04:00
parent a789201bfe
commit 9e8b10215b
2 changed files with 18 additions and 11 deletions

View file

@ -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,
)

View file

@ -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: