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    pub gso_segs: u32,
43    pub gso_size: u32,
44}
45
46#[derive(Clone, FromBytes, IntoBytes, Immutable, KnownLayout)]
47#[repr(C)]
48pub struct TestProgramState {
49    pub config: TestConfig,
50    pub _padding: [u8; 4],
51    pub result: TestResult,
52}
53
54// LINT.ThenChange(//src/connectivity/network/testing/ebpf_test_util/ebpf/ebpf_test_progs.c)
55
56pub struct TestProgramDefinition {
57    program: ebpf::VerifiedEbpfProgram,
58    maps: Vec<ebpf_loader::MapDefinition>,
59}
60
61impl TestProgramDefinition {
62    /// Loads the test program, verifies it for the specified `program_type`.
63    pub fn load(program_type: ProgramType) -> Self {
64        let prog =
65            ebpf_loader::load_ebpf_program("/pkg/data/ebpf_test_progs.o", ".text", "skb_test_prog")
66                .expect("Failed to load test prog");
67        let maps_schema = prog.maps.iter().map(|m| m.schema).collect();
68        let calling_context = program_type
69            .create_calling_context(AttachType::Unspecified, maps_schema)
70            .expect("Failed to create CallingContext");
71        let program =
72            ebpf::verify_program(prog.code, calling_context, &mut ebpf::NullVerifierLogger)
73                .expect("Failed to verify loaded program");
74        Self { program, maps: prog.maps }
75    }
76
77    /// Initializes all maps used by the program.
78    pub fn instantiate(&self) -> TestProgram {
79        let maps = self
80            .maps
81            .iter()
82            .map(|def| ebpf_api::Map::new(def.schema, &def.name()).expect("Failed to create a map"))
83            .collect();
84
85        let (handle, server_handle) = zx::EventPair::create();
86        let handle = febpf::ProgramHandle { handle };
87        TestProgram { program: self.program.clone(), maps, handle, server_handle }
88    }
89}
90
91#[derive(Debug)]
92pub struct TestProgram {
93    handle: febpf::ProgramHandle,
94    server_handle: zx::EventPair,
95    program: ebpf::VerifiedEbpfProgram,
96    maps: Vec<ebpf_api::PinnedMap>,
97}
98
99impl TestProgram {
100    pub fn maps(&self) -> &[ebpf_api::PinnedMap] {
101        &self.maps
102    }
103
104    pub fn get_fidl_program(&self) -> febpf::VerifiedProgram {
105        let code: Vec<u64> =
106            <[u64]>::ref_from_bytes(self.program.code().as_bytes()).unwrap().to_owned();
107        let struct_access_instructions = self
108            .program
109            .struct_access_instructions()
110            .iter()
111            .map(|s| febpf::StructAccess {
112                pc: s.pc.try_into().unwrap(),
113                struct_id: StructId::try_from(&s.memory_id).unwrap().into(),
114                field_offset: s.field_offset.try_into().unwrap(),
115                is_32_bit_ptr_load: s.is_32_bit_ptr_load,
116            })
117            .collect();
118        febpf::VerifiedProgram {
119            code: Some(code),
120            struct_access_instructions: Some(struct_access_instructions),
121            maps: Some(self.maps.iter().map(|m| m.share().expect("share map")).collect()),
122            __source_breaking: Default::default(),
123        }
124    }
125
126    pub fn get_program_handle(&self) -> febpf::ProgramHandle {
127        febpf::ProgramHandle {
128            handle: self
129                .handle
130                .handle
131                .duplicate_handle(zx::Rights::SAME_RIGHTS)
132                .expect("duplicate handle"),
133        }
134    }
135
136    pub fn get_program_id(&self) -> febpf::ProgramId {
137        febpf::ProgramId { id: self.handle.handle.koid().expect("get koid").raw_koid() }
138    }
139
140    // Mark the program defunct and return the server handle.
141    pub fn mark_defunct(self) -> zx::EventPair {
142        let Self { handle, server_handle, .. } = self;
143        handle
144            .handle
145            .signal(
146                zx::Signals::empty(),
147                zx::Signals::from_bits_truncate(febpf::PROGRAM_DEFUNCT_SIGNAL),
148            )
149            .expect("signal EventPair");
150        server_handle
151    }
152
153    pub fn read_test_state(&self) -> TestProgramState {
154        let state = self.maps[0].load(&[0; 4]).expect("retrieve test state");
155        TestProgramState::ref_from_bytes(&state).expect("convert test state struct").clone()
156    }
157
158    pub fn read_test_result(&self) -> TestResult {
159        self.read_test_state().result
160    }
161
162    pub fn write_test_config(&self, config: TestConfig) {
163        let mut state = self.read_test_state();
164        state.config = config;
165        self.maps[0].update(&[0; 4], state.as_mut_bytes().into(), 0).expect("store test state");
166    }
167}