tests/test_program.py
2.4 KB · 63 lines · python Raw
1 """Tests for BPFProgram / BPFProgramMetadata."""
2
3 from __future__ import annotations
4
5 import base64
6
7 from pqc_ebpf_attestation import BPFProgram, BPFProgramMetadata, BPFProgramType
8
9
10 def test_hash_bytecode_is_deterministic() -> None:
11 data = b"some eBPF ELF bytes"
12 h1 = BPFProgram.hash_bytecode(data)
13 h2 = BPFProgram.hash_bytecode(data)
14 assert h1 == h2
15 assert len(h1) == 64 # sha3-256 hex
16
17
18 def test_hash_differs_with_content() -> None:
19 h1 = BPFProgram.hash_bytecode(b"payload-a")
20 h2 = BPFProgram.hash_bytecode(b"payload-b")
21 assert h1 != h2
22
23
24 def test_from_bytes_and_from_file(tmp_path, sample_bpf_metadata, sample_bytecode) -> None:
25 prog_bytes = BPFProgram.from_bytes(sample_bpf_metadata, sample_bytecode)
26 assert prog_bytes.bytecode == sample_bytecode
27 assert prog_bytes.bytecode_size == len(sample_bytecode)
28 assert prog_bytes.bytecode_hash == BPFProgram.hash_bytecode(sample_bytecode)
29
30 bpf_path = tmp_path / "program.bpf.o"
31 bpf_path.write_bytes(sample_bytecode)
32 prog_file = BPFProgram.from_file(sample_bpf_metadata, str(bpf_path))
33 assert prog_file.bytecode == sample_bytecode
34 assert prog_file.bytecode_hash == prog_bytes.bytecode_hash
35
36
37 def test_canonical_manifest_bytes_is_deterministic(sample_bpf_metadata, sample_bytecode) -> None:
38 prog_a = BPFProgram.from_bytes(sample_bpf_metadata, sample_bytecode)
39 prog_b = BPFProgram.from_bytes(sample_bpf_metadata, sample_bytecode)
40 assert prog_a.canonical_manifest_bytes() == prog_b.canonical_manifest_bytes()
41
42 # Changing the bytecode also changes the canonical manifest (via hash+size).
43 prog_c = BPFProgram.from_bytes(sample_bpf_metadata, sample_bytecode + b"\x00")
44 assert prog_a.canonical_manifest_bytes() != prog_c.canonical_manifest_bytes()
45
46
47 def test_to_dict_with_and_without_bytecode(sample_bpf_metadata, sample_bytecode) -> None:
48 prog = BPFProgram.from_bytes(sample_bpf_metadata, sample_bytecode)
49
50 d_noby = prog.to_dict(include_bytecode=False)
51 assert "bytecode_base64" not in d_noby
52 assert d_noby["bytecode_hash"] == prog.bytecode_hash
53 assert d_noby["metadata"]["program_type"] == BPFProgramType.KPROBE.value
54
55 d_by = prog.to_dict(include_bytecode=True)
56 assert "bytecode_base64" in d_by
57 assert base64.b64decode(d_by["bytecode_base64"]) == sample_bytecode
58
59
60 def test_metadata_to_dict_uses_string_enum() -> None:
61 meta = BPFProgramMetadata(name="x", program_type=BPFProgramType.XDP)
62 assert meta.to_dict()["program_type"] == "xdp"
63