ZKAC/tests/test_pir.py
everbarry 6e67836e95 v0.4
2026-04-18 01:06:12 +02:00

43 lines
1.0 KiB
Python

"""Unit tests for XOR PIR helpers (Chor-style two-server PIR)."""
from zkac_cli import pir
def test_pad_roundtrip():
d = {
"grant_id": "abc",
"claimed": False,
"eph_pk_b64": "eA==",
"ciphertext_b64": "cA==",
}
p = pir.pad_grant_record(d)
assert len(p) == pir.PIR_RECORD_BYTES
assert pir.unpad_grant_record(p) == d
def test_xor_pir_row_recovery():
rows = [
{
"grant_id": f"id{i}",
"claimed": False,
"eph_pk_b64": f"e{i}==",
"ciphertext_b64": f"c{i}==",
}
for i in range(5)
]
want_i = 3
sa, sb = pir.pir_query_indices(len(rows), want_i)
def fold(idxs):
pads = [pir.pad_grant_record(rows[j]) for j in sorted(set(idxs))]
return pir.xor_bytes_many(pads)
xa = fold(sa)
xb = fold(sb)
got = pir.pir_recover(xa, xb)
assert got == pir.pad_grant_record(rows[want_i])
def test_pir_fold_empty_is_zero_block():
assert pir.xor_bytes_many([]) == b"\x00" * pir.PIR_RECORD_BYTES