diff --git a/plugins/platforms/matrix/adapter.py b/plugins/platforms/matrix/adapter.py index d36e262e1a657..990df344db7c0 100644 --- a/plugins/platforms/matrix/adapter.py +++ b/plugins/platforms/matrix/adapter.py @@ -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 ``: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 diff --git a/tests/gateway/test_matrix.py b/tests/gateway/test_matrix.py index c32ecf1d34e84..3fe6efcae435d 100644 --- a/tests/gateway/test_matrix.py +++ b/tests/gateway/test_matrix.py @@ -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