"""Tests for the query endpoint. Run: python3 -m unittest discover . A miniature PKI is minted with openssl per test class: a root CA, a server cert, an owner client cert, a grantee client cert, an expired- grant client cert, and a client cert from a DIFFERENT CA. The server binds loopback here ONLY because client and server are the same test process; production binds the mesh interface (docs/serve.md). """ import http.client import json import os import shutil import ssl import subprocess import tempfile import threading import unittest from . import embed as em from . import run as rn from . import serve as sv from .catalog import catalog from .extract import extract from .sanitize import sanitize from .test_embed import _Stub as EmbedStub # deterministic bag-of-words embedder from .walk import walk HAVE_OPENSSL = bool(shutil.which("openssl")) import datetime as _dt TODAY = _dt.date(2026, 8, 28) def sh(*argv): subprocess.run(list(argv), check=True, capture_output=True) def mint_pki(d): """root CA + server + clients; returns dict of paths.""" p = lambda n: os.path.join(d, n) def root(name, cn): sh("openssl", "ecparam", "-name", "prime256v1", "-genkey", "-noout", "-out", p(name + ".key")) sh("openssl", "req", "-new", "-x509", "-key", p(name + ".key"), "-out", p(name + ".crt"), "-days", "2", "-subj", "/CN=" + cn, "-addext", "basicConstraints=critical,CA:TRUE", "-addext", "keyUsage=critical,keyCertSign,cRLSign") root("ca", "test-root") root("other-ca", "other-root") def leaf(name, cn, ca="ca", san=None): sh("openssl", "ecparam", "-name", "prime256v1", "-genkey", "-noout", "-out", p(name + ".key")) sh("openssl", "req", "-new", "-key", p(name + ".key"), "-out", p(name + ".csr"), "-subj", "/CN=" + cn) ext = p(name + ".ext") with open(ext, "w") as fh: fh.write("subjectAltName=%s\nbasicConstraints=CA:FALSE\n" "extendedKeyUsage=serverAuth,clientAuth\n" % (san or "DNS:" + cn)) sh("openssl", "x509", "-req", "-in", p(name + ".csr"), "-CA", p(ca + ".crt"), "-CAkey", p(ca + ".key"), "-CAcreateserial", "-out", p(name + ".crt"), "-days", "2", "-extfile", ext) return p(name + ".crt"), p(name + ".key") out = {"ca": p("ca.crt"), "other_ca": p("other-ca.crt")} out["server"] = leaf("server", "localhost", san="DNS:localhost,IP:127.0.0.1") out["owner"] = leaf("owner", "owner-host") out["partner"] = leaf("partner", "partner-host") out["expired"] = leaf("expired", "old-host") out["stranger"] = leaf("stranger", "owner-host", ca="other-ca") # right CN, wrong CA out["nobody"] = leaf("nobody", "nobody-host") return out @unittest.skipUnless(HAVE_OPENSSL, "openssl not on PATH") class ServeTests(unittest.TestCase): @classmethod def setUpClass(cls): cls.work = tempfile.mkdtemp(prefix="fa-serve-") cls.pki = mint_pki(cls.work) cls.root = os.path.join(cls.work, "corpus") os.makedirs(cls.root) for name, text in (("apples.txt", "apples pears plums orchard harvest m@example.com\n"), ("routers.txt", "router switch firewall vlan network segment\n")): with open(os.path.join(cls.root, name), "w", encoding="utf-8") as fh: fh.write(text) cls.out = os.path.join(cls.work, "out") cls.manifest = os.path.join(cls.work, "m.yaml") cls.write_manifest(grants=[{"host": "partner-host", "ops": ["search"], "until": "2026-12-31"}, {"host": "old-host", "ops": ["search", "read_range"], "until": "2026-01-01"}]) m = cls.load() # embedding stub from http.server import ThreadingHTTPServer EmbedStub.dim = 64 cls.estub = ThreadingHTTPServer(("127.0.0.1", 0), EmbedStub) threading.Thread(target=cls.estub.serve_forever, daemon=True).start() cls.env = {"EMBED_URL": "http://127.0.0.1:%d" % cls.estub.server_address[1], "FREEAGENT_TLS_CERT": cls.pki["server"][0], "FREEAGENT_TLS_KEY": cls.pki["server"][1], "FREEAGENT_CA_ROOT": cls.pki["ca"]} walk(m, cls.out) extract(m, cls.out, use_external=False) sanitize(m, cls.out, env={}) em.chunk(m, cls.out) em.embed(m, cls.out, env=cls.env) catalog(m, cls.out, env={}) cls.srv = sv.build_server(cls.manifest, cls.out, env=cls.env, bind="127.0.0.1:0", today=TODAY) cls.port = cls.srv.server_address[1] threading.Thread(target=cls.srv.serve_forever, daemon=True).start() @classmethod def tearDownClass(cls): cls.srv.shutdown() cls.srv.server_close() cls.estub.shutdown() shutil.rmtree(cls.work, ignore_errors=True) @classmethod def write_manifest(cls, grants, mode="flag"): import yaml doc = {"freeagent": 1, "id": "t", "owner": {"host": "owner-host"}, "scope": {"roots": [cls.root]}, "mask": ["read", "search", "index"], "stages": ["walk", "extract", "sanitize", "chunk", "embed", "catalog", "seal"], "sanitize": {"mode": mode}, "embed": {"model": "stub-model", "dim": 64, "chunk": {"size_tokens": 8, "overlap_tokens": 2}}, "sweep": "manual", "sink": {"catalog": "host", "vectors": "local", "plaintext": "never"}, "grants": grants} with open(cls.manifest, "w", encoding="utf-8") as fh: yaml.safe_dump(doc, fh, sort_keys=False) rn.sign_manifest_file(cls.manifest, cls.pki["owner"][1], cls.pki["owner"][0]) @classmethod def load(cls): from .manifest import load_manifest return load_manifest(cls.manifest) # ── client ─────────────────────────────────────────────────────────── def client(self, who, ca=None): ctx = ssl.create_default_context(cafile=ca or self.pki["ca"]) ctx.check_hostname = False # test server cert is for localhost/127.0.0.1; keep simple if who: ctx.load_cert_chain(*self.pki[who]) return http.client.HTTPSConnection("127.0.0.1", self.port, context=ctx, timeout=10) def call(self, who, method, path, body=None): c = self.client(who) data = json.dumps(body).encode() if body is not None else None c.request(method, path, body=data, headers={"Content-Type": "application/json"}) r = c.getresponse() out = json.loads(r.read().decode("utf-8")) c.close() return r.status, out def audit(self): with open(os.path.join(self.out, "serve.audit.jsonl"), encoding="utf-8") as fh: return [json.loads(l) for l in fh] # ── tests ──────────────────────────────────────────────────────────── def test_owner_everything(self): st, body = self.call("owner", "GET", "/health") self.assertEqual((st, body["ok"], body["serving"]), (200, True, True)) st, body = self.call("owner", "POST", "/search", {"query": "firewall vlan", "k": 2}) self.assertEqual(st, 200) self.assertEqual(len(body["hits"]), 2) self.assertIn("router", body["hits"][0]["text"]) doc = body["hits"][0]["doc"] st, body = self.call("owner", "POST", "/read_range", {"doc": doc, "start": 0, "end": 20}) self.assertEqual(st, 200) self.assertEqual(body["text"], "router switch firewa") self.assertEqual(body["total"], 44) st, body = self.call("owner", "GET", "/catalog") self.assertEqual(st, 200) self.assertEqual({c["path"] for c in body["cards"]}, {"apples.txt", "routers.txt"}) # sanitized store only: the apples card head carries a token, not the address apples = [c for c in body["cards"] if c["path"] == "apples.txt"][0] self.assertIn("[PII_EMAIL_1]", apples["head"]) st, body = self.call("owner", "GET", "/status") self.assertEqual(st, 200) # the map: directory cards, root "" aggregates everything st, body = self.call("owner", "GET", "/catalog/dirs") self.assertEqual((st, body["kind"]), (200, "dirs")) root = [c for c in body["cards"] if c["path"] == ""][0] self.assertEqual(root["files"], 2) self.assertEqual(root["entities"], {"EMAIL": 1}) self.assertEqual(self.call("partner", "GET", "/catalog/dirs")[0], 403) # paging: total/offset/limit/next_offset; prefix filters by path st, body = self.call("owner", "GET", "/catalog?limit=1") self.assertEqual((st, body["total"], len(body["cards"]), body["next_offset"]), (200, 2, 1, 1)) st, body = self.call("owner", "GET", "/catalog?limit=1&offset=1") self.assertEqual((len(body["cards"]), body["next_offset"]), (1, None)) st, body = self.call("owner", "GET", "/catalog?prefix=rou") self.assertEqual([c["path"] for c in body["cards"]], ["routers.txt"]) self.assertEqual(self.call("owner", "GET", "/catalog?limit=x")[0], 400) def test_grantee_gets_only_its_ops(self): st, body = self.call("partner", "POST", "/search", {"query": "apples"}) self.assertEqual(st, 200) self.assertIn("[PII_EMAIL_1]", body["hits"][0]["text"]) # sanitized, always doc = body["hits"][0]["doc"] st, body = self.call("partner", "POST", "/read_range", {"doc": doc, "start": 0, "end": 5}) self.assertEqual(st, 403) self.assertIn("partner-host", body["error"]) self.assertEqual(self.call("partner", "GET", "/catalog")[0], 403) self.assertEqual(self.call("partner", "GET", "/status")[0], 403) self.assertEqual(self.call("partner", "GET", "/health")[0], 200) def test_expired_grant_and_unknown_cn(self): self.assertEqual(self.call("expired", "POST", "/search", {"query": "x"})[0], 403) self.assertEqual(self.call("nobody", "POST", "/search", {"query": "x"})[0], 403) self.assertEqual(self.call("nobody", "GET", "/health")[0], 200) # chained, so alive def test_no_cert_or_foreign_ca_never_reaches_a_handler(self): with self.assertRaises((ssl.SSLError, ConnectionError, OSError)): self.call(None, "GET", "/health") with self.assertRaises((ssl.SSLError, ConnectionError, OSError)): self.call("stranger", "GET", "/health") # CN 'owner-host' but signed by another CA # and no audit line was written for either self.assertFalse(any(a["cn"] is None for a in self.audit())) def test_read_range_is_bounded_and_sanitized_only(self): st, body = self.call("owner", "POST", "/search", {"query": "apples", "k": 1}) doc = body["hits"][0]["doc"] st, body = self.call("owner", "POST", "/read_range", {"doc": doc, "start": 0, "end": 10 ** 9}) self.assertEqual(st, 200) self.assertLessEqual(body["end"] - body["start"], sv.READ_RANGE_MAX) self.assertNotIn("m@example.com", body["text"]) # one id per document: the extract and sanitized stores share the # content hash, and the route reads ONLY the sanitized directory — # so the extract-side plaintext for the same id is never what comes back with open(os.path.join(self.out, "extract.jsonl"), encoding="utf-8") as fh: ext_sha = [json.loads(l) for l in fh if "apples" in l][0]["sha256"] st, body = self.call("owner", "POST", "/read_range", {"doc": ext_sha, "start": 0, "end": 200}) self.assertEqual(st, 200) self.assertIn("[PII_EMAIL_1]", body["text"]) self.assertNotIn("m@example.com", body["text"]) with open(os.path.join(self.out, "extract", ext_sha[:2], ext_sha + ".txt"), encoding="utf-8") as fh: self.assertIn("m@example.com", fh.read()) # the plaintext exists on disk; it has no route st, _ = self.call("owner", "POST", "/read_range", {"doc": "0" * 64, "start": 0, "end": 10}) self.assertEqual(st, 404) self.assertEqual(self.call("owner", "POST", "/read_range", {"doc": "zz", "start": 0, "end": 1})[0], 400) self.assertEqual(self.call("owner", "POST", "/search", {"query": ""})[0], 400) self.assertEqual(self.call("owner", "POST", "/nope", {})[0], 404) GRANTS = [{"host": "partner-host", "ops": ["search"], "until": "2026-12-31"}, {"host": "old-host", "ops": ["search", "read_range"], "until": "2026-01-01"}] def modes(self): """(owner search, owner read_range, partner search) sanitize_mode for the apples doc.""" st, body = self.call("owner", "POST", "/search", {"query": "apples", "k": 1}) self.assertEqual(st, 200) doc = body["hits"][0]["doc"] st, rr = self.call("owner", "POST", "/read_range", {"doc": doc, "start": 0, "end": 10}) self.assertEqual(st, 200) st, pb = self.call("partner", "POST", "/search", {"query": "apples", "k": 1}) self.assertEqual(st, 200) return doc, (body["sanitize_mode"], rr["sanitize_mode"], pb["sanitize_mode"]) def rewrite_index(self, fn): """Rewrite sanitize.jsonl through fn(record); returns the original bytes.""" path = os.path.join(self.out, "sanitize.jsonl") with open(path, "rb") as fh: orig = fh.read() recs = [fn(json.loads(l)) for l in orig.decode("utf-8").splitlines() if l.strip()] with open(path, "w", encoding="utf-8") as fh: fh.write("".join(json.dumps(r, sort_keys=True) + "\n" for r in recs)) return orig def restore_index(self, orig): with open(os.path.join(self.out, "sanitize.jsonl"), "wb") as fh: fh.write(orig) def test_a_reply_names_the_mode_that_wrote_its_text_not_the_manifests(self): # The store was swept in flag mode. Switching the manifest to # sterilize changes what the NEXT sweep does, not the text already # served: until a sweep rewrites it, every reply still says flag, so # the bastion never calls flag-mode text sanitized. try: for mode in ("flag", "sterilize", "flag"): self.write_manifest(grants=self.GRANTS, mode=mode) self.assertEqual(self.modes()[1], ("flag", "flag", "flag"), mode) finally: self.write_manifest(grants=self.GRANTS) def test_a_store_the_sanitize_stage_wrote_in_sterilize_mode_is_named_sterilize(self): orig = self.rewrite_index(lambda r: dict(r, mode="sterilize")) try: self.write_manifest(grants=self.GRANTS, mode="sterilize") self.assertEqual(self.modes()[1], ("sterilize", "sterilize", "sterilize")) finally: self.restore_index(orig) self.write_manifest(grants=self.GRANTS) def test_text_the_sanitize_record_does_not_vouch_for_is_unknown(self): doc, modes = self.modes() self.assertEqual(modes, ("flag", "flag", "flag")) spath = sv.sanitized_store_path(self.out, doc) st = os.stat(spath) index = os.path.join(self.out, "sanitize.jsonl") # a sweep still running has rewritten the file after the index later = os.stat(index).st_mtime_ns + 10 ** 9 os.utime(spath, ns=(later, later)) try: self.assertEqual(self.modes()[1], ("unknown", "unknown", "unknown")) finally: os.utime(spath, ns=(st.st_atime_ns, st.st_mtime_ns)) # the chunks a search reads were cut from text the store no longer holds with open(spath, encoding="utf-8") as fh: text = fh.read() with open(spath, "w", encoding="utf-8") as fh: fh.write("X" * len(text)) os.utime(spath, ns=(st.st_atime_ns, st.st_mtime_ns)) try: got = self.modes()[1] self.assertEqual((got[0], got[2]), ("unknown", "unknown")) finally: with open(spath, "w", encoding="utf-8") as fh: fh.write(text) os.utime(spath, ns=(st.st_atime_ns, st.st_mtime_ns)) # the last run did not write it (stopped: not_run), so the file is an older run's orig = self.rewrite_index(lambda r: dict(r, status="not_run") if r.get("sha256") == doc else r) try: self.assertEqual(self.modes()[1], ("unknown", "unknown", "unknown")) finally: self.restore_index(orig) self.assertEqual(self.modes()[1], ("flag", "flag", "flag")) def test_sanitize_mode_of_a_manifest_without_the_stage_is_none(self): self.assertEqual(sv.sanitize_mode(None), "none") self.assertEqual(sv.sanitize_mode(type("M", (), {"sanitize": None})()), "none") def test_silent_tcp_connection_does_not_stall_other_clients(self): # with the handshake on the accept loop one idle raw TCP # connection blocked every client; now the handshake runs in the # connection's own thread with its own timeout import socket staller = socket.create_connection(("127.0.0.1", self.port)) try: st, body = self.call("owner", "GET", "/health") self.assertEqual((st, body["ok"]), (200, True)) finally: staller.close() def test_audit_hashes_query_never_logs_it(self): self.call("owner", "POST", "/search", {"query": "secret question about marisol"}) lines = self.audit() last = [a for a in lines if a["op"] == "search"][-1] self.assertEqual(last["cn"], "owner-host") self.assertEqual(len(last["q"]), 16) blob = json.dumps(lines) self.assertNotIn("marisol", blob) def test_revoke_grant_by_resigning(self): self.assertEqual(self.call("partner", "POST", "/search", {"query": "apples"})[0], 200) self.write_manifest(grants=[]) # drop the grant, re-sign self.assertEqual(self.call("partner", "POST", "/search", {"query": "apples"})[0], 403) self.assertEqual(self.call("owner", "POST", "/search", {"query": "apples"})[0], 200) # tamper without re-signing -> the endpoint refuses everything but health with open(self.manifest, encoding="utf-8") as fh: text = fh.read() with open(self.manifest, "w", encoding="utf-8") as fh: fh.write(text.replace("sweep: manual", "sweep: daily")) st, body = self.call("owner", "POST", "/search", {"query": "apples"}) self.assertEqual(st, 503) self.assertIn("gate", body["error"]) self.assertEqual(self.call("owner", "GET", "/health")[1]["serving"], False) # restore self.write_manifest(grants=[{"host": "partner-host", "ops": ["search"], "until": "2026-12-31"}, {"host": "old-host", "ops": ["search", "read_range"], "until": "2026-01-01"}]) self.assertEqual(self.call("owner", "POST", "/search", {"query": "apples"})[0], 200) class UnitTests(unittest.TestCase): def test_policy_fails_closed_when_the_manifest_cannot_be_read(self): # Review: an unreadable manifest (PermissionError, bad encoding) is # not a ManifestError, and it escaped after _mtime advanced — leaving # the PREVIOUS manifest in force. from unittest import mock work = tempfile.mkdtemp(prefix="fa-pol-") self.addCleanup(shutil.rmtree, work, ignore_errors=True) man = os.path.join(work, "m.yaml") with open(man, "w") as fh: fh.write("x: 1\n") good = object() with mock.patch.object(sv, "load_manifest", return_value=good), \ mock.patch.object(sv, "verify_signature", return_value={}): pol = sv.Policy(man, {}) self.assertIs(pol.manifest, good) os.utime(man, ns=(1, 1)) with mock.patch.object(sv, "load_manifest", side_effect=PermissionError("denied")): pol.reload() self.assertIsNone(pol.manifest) self.assertIn("manifest unreadable", pol.gate_error) # skipped rather than faked: the port file is written by build_server # only after it has loaded a real certificate chain and verified a signed # manifest, and minting those needs openssl @unittest.skipUnless(HAVE_OPENSSL, "openssl not on PATH: needed to mint the test PKI and sign the manifest") def test_port_file_written_after_bind(self): # a supervisor binds :0 and reads the port back — one atomic file work = tempfile.mkdtemp(prefix="fa-pf-") try: pki = mint_pki(work) man = os.path.join(work, "m.yaml") import yaml with open(man, "w") as fh: yaml.safe_dump({"freeagent": 1, "id": "p", "owner": {"host": "owner-host"}, "scope": {"roots": [work]}, "mask": ["read"], "stages": ["walk"], "sweep": "manual", "sink": {"catalog": "local", "vectors": "local", "plaintext": "never"}}, fh) rn.sign_manifest_file(man, pki["owner"][1], pki["owner"][0]) env = {"FREEAGENT_TLS_CERT": pki["server"][0], "FREEAGENT_TLS_KEY": pki["server"][1], "FREEAGENT_CA_ROOT": pki["ca"]} pf = os.path.join(work, "port") srv = sv.build_server(man, os.path.join(work, "out"), env=env, bind="127.0.0.1:0", port_file=pf) try: with open(pf) as fh: self.assertEqual(int(fh.read().strip()), srv.server_address[1]) finally: srv.server_close() finally: shutil.rmtree(work, ignore_errors=True) def test_parse_bind(self): self.assertEqual(sv.parse_bind("[::]:8470"), ("::", 8470)) self.assertEqual(sv.parse_bind("10.0.0.5:9"), ("10.0.0.5", 9)) self.assertEqual(sv.parse_bind("[2001:db8::31]:8470"), ("2001:db8::31", 8470)) with self.assertRaises(sv.ServeError): sv.parse_bind("nope") def test_peer_cn(self): self.assertEqual(sv.peer_cn({"subject": ((("commonName", "x"),),)}), "x") self.assertIsNone(sv.peer_cn({})) self.assertIsNone(sv.peer_cn(None)) def test_ssl_context_requires_all_three(self): with self.assertRaises(sv.ServeError): sv.build_ssl_context("", "", "") if __name__ == "__main__": unittest.main()