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