|
| 1 | +from __future__ import annotations |
| 2 | + |
| 3 | +import json |
| 4 | +import os |
| 5 | +import pathlib |
| 6 | +import tempfile |
| 7 | +import unittest |
| 8 | +import uuid |
| 9 | + |
| 10 | +import failures |
| 11 | +from artifacts_pg import PostgresArtifactStore |
| 12 | +from research.loop_repository import ResearchLoopRepository |
| 13 | +from research.provider_registry import FileAiProviderRegistry |
| 14 | + |
| 15 | +try: |
| 16 | + import psycopg |
| 17 | +except ImportError: # pragma: no cover - integration environment decides |
| 18 | + psycopg = None |
| 19 | + |
| 20 | + |
| 21 | +ROOT = pathlib.Path(__file__).resolve().parents[2] |
| 22 | +EXAMPLE = json.loads( |
| 23 | + (ROOT / "deploy/ai/providers.example.json").read_text(encoding="utf-8") |
| 24 | +) |
| 25 | + |
| 26 | + |
| 27 | +@unittest.skipUnless(os.environ.get("DIRAC_TEST_DSN") and psycopg, |
| 28 | + "requires isolated PostgreSQL DIRAC_TEST_DSN with migrations 000-050") |
| 29 | +class ResearchLoopSecurityTests(unittest.TestCase): |
| 30 | + @classmethod |
| 31 | + def setUpClass(cls): |
| 32 | + cls.dsn = os.environ["DIRAC_TEST_DSN"] |
| 33 | + cls.connect = staticmethod(lambda: psycopg.connect(cls.dsn)) |
| 34 | + with cls.connect() as conn, conn.cursor() as cur: |
| 35 | + cur.execute("SELECT current_database()") |
| 36 | + if "test" not in str(cur.fetchone()[0]).lower(): |
| 37 | + raise RuntimeError("security tests refuse a database without 'test' in its name") |
| 38 | + code = "SEC-" + uuid.uuid4().hex[:10] |
| 39 | + cls.actor = {"kind": "human", "id": code} |
| 40 | + cur.execute( |
| 41 | + "INSERT INTO design.project(code,name) VALUES (%s,%s) RETURNING id", |
| 42 | + (code, "Research-loop security test"), |
| 43 | + ) |
| 44 | + cls.program_id = str(cur.fetchone()[0]) |
| 45 | + cur.execute( |
| 46 | + "INSERT INTO design.campaign " |
| 47 | + "(program_id,name,objective,created_by_kind,created_by_id) " |
| 48 | + "VALUES (%s,%s,%s,'human',%s) RETURNING id", |
| 49 | + (cls.program_id, code, "verify private loop surfaces", code), |
| 50 | + ) |
| 51 | + cls.campaign_id = str(cur.fetchone()[0]) |
| 52 | + cls.repository = ResearchLoopRepository(cls.connect) |
| 53 | + cls.loop = cls.repository.create( |
| 54 | + request_key="security-loop", program_id=cls.program_id, |
| 55 | + campaign_id=cls.campaign_id, actor=cls.actor, |
| 56 | + intent="Verify tenant isolation and provider-secret redaction", |
| 57 | + autonomy_class="A2", provider_profile_id="qwen-local-isolated", |
| 58 | + provider_profile_digest=b"p" * 32, |
| 59 | + prompt_release_id="fep-action-proposal-v1", |
| 60 | + prompt_release_digest=b"r" * 32, |
| 61 | + action_catalog_digest=b"a" * 32, data_classification="restricted", |
| 62 | + policy={"cloud_egress_approved": False}, |
| 63 | + budget_remaining={"reasoner_calls": 1, "fep_runsets": 0}, |
| 64 | + ) |
| 65 | + cls.store = PostgresArtifactStore(cls.connect) |
| 66 | + cls.artifact = cls.store.put( |
| 67 | + b'{"claim_boundary":"model_proposal_not_scientific_evidence"}', |
| 68 | + role="research.context_snapshot", media_type="application/json", |
| 69 | + metadata={"visibility": "private"}, |
| 70 | + ) |
| 71 | + cls.repository.link_artifact( |
| 72 | + run_id=cls.loop["run_id"], artifact_id=cls.artifact.id, |
| 73 | + role="research.context_snapshot", data_classification="restricted", |
| 74 | + ) |
| 75 | + with cls.connect() as conn, conn.cursor() as cur: |
| 76 | + cur.execute( |
| 77 | + "UPDATE app.research_loop_state SET state='paused' WHERE run_id=%s", |
| 78 | + (cls.loop["run_id"],), |
| 79 | + ) |
| 80 | + |
| 81 | + def test_loop_artifact_uuid_and_digest_are_not_authorization_capabilities(self): |
| 82 | + owner = self.store.read_authorized(self.artifact.id, self.actor) |
| 83 | + self.assertEqual(owner[1], |
| 84 | + b'{"claim_boundary":"model_proposal_not_scientific_evidence"}') |
| 85 | + intruder = {"kind": "human", "id": "intruder-" + uuid.uuid4().hex} |
| 86 | + for address in (self.artifact.id, "sha256:" + self.artifact.sha256): |
| 87 | + with self.subTest(address=address), self.assertRaises(failures.DiracNotFound): |
| 88 | + self.store.read_authorized(address, intruder) |
| 89 | + |
| 90 | + def test_provider_secret_is_absent_from_every_loop_persisted_surface(self): |
| 91 | + secret = "research-loop-secret-" + uuid.uuid4().hex |
| 92 | + profile = dict(EXAMPLE["profiles"][1]) |
| 93 | + with tempfile.NamedTemporaryFile(mode="w", suffix=".json", encoding="utf-8") as source: |
| 94 | + json.dump({"profiles": [profile]}, source) |
| 95 | + source.flush() |
| 96 | + registry = FileAiProviderRegistry(source.name, environ={ |
| 97 | + "DIRAC_LOCAL_QWEN_BASE_URL": "http://127.0.0.1:8000/v1", |
| 98 | + "DIRAC_LOCAL_QWEN_API_KEY": secret, |
| 99 | + }) |
| 100 | + resolved = registry.resolve("qwen-local-isolated") |
| 101 | + public_surfaces = { |
| 102 | + "repr": repr(resolved), |
| 103 | + "public": resolved.to_public_dict(), |
| 104 | + "attestation": registry.attest( |
| 105 | + resolved.profile_id, resolved.profile_digest, "restricted"), |
| 106 | + } |
| 107 | + self.assertNotIn(secret, json.dumps(public_surfaces, sort_keys=True)) |
| 108 | + |
| 109 | + with self.connect() as conn, conn.cursor() as cur: |
| 110 | + cur.execute( |
| 111 | + "SELECT row_to_json(ls)::text FROM app.research_loop_state ls " |
| 112 | + "WHERE run_id=%s UNION ALL " |
| 113 | + "SELECT row_to_json(le)::text FROM app.research_loop_event le " |
| 114 | + "WHERE run_id=%s UNION ALL " |
| 115 | + "SELECT row_to_json(la)::text FROM app.research_loop_artifact la " |
| 116 | + "WHERE run_id=%s", |
| 117 | + (self.loop["run_id"],) * 3, |
| 118 | + ) |
| 119 | + persisted = "\n".join(str(row[0]) for row in cur.fetchall()) |
| 120 | + cur.execute( |
| 121 | + "SELECT b.bytes,a.metadata::text FROM app.research_loop_artifact la " |
| 122 | + "JOIN app.artifact a ON a.id=la.artifact_id " |
| 123 | + "JOIN app.blob b ON b.sha256=a.blob_sha256 WHERE la.run_id=%s", |
| 124 | + (self.loop["run_id"],), |
| 125 | + ) |
| 126 | + for raw, metadata in cur.fetchall(): |
| 127 | + persisted += bytes(raw).decode("utf-8", errors="replace") + str(metadata) |
| 128 | + self.assertNotIn(secret, persisted) |
| 129 | + |
| 130 | + |
| 131 | +if __name__ == "__main__": |
| 132 | + unittest.main() |
0 commit comments