fix(matrix): migrate crypto store when the Olm pickle key changes
The Olm account pickle key is derived from the account ID plus the configured device ID (acct:device_id). If the crypto store's account was created before MATRIX_DEVICE_ID was set — e.g. the very first password login, where the device ID is only known after connecting — it gets pickled under "<acct>:default". Setting MATRIX_DEVICE_ID afterwards (a reasonable thing to do once you know the device ID you want to pin) changes the derived pickle key, and every subsequent unpickle attempt fails with BAD_ACCOUNT_KEY. In optional-E2EE mode that failure is swallowed and encryption silently stays disabled instead of surfacing an actionable error. _migrate_legacy_crypto_pickle() detects the BAD_ACCOUNT_KEY failure, tries the known legacy pickle keys, and re-pickles the account (plus every stored olm/megolm session — sessions share the same pickle key, so migrating only the account would leave them unreadable on the next decrypt and silently break key sharing with peers) under the current key. It only reports failure when no known key can unpickle the account, with a log message pointing at what changed. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
ckaznocha
kshitij
parent
907bdc7856
commit
3a43422065
|
|
@ -1360,6 +1360,108 @@ class MatrixAdapter(BasePlatformAdapter):
|
|||
await crypto_store.delete()
|
||||
return True
|
||||
|
||||
async def _migrate_legacy_crypto_pickle(
|
||||
self, crypto_store: Any, crypto_db: Any, acct_id: str, pickle_key: str
|
||||
) -> bool:
|
||||
"""Re-pickle the Olm account when the pickle key changed.
|
||||
|
||||
The pickle key embeds the *configured* device ID. If the account was
|
||||
created before MATRIX_DEVICE_ID was set (e.g. first password login,
|
||||
where the device ID is only known after connecting), the store was
|
||||
pickled under ``<acct>:default``; setting MATRIX_DEVICE_ID afterwards
|
||||
changes the key and unpickling fails with BAD_ACCOUNT_KEY — which in
|
||||
optional-E2EE mode silently disables encryption. Try known legacy
|
||||
keys and re-pickle under the current one. Returns False only when an
|
||||
account exists but no key can unpickle it.
|
||||
"""
|
||||
try:
|
||||
await crypto_store.get_account()
|
||||
return True
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
from mautrix.crypto.store.asyncpg import PgCryptoStore
|
||||
|
||||
for legacy_key in (f"{acct_id}:default", acct_id):
|
||||
if legacy_key == pickle_key:
|
||||
continue
|
||||
legacy_store = PgCryptoStore(
|
||||
account_id=acct_id, pickle_key=legacy_key, db=crypto_db
|
||||
)
|
||||
try:
|
||||
account = await legacy_store.get_account()
|
||||
except Exception:
|
||||
continue
|
||||
if account is None:
|
||||
continue
|
||||
await crypto_store.put_account(account)
|
||||
await self._repickle_crypto_sessions(
|
||||
crypto_db, acct_id, legacy_key, pickle_key
|
||||
)
|
||||
logger.info(
|
||||
"Matrix: re-pickled crypto store account and sessions under "
|
||||
"the current pickle key (device ID was configured after the "
|
||||
"account was created)"
|
||||
)
|
||||
return True
|
||||
|
||||
logger.error(
|
||||
"Matrix: crypto store account exists but cannot be unpickled "
|
||||
"with the current or any legacy pickle key. If MATRIX_DEVICE_ID "
|
||||
"was changed manually, restore its previous value."
|
||||
)
|
||||
return False
|
||||
|
||||
async def _repickle_crypto_sessions(
|
||||
self, crypto_db: Any, acct_id: str, legacy_key: str, pickle_key: str
|
||||
) -> None:
|
||||
"""Re-pickle olm/megolm session blobs alongside the account.
|
||||
|
||||
The account and every stored session share the pickle key; migrating
|
||||
only the account leaves sessions unreadable (BAD_ACCOUNT_KEY on the
|
||||
next olm decrypt), which breaks key sharing with peers.
|
||||
"""
|
||||
import olm as olm_lib
|
||||
|
||||
tables = {
|
||||
"crypto_olm_session": olm_lib.Session,
|
||||
"crypto_megolm_inbound_session": olm_lib.InboundGroupSession,
|
||||
"crypto_megolm_outbound_session": olm_lib.OutboundGroupSession,
|
||||
}
|
||||
for table, session_cls in tables.items():
|
||||
rows = await crypto_db.fetch(
|
||||
f"SELECT session_id, session FROM {table} WHERE account_id=$1",
|
||||
acct_id,
|
||||
)
|
||||
for row in rows:
|
||||
blob = row["session"]
|
||||
if blob is None:
|
||||
continue
|
||||
pickled = bytes(blob)
|
||||
try:
|
||||
session_cls.from_pickle(pickled, pickle_key)
|
||||
continue # already readable with the current key
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
session = session_cls.from_pickle(pickled, legacy_key)
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Matrix: dropping unreadable %s row %s during pickle "
|
||||
"key migration: %s",
|
||||
table,
|
||||
row["session_id"],
|
||||
exc,
|
||||
)
|
||||
continue
|
||||
await crypto_db.execute(
|
||||
f"UPDATE {table} SET session=$1 "
|
||||
"WHERE account_id=$2 AND session_id=$3",
|
||||
session.pickle(pickle_key),
|
||||
acct_id,
|
||||
row["session_id"],
|
||||
)
|
||||
|
||||
async def _verify_device_keys_on_server(self, client: Any, olm: Any) -> bool:
|
||||
"""Verify our device keys are on the homeserver after loading crypto state.
|
||||
|
||||
|
|
@ -1669,6 +1771,10 @@ class MatrixAdapter(BasePlatformAdapter):
|
|||
)
|
||||
await crypto_store.put_device_id(client.device_id)
|
||||
|
||||
await self._migrate_legacy_crypto_pickle(
|
||||
crypto_store, crypto_db, _acct_id, _pickle_key
|
||||
)
|
||||
|
||||
crypto_state = _CryptoStateStore(state_store, self._joined_rooms, client)
|
||||
olm = OlmMachine(client, crypto_store, crypto_state)
|
||||
olm.share_keys_min_trust = TrustState.UNVERIFIED
|
||||
|
|
|
|||
|
|
@ -3099,3 +3099,92 @@ class TestCryptoStoreResetOnDeviceChange:
|
|||
assert "MATRIX_DEVICE_ID=DEVICE_A" in caplog.text
|
||||
|
||||
await adapter.disconnect()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Crypto store pickle-key migration
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestCryptoPickleKeyMigration:
|
||||
@pytest.mark.asyncio
|
||||
async def test_account_loads_fine_no_migration(self):
|
||||
adapter = _make_adapter()
|
||||
store = MagicMock()
|
||||
store.get_account = AsyncMock(return_value=MagicMock())
|
||||
assert await adapter._migrate_legacy_crypto_pickle(
|
||||
store, MagicMock(), "@bot:example.org", "@bot:example.org:DEV"
|
||||
) is True
|
||||
store.put_account.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_migrates_from_default_pickle_key(self, caplog):
|
||||
import logging
|
||||
adapter = _make_adapter()
|
||||
store = MagicMock()
|
||||
store.get_account = AsyncMock(side_effect=RuntimeError("BAD_ACCOUNT_KEY"))
|
||||
store.put_account = AsyncMock()
|
||||
|
||||
legacy_account = MagicMock()
|
||||
created = []
|
||||
|
||||
class FakePgCryptoStore:
|
||||
def __init__(self, account_id, pickle_key, db):
|
||||
self.pickle_key = pickle_key
|
||||
created.append(pickle_key)
|
||||
|
||||
async def get_account(self):
|
||||
if self.pickle_key == "@bot:example.org:default":
|
||||
return legacy_account
|
||||
raise RuntimeError("BAD_ACCOUNT_KEY")
|
||||
|
||||
crypto_db = MagicMock()
|
||||
crypto_db.fetch = AsyncMock(return_value=[])
|
||||
crypto_db.execute = AsyncMock()
|
||||
|
||||
fake_mod = types.ModuleType("mautrix.crypto.store.asyncpg")
|
||||
fake_mod.PgCryptoStore = FakePgCryptoStore
|
||||
with patch.dict(sys.modules, {"mautrix.crypto.store.asyncpg": fake_mod}), \
|
||||
caplog.at_level(logging.INFO):
|
||||
result = await adapter._migrate_legacy_crypto_pickle(
|
||||
store, crypto_db, "@bot:example.org", "@bot:example.org:NEWDEV"
|
||||
)
|
||||
|
||||
assert result is True
|
||||
store.put_account.assert_awaited_once_with(legacy_account)
|
||||
assert "re-pickled crypto store account" in caplog.text
|
||||
assert "@bot:example.org:default" in created
|
||||
# session re-pickle pass must sweep all three session tables
|
||||
queried = " ".join(str(c.args[0]) for c in crypto_db.fetch.await_args_list)
|
||||
for table in (
|
||||
"crypto_olm_session",
|
||||
"crypto_megolm_inbound_session",
|
||||
"crypto_megolm_outbound_session",
|
||||
):
|
||||
assert table in queried
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unrecoverable_pickle_logs_error(self, caplog):
|
||||
import logging
|
||||
adapter = _make_adapter()
|
||||
store = MagicMock()
|
||||
store.get_account = AsyncMock(side_effect=RuntimeError("BAD_ACCOUNT_KEY"))
|
||||
store.put_account = AsyncMock()
|
||||
|
||||
class FakePgCryptoStore:
|
||||
def __init__(self, account_id, pickle_key, db):
|
||||
pass
|
||||
|
||||
async def get_account(self):
|
||||
raise RuntimeError("BAD_ACCOUNT_KEY")
|
||||
|
||||
fake_mod = types.ModuleType("mautrix.crypto.store.asyncpg")
|
||||
fake_mod.PgCryptoStore = FakePgCryptoStore
|
||||
with patch.dict(sys.modules, {"mautrix.crypto.store.asyncpg": fake_mod}), \
|
||||
caplog.at_level(logging.ERROR):
|
||||
result = await adapter._migrate_legacy_crypto_pickle(
|
||||
store, MagicMock(), "@bot:example.org", "@bot:example.org:NEWDEV"
|
||||
)
|
||||
|
||||
assert result is False
|
||||
store.put_account.assert_not_awaited()
|
||||
assert "cannot be unpickled" in caplog.text
|
||||
|
|
|
|||
Loading…
Reference in New Issue