Skip to main content

ebpf_test_util/
lib.rs

1// Copyright 2025 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5use ebpf_api::{AttachType, ProgramType, StructId};
6use fidl_fuchsia_ebpf as febpf;
7use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout};
8
9// The structs below must match the structs in `ebpf_test_progs.c`.
10// LINT.IfChange
11
12/// Configuration for the test program. If either field is not zero then the
13/// test program will try to match these fields against the corresponding
14/// fields in UDP packets. In that case, if a packet doesn't match or it's not
15/// a UDP packet the program returns 0 without updating the `TestResult`. If
16/// both fields are zero or the packet matches then `TestResult` is updated
17/// and the program returns 1.
18#[derive(Clone, FromBytes, IntoBytes, Immutable, KnownLayout)]
19#[repr(C)]
20pub struct TestConfig {
21    /// Source port to match. Any port matches if set to 0.
22    pub src_port: u16,
23    /// Destination port to match. Any port matches if set to 0.
24    pub dst_port: u16,
25}
26
27/// Struct used to store results of the last invocation of the test program
28/// to test. The program copies the corresponding fields from the packet to
29/// this struct. Test can read it using `TestProgram::read_test_result()`.
30#[derive(Clone, FromBytes, IntoBytes, Immutable, KnownLayout)]
31#[repr(C)]
32pub struct TestResult {
33    pub cookie: u64,
34    pub uid: u32,
35    pub ifindex: u32,
36    pub ether_type: u32,
37    pub mark: u32,
38    pub src_port: u16,
39    pub dst_port: u16,
40    pub ip_proto: u8,
41    pub _padding: [u8; 3],
42}
43
44#[derive(Clone, FromBytes, IntoBytes, Immutable, KnownLayout)]
45#[repr(C)]
46pub struct TestProgramState {
47    pub config: TestConfig,
48    pub _padding: [u8; 4],
49    pub result: TestResult,
50}
51
52// LINT.ThenChange(//src/connectivity/network/testing/ebpf_test_util/ebpf/ebpf_test_progs.c)
53
54pub struct TestProgramDefinition {
55    program: ebpf::VerifiedEbpfProgram,
56    maps: Vec<ebpf_loader::MapDefinition>,
57}
58
59impl TestProgramDefinition {
60    /// Loads the test program, verifies it for the specified `program_type`.
61    pub fn load(program_type: ProgramType) -> Self {
62        let prog =
63            ebpf_loader::load_ebpf_program("/pkg/data/ebpf_test_progs.o", ".text", "skb_test_prog")
64                .expect("Failed to load test prog");
65        let maps_schema = prog.maps.iter().map(|m| m.schema).collect();
66        let calling_context = program_type
67            .create_calling_context(AttachType::Unspecified, maps_schema)
68            .expect("Failed to create CallingContext");
69        let program =
70            ebpf::verify_program(prog.code, calling_context, &mut ebpf::NullVerifierLogger)
71                .expect("Failed to verify loaded program");
72        Self { program, maps: prog.maps }
73    }
74
75    /// Initializes all maps used by the program.
76    pub fn instantiate(&self) -> TestProgram {
77        let maps = self
78            .maps
79            .iter()
80            .map(|def| ebpf_api::Map::new(def.schema, &def.name()).expect("Failed to create a map"))
81            .collect();
82
83        let (handle, server_handle) = zx::EventPair::create();
84        let handle = febpf::ProgramHandle { handle };
85        TestProgram { program: self.program.clone(), maps, handle, server_handle }
86    }
87}
88
89#[derive(Debug)]
90pub struct TestProgram {
91    handle: febpf::ProgramHandle,
92    server_handle: zx::EventPair,
93    program: ebpf::VerifiedEbpfProgram,
94    maps: Vec<ebpf_api::PinnedMap>,
95}
96
97impl TestProgram {
98    pub fn maps(&self) -> &[ebpf_api::PinnedMap] {
99        &self.maps
100    }
101
102    pub fn get_fidl_program(&self) -> febpf::VerifiedProgram {
103        let code: Vec<u64> =
104            <[u64]>::ref_from_bytes(self.program.code().as_bytes()).unwrap().to_owned();
105        let struct_access_instructions = self
106            .program
107            .struct_access_instructions()
108            .iter()
109            .map(|s| febpf::StructAccess {
110                pc: s.pc.try_into().unwrap(),
111                struct_id: StructId::try_from(&s.memory_id).unwrap().into(),
112                field_offset: s.field_offset.try_into().unwrap(),
113                is_32_bit_ptr_load: s.is_32_bit_ptr_load,
114            })
115            .collect();
116        febpf::VerifiedProgram {
117            code: Some(code),
118            struct_access_instructions: Some(struct_access_instructions),
119            maps: Some(self.maps.iter().map(|m| m.share().expect("share map")).collect()),
120            __source_breaking: Default::default(),
121        }
122    }
123
124    pub fn get_program_handle(&self) -> febpf::ProgramHandle {
125        febpf::ProgramHandle {
126            handle: self
127                .handle
128                .handle
129                .duplicate_handle(zx::Rights::SAME_RIGHTS)
130                .expect("duplicate handle"),
131        }
132    }
133
134    pub fn get_program_id(&self) -> febpf::ProgramId {
135        febpf::ProgramId { id: self.handle.handle.koid().expect("get koid").raw_koid() }
136    }
137
138    // Mark the program defunct and return the server handle.
139    pub fn mark_defunct(self) -> zx::EventPair {
140        let Self { handle, server_handle, .. } = self;
141        handle
142            .handle
143            .signal(
144                zx::Signals::empty(),
145                zx::Signals::from_bits_truncate(febpf::PROGRAM_DEFUNCT_SIGNAL),
146            )
147            .expect("signal EventPair");
148        server_handle
149    }
150
151    pub fn read_test_state(&self) -> TestProgramState {
152        let state = self.maps[0].load(&[0; 4]).expect("retrieve test state");
153        TestProgramState::ref_from_bytes(&state).expect("convert test state struct").clone()
154    }
155
156    pub fn read_test_result(&self) -> TestResult {
157        self.read_test_state().result
158    }
159
160    pub fn write_test_config(&self, config: TestConfig) {
161        let mut state = self.read_test_state();
162        state.config = config;
163        self.maps[0].update(&[0; 4], state.as_mut_bytes().into(), 0).expect("store test state");
164    }
165}