fix(plugin): backfill certified artifact digest

src/langbot/pkg/persistence/alembic/versions/0032_certification_artifact_digest_backfill.py

@@ -0,0 +1,138 @@
+"""Backfill raw artifact digests for legacy shared certificates.
+
+Revision ID: 0032_cert_artifact_digest
+Revises: 0031_merge_totp_assistant
+
+Older certified-plugin rows persisted every shared-admission fact except the raw
+archive digest.  Placement now requires that digest to match the installation's
+immutable artifact digest, so this migration fills only records whose remaining
+facts already prove the exact legacy shared-admission shape.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+
+import sqlalchemy as sa
+from alembic import op
+
+revision = '0032_cert_artifact_digest'
+down_revision = '0031_merge_totp_assistant'
+branch_labels = None
+depends_on = None
+
+_TABLE = 'plugin_settings'
+_SHARED_RUNTIME = 'shared-runtime-v1'
+_SHARED_ADMISSION_CODE = 'CERTIFIED_PLUGIN_SHARED_ELIGIBLE'
+_LOWER_HEX = frozenset('0123456789abcdef')
+_HEX = frozenset('0123456789abcdefABCDEF')
+
+
+def _is_digest(value: object, *, lowercase: bool) -> bool:
+    allowed = _LOWER_HEX if lowercase else _HEX
+    return isinstance(value, str) and len(value) == 64 and all(character in allowed for character in value)
+
+
+def _eligible_certification(install_info: object, artifact_digest: object) -> Mapping[object, object] | None:
+    if not _is_digest(artifact_digest, lowercase=True) or not isinstance(install_info, Mapping):
+        return None
+    certification = install_info.get('_certification')
+    if not isinstance(certification, Mapping) or 'artifact_digest' in certification:
+        return None
+    certificate_id = certification.get('certificate_id')
+    if not isinstance(certificate_id, str) or not certificate_id.strip():
+        return None
+    if not _is_digest(certification.get('normalized_digest'), lowercase=False):
+        return None
+    required_facts = {
+        'verification': 'valid',
+        'certificate_runtime_profile': _SHARED_RUNTIME,
+        'runtime_profile': _SHARED_RUNTIME,
+        'admission_code': _SHARED_ADMISSION_CODE,
+    }
+    if not all(certification.get(key) == value for key, value in required_facts.items()):
+        return None
+    return certification
+
+
+def _suspend_postgres_rls(conn: sa.Connection) -> tuple[bool, bool]:
+    if conn.dialect.name != 'postgresql':
+        return False, False
+    state = conn.execute(
+        sa.text('SELECT relrowsecurity, relforcerowsecurity FROM pg_class WHERE oid = to_regclass(:table_name)'),
+        {'table_name': _TABLE},
+    ).one()
+    rls_enabled, rls_forced = bool(state.relrowsecurity), bool(state.relforcerowsecurity)
+    if rls_forced:
+        conn.execute(sa.text(f'ALTER TABLE {_TABLE} NO FORCE ROW LEVEL SECURITY'))
+    if rls_enabled:
+        conn.execute(sa.text(f'ALTER TABLE {_TABLE} DISABLE ROW LEVEL SECURITY'))
+    return rls_enabled, rls_forced
+
+
+def _restore_postgres_rls(conn: sa.Connection, state: tuple[bool, bool]) -> None:
+    if conn.dialect.name != 'postgresql':
+        return
+    rls_enabled, rls_forced = state
+    if rls_enabled:
+        conn.execute(sa.text(f'ALTER TABLE {_TABLE} ENABLE ROW LEVEL SECURITY'))
+    if rls_forced:
+        conn.execute(sa.text(f'ALTER TABLE {_TABLE} FORCE ROW LEVEL SECURITY'))
+
+
+def backfill_certification_artifact_digests(conn: sa.Connection) -> None:
+    inspector = sa.inspect(conn)
+    if _TABLE not in inspector.get_table_names():
+        return
+    columns = {column['name'] for column in inspector.get_columns(_TABLE)}
+    required_columns = {
+        'workspace_uuid',
+        'plugin_author',
+        'plugin_name',
+        'artifact_digest',
+        'install_info',
+    }
+    if not required_columns <= columns:
+        return
+
+    plugin_settings = sa.table(
+        _TABLE,
+        sa.column('workspace_uuid', sa.String(36)),
+        sa.column('plugin_author', sa.String(255)),
+        sa.column('plugin_name', sa.String(255)),
+        sa.column('artifact_digest', sa.String(64)),
+        sa.column('install_info', sa.JSON()),
+    )
+    rows = conn.execute(sa.select(plugin_settings)).mappings().all()
+    for row in rows:
+        certification = _eligible_certification(row['install_info'], row['artifact_digest'])
+        if certification is None:
+            continue
+        updated_certification = dict(certification)
+        updated_certification['artifact_digest'] = row['artifact_digest']
+        updated_install_info = dict(row['install_info'])
+        updated_install_info['_certification'] = updated_certification
+        conn.execute(
+            plugin_settings.update()
+            .where(plugin_settings.c.workspace_uuid == row['workspace_uuid'])
+            .where(plugin_settings.c.plugin_author == row['plugin_author'])
+            .where(plugin_settings.c.plugin_name == row['plugin_name'])
+            .values(install_info=updated_install_info)
+        )
+
+
+def upgrade() -> None:
+    conn = op.get_bind()
+    if _TABLE not in sa.inspect(conn).get_table_names():
+        return
+    rls_state = _suspend_postgres_rls(conn)
+    try:
+        backfill_certification_artifact_digests(conn)
+    finally:
+        _restore_postgres_rls(conn, rls_state)
+
+
+def downgrade() -> None:
+    # The source digest cannot be distinguished from a digest persisted by
+    # current Core, so downgrade intentionally preserves the safe enrichment.
+    pass

tests/integration/persistence/test_migration_branch_convergence.py

@@ -16,6 +16,7 @@ import uuid
 import pytest
 import sqlalchemy as sa
 from sqlalchemy.ext.asyncio import create_async_engine
+from langbot_plugin.entities.io.context import PluginExecutionMode
 
 from langbot.pkg.entity import persistence
 from langbot.pkg.entity.persistence.base import Base
@@ -25,6 +26,7 @@ from langbot.pkg.persistence.alembic_runner import (
     run_alembic_downgrade,
     run_alembic_upgrade,
 )
+from langbot.pkg.plugin.certification import execution_mode_for_persisted_installation
 from langbot.pkg.utils import importutil
 
 
@@ -353,6 +355,143 @@ async def test_empty_database_startup_schema_then_real_migrations(convergence_en
     assert await get_alembic_current(engine) == get_alembic_head()
 
 
+def _legacy_shared_certification(**overrides):
+    certification = {
+        'normalized_digest': 'B' * 64,
+        'verification': 'valid',
+        'certificate_runtime_profile': 'shared-runtime-v1',
+        'certificate_id': 'ed25519:trusted-issuer',
+        'runtime_profile': 'shared-runtime-v1',
+        'admission_code': 'CERTIFIED_PLUGIN_SHARED_ELIGIBLE',
+        'preserved_certificate_key': {'nested': True},
+    }
+    certification.update(overrides)
+    return certification
+
+
+@pytest.mark.asyncio
+async def test_certification_artifact_digest_backfill_is_safe_and_enables_shared_placement(convergence_engine):
+    engine = convergence_engine
+    async with engine.begin() as conn:
+        await conn.run_sync(Base.metadata.create_all)
+    await run_alembic_upgrade(engine, '0031_merge_totp_assistant')
+
+    valid_digest = 'a' * 64
+    rows = {
+        'eligible-a': (_legacy_shared_certification(), valid_digest),
+        'eligible-b': (_legacy_shared_certification(), valid_digest),
+        'present-matching': (_legacy_shared_certification(artifact_digest=valid_digest), valid_digest),
+        'present-mismatched': (_legacy_shared_certification(artifact_digest='c' * 64), valid_digest),
+        'invalid': (_legacy_shared_certification(verification='invalid'), valid_digest),
+        'incomplete': (_legacy_shared_certification(certificate_id='   '), valid_digest),
+        'missing-fact': (
+            {
+                key: value
+                for key, value in _legacy_shared_certification().items()
+                if key != 'certificate_runtime_profile'
+            },
+            valid_digest,
+        ),
+        'malformed-normalized': (_legacy_shared_certification(normalized_digest='g' * 64), valid_digest),
+        'dedicated-certificate': (
+            _legacy_shared_certification(certificate_runtime_profile='dedicated'),
+            valid_digest,
+        ),
+        'dedicated': (_legacy_shared_certification(runtime_profile='dedicated'), valid_digest),
+        'wrong-admission': (_legacy_shared_certification(admission_code='SHARED_ELIGIBLE'), valid_digest),
+        'uppercase-row-digest': (_legacy_shared_certification(), 'A' * 64),
+        'nonhex-row-digest': (_legacy_shared_certification(), 'g' * 64),
+    }
+    plugin_settings = Base.metadata.tables['plugin_settings']
+    workspaces = Base.metadata.tables['workspaces']
+    async with engine.begin() as conn:
+        for index, (name, (certification, artifact_digest)) in enumerate(rows.items(), start=1):
+            workspace_uuid = f'41100000-0000-4000-8000-{index:012d}'
+            await conn.execute(
+                workspaces.insert().values(
+                    uuid=workspace_uuid,
+                    instance_uuid=f'cert-backfill-{index}',
+                    name=name,
+                    slug=name,
+                )
+            )
+            await conn.execute(
+                plugin_settings.insert().values(
+                    workspace_uuid=workspace_uuid,
+                    plugin_author='langbot',
+                    plugin_name=name,
+                    installation_uuid=f'51100000-0000-4000-8000-{index:012d}',
+                    artifact_digest=artifact_digest,
+                    runtime_revision=1,
+                    install_info={
+                        '_certification': certification,
+                        'preserved_install_key': ['keep', {'nested': True}],
+                    },
+                )
+            )
+
+    await run_alembic_upgrade(engine, 'head')
+    # Re-running the data revision itself must also be harmless.
+    migration = __import__(
+        'langbot.pkg.persistence.alembic.versions.0032_certification_artifact_digest_backfill',
+        fromlist=['upgrade'],
+    )
+    async with engine.begin() as conn:
+        await conn.run_sync(lambda sync: migration.backfill_certification_artifact_digests(sync))
+        stored = {
+            name: (artifact_digest, install_info)
+            for name, artifact_digest, install_info in (
+                await conn.execute(
+                    sa.select(
+                        plugin_settings.c.plugin_name,
+                        plugin_settings.c.artifact_digest,
+                        plugin_settings.c.install_info,
+                    )
+                )
+            ).all()
+        }
+
+    for name in ('eligible-a', 'eligible-b'):
+        eligible_digest, eligible_info = stored[name]
+        assert eligible_info['preserved_install_key'] == ['keep', {'nested': True}]
+        assert eligible_info['_certification']['preserved_certificate_key'] == {'nested': True}
+        assert eligible_info['_certification']['artifact_digest'] == eligible_digest == valid_digest
+        assert (
+            execution_mode_for_persisted_installation(
+                artifact_digest=eligible_digest,
+                install_info=eligible_info,
+            )
+            is PluginExecutionMode.SHARED_CERTIFIED
+        )
+
+    for name, (original_certification, artifact_digest) in rows.items():
+        if name in {'eligible-a', 'eligible-b'}:
+            continue
+        stored_digest, stored_info = stored[name]
+        assert stored_digest == artifact_digest
+        assert stored_info['_certification'] == original_certification
+        if name != 'present-matching':
+            assert (
+                execution_mode_for_persisted_installation(
+                    artifact_digest=stored_digest,
+                    install_info=stored_info,
+                )
+                is PluginExecutionMode.DEDICATED
+            )
+
+
+@pytest.mark.asyncio
+async def test_certification_artifact_digest_backfill_accepts_fresh_current_schema(convergence_engine):
+    engine = convergence_engine
+    async with engine.begin() as conn:
+        await conn.run_sync(Base.metadata.create_all)
+
+    await run_alembic_upgrade(engine, 'head')
+    await run_alembic_upgrade(engine, 'head')
+
+    assert await get_alembic_current(engine) == get_alembic_head()
+
+
 @pytest.mark.asyncio
 async def test_merge_only_downgrade_preserves_both_branch_schemas(convergence_engine):
     engine = convergence_engine

tests/integration/persistence/test_migrations_postgres.py

@@ -329,6 +329,76 @@ class TestPostgreSQLMigrationBaseline:
         rev = await get_alembic_current(postgres_engine)
         assert rev == '0001_baseline'
 
+    @pytest.mark.asyncio
+    async def test_certification_artifact_digest_backfill_updates_only_complete_legacy_shared_rows(
+        self,
+        postgres_engine,
+        clean_tables,
+        clean_alembic_version,
+    ):
+        async with postgres_engine.begin() as conn:
+            await conn.run_sync(Base.metadata.create_all)
+        await run_alembic_stamp(postgres_engine, '0031_merge_totp_assistant')
+
+        plugin_settings = Base.metadata.tables['plugin_settings']
+        workspaces = Base.metadata.tables['workspaces']
+        digest = 'a' * 64
+        complete = {
+            'normalized_digest': 'B' * 64,
+            'verification': 'valid',
+            'certificate_runtime_profile': 'shared-runtime-v1',
+            'certificate_id': 'ed25519:trusted-issuer',
+            'runtime_profile': 'shared-runtime-v1',
+            'admission_code': 'CERTIFIED_PLUGIN_SHARED_ELIGIBLE',
+            'preserved': {'key': True},
+        }
+        rows = {
+            'eligible': complete,
+            'invalid': {**complete, 'verification': 'invalid'},
+            'present': {**complete, 'artifact_digest': 'c' * 64},
+        }
+        async with postgres_engine.begin() as conn:
+            for index, (name, certification) in enumerate(rows.items(), start=1):
+                workspace_uuid = f'61100000-0000-4000-8000-{index:012d}'
+                await conn.execute(
+                    workspaces.insert().values(
+                        uuid=workspace_uuid,
+                        instance_uuid=f'postgres-cert-backfill-{index}',
+                        name=name,
+                        slug=name,
+                    )
+                )
+                await conn.execute(
+                    plugin_settings.insert().values(
+                        workspace_uuid=workspace_uuid,
+                        plugin_author='langbot',
+                        plugin_name=name,
+                        installation_uuid=f'71100000-0000-4000-8000-{index:012d}',
+                        artifact_digest=digest,
+                        runtime_revision=1,
+                        install_info={'_certification': certification, 'preserved': ['install-info']},
+                    )
+                )
+
+        await run_alembic_upgrade(postgres_engine, 'head')
+
+        async with postgres_engine.connect() as conn:
+            stored = dict(
+                (await conn.execute(sa.select(plugin_settings.c.plugin_name, plugin_settings.c.install_info))).all()
+            )
+        assert stored['eligible'] == {
+            '_certification': {**complete, 'artifact_digest': digest},
+            'preserved': ['install-info'],
+        }
+        assert stored['invalid'] == {
+            '_certification': rows['invalid'],
+            'preserved': ['install-info'],
+        }
+        assert stored['present'] == {
+            '_certification': rows['present'],
+            'preserved': ['install-info'],
+        }
+
     @pytest.mark.asyncio
     async def test_fresh_postgres_schema_accepts_application_casefold_identity(
         self,