Skip to main content

ebpf_api/
helpers.rs

1// Copyright 2024 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 crate::MapKey;
6use crate::maps::{Map, MapValueRef, RingBuffer, RingBufferWakeupPolicy};
7use ebpf::{BpfValue, EbpfBufferPtr, EbpfHelperImpl, EbpfProgramContext, FromBpfValue, HelperSet};
8use inspect_stubs::track_stub;
9use linux_uapi::{
10    BPF_SK_STORAGE_GET_F_CREATE, bpf_func_id_BPF_FUNC_get_current_pid_tgid,
11    bpf_func_id_BPF_FUNC_get_current_uid_gid, bpf_func_id_BPF_FUNC_get_netns_cookie,
12    bpf_func_id_BPF_FUNC_get_retval, bpf_func_id_BPF_FUNC_get_smp_processor_id,
13    bpf_func_id_BPF_FUNC_get_socket_cookie, bpf_func_id_BPF_FUNC_get_socket_uid,
14    bpf_func_id_BPF_FUNC_ktime_get_boot_ns, bpf_func_id_BPF_FUNC_ktime_get_coarse_ns,
15    bpf_func_id_BPF_FUNC_ktime_get_ns, bpf_func_id_BPF_FUNC_map_delete_elem,
16    bpf_func_id_BPF_FUNC_map_lookup_elem, bpf_func_id_BPF_FUNC_map_update_elem,
17    bpf_func_id_BPF_FUNC_probe_read_str, bpf_func_id_BPF_FUNC_probe_read_user,
18    bpf_func_id_BPF_FUNC_probe_read_user_str, bpf_func_id_BPF_FUNC_ringbuf_discard,
19    bpf_func_id_BPF_FUNC_ringbuf_reserve, bpf_func_id_BPF_FUNC_ringbuf_submit,
20    bpf_func_id_BPF_FUNC_set_retval, bpf_func_id_BPF_FUNC_sk_fullsock,
21    bpf_func_id_BPF_FUNC_sk_lookup_tcp, bpf_func_id_BPF_FUNC_sk_lookup_udp,
22    bpf_func_id_BPF_FUNC_sk_release, bpf_func_id_BPF_FUNC_sk_storage_get,
23    bpf_func_id_BPF_FUNC_skb_load_bytes, bpf_func_id_BPF_FUNC_skb_load_bytes_relative,
24    bpf_func_id_BPF_FUNC_trace_printk, bpf_map_type_BPF_MAP_TYPE_RINGBUF,
25    bpf_map_type_BPF_MAP_TYPE_SK_STORAGE, gid_t, pid_t, uid_t,
26};
27use smallvec::SmallVec;
28use zerocopy::IntoBytes as _;
29
30pub trait MapsContext<'a> {
31    fn on_map_access(&mut self, map: &Map);
32    fn add_value_ref(&mut self, map_ref: MapValueRef<'a>);
33}
34
35pub trait MapsProgramContext: EbpfProgramContext {
36    fn on_map_access(context: &mut Self::RunContext<'_>, map: &Map);
37    fn add_value_ref<'a>(context: &mut Self::RunContext<'a>, map_ref: MapValueRef<'a>);
38}
39
40impl<C: EbpfProgramContext> MapsProgramContext for C
41where
42    for<'a> C::RunContext<'a>: MapsContext<'a>,
43{
44    fn on_map_access(context: &mut Self::RunContext<'_>, map: &Map) {
45        context.on_map_access(map);
46    }
47
48    fn add_value_ref<'a>(context: &mut Self::RunContext<'a>, map_ref: MapValueRef<'a>) {
49        context.add_value_ref(map_ref);
50    }
51}
52
53fn bpf_map_lookup_elem<'a, C: MapsProgramContext>(
54    context: &mut C::RunContext<'a>,
55    map: BpfValue,
56    key: BpfValue,
57    _: BpfValue,
58    _: BpfValue,
59    _: BpfValue,
60) -> BpfValue {
61    // SAFETY: The `map` must be a reference to a `Map` object kept alive by the program itself.
62    let map: &Map = unsafe { &*map.as_ptr::<Map>() };
63
64    // SAFETY: safety is ensured by the verifier.
65    let key = unsafe { EbpfBufferPtr::new(key.as_ptr::<u8>(), map.schema.key_size as usize) };
66    let key: MapKey = key.load();
67
68    C::on_map_access(context, map);
69
70    let Some(value_ref) = map.lookup(&key) else {
71        return BpfValue::default();
72    };
73
74    let result: BpfValue = value_ref.ptr().raw_ptr().into();
75
76    // If this is a map with ref-counted elements then save the reference for
77    // the lifetime of the program.
78    if value_ref.is_ref_counted() {
79        C::add_value_ref(context, value_ref);
80    }
81
82    result
83}
84
85fn bpf_map_update_elem<C: MapsProgramContext>(
86    context: &mut C::RunContext<'_>,
87    map: BpfValue,
88    key: BpfValue,
89    value: BpfValue,
90    flags: BpfValue,
91    _: BpfValue,
92) -> BpfValue {
93    // SAFETY: The `map` must be a reference to a `Map` object kept alive by the program itself.
94    let map: &Map = unsafe { &*map.as_ptr::<Map>() };
95
96    // TODO(https://fxbug.dev/496639039): This should be checked by the verifier.
97    if map.schema.map_type == bpf_map_type_BPF_MAP_TYPE_SK_STORAGE {
98        return BpfValue::default();
99    }
100
101    // SAFETY: safety is ensured by the verifier.
102    let key = unsafe { EbpfBufferPtr::new(key.as_ptr::<u8>(), map.schema.key_size as usize) };
103    let key: MapKey = key.load();
104
105    // SAFETY: safety is ensured by the verifier.
106    let value = unsafe { EbpfBufferPtr::new(value.as_ptr::<u8>(), map.schema.value_size as usize) };
107    let flags = flags.as_u64();
108
109    C::on_map_access(context, map);
110
111    map.update(&key, value, flags).map(|_| 0).unwrap_or(u64::MAX).into()
112}
113
114fn bpf_map_delete_elem<C: MapsProgramContext>(
115    context: &mut C::RunContext<'_>,
116    map: BpfValue,
117    key: BpfValue,
118    _: BpfValue,
119    _: BpfValue,
120    _: BpfValue,
121) -> BpfValue {
122    // SAFETY: The `map` must be a reference to a `Map` object kept alive by the program itself.
123    let map: &Map = unsafe { &*map.as_ptr::<Map>() };
124
125    // TODO(https://fxbug.dev/496639039): This should be checked by the verifier.
126    if map.schema.map_type == bpf_map_type_BPF_MAP_TYPE_SK_STORAGE {
127        return BpfValue::default();
128    }
129
130    // SAFETY: safety is ensured by the verifier.
131    let key = unsafe { EbpfBufferPtr::new(key.as_ptr::<u8>(), map.schema.key_size as usize) };
132    let key: MapKey = key.load();
133
134    C::on_map_access(context, map);
135
136    map.delete(&key).map(|_| 0).unwrap_or(u64::MAX).into()
137}
138
139fn bpf_trace_printk<C: EbpfProgramContext>(
140    _context: &mut C::RunContext<'_>,
141    _fmt: BpfValue,
142    _fmt_size: BpfValue,
143    _: BpfValue,
144    _: BpfValue,
145    _: BpfValue,
146) -> BpfValue {
147    track_stub!(TODO("https://fxbug.dev/287120494"), "bpf_trace_printk");
148    0.into()
149}
150
151fn bpf_ktime_get_ns<C: EbpfProgramContext>(
152    _context: &mut C::RunContext<'_>,
153    _: BpfValue,
154    _: BpfValue,
155    _: BpfValue,
156    _: BpfValue,
157    _: BpfValue,
158) -> BpfValue {
159    zx::MonotonicInstant::get().into_nanos().into()
160}
161
162fn bpf_ringbuf_reserve<C: EbpfProgramContext>(
163    _context: &mut C::RunContext<'_>,
164    map: BpfValue,
165    size: BpfValue,
166    flags: BpfValue,
167    _: BpfValue,
168    _: BpfValue,
169) -> BpfValue {
170    // SAFETY: The safety of the operation is ensured by the bpf verifier. The `map` must be a
171    // reference to a `Map` object kept alive by the program itself.
172    let map: &Map = unsafe { &*map.as_ptr::<Map>() };
173
174    // Map type is checked by the verifier.
175    assert!(map.schema.map_type == bpf_map_type_BPF_MAP_TYPE_RINGBUF);
176
177    let Ok(size) = u32::try_from(size) else {
178        return BpfValue::default();
179    };
180    let flags = u64::from(flags);
181    map.ringbuf_reserve(size, flags).map(BpfValue::from).unwrap_or_else(|_| BpfValue::default())
182}
183
184fn bpf_ringbuf_submit<C: EbpfProgramContext>(
185    _context: &mut C::RunContext<'_>,
186    data: BpfValue,
187    flags: BpfValue,
188    _: BpfValue,
189    _: BpfValue,
190    _: BpfValue,
191) -> BpfValue {
192    let flags = RingBufferWakeupPolicy::from(flags);
193
194    // SAFETY: The safety of the operation is ensured by the bpf verifier. The data has to come from
195    // the result of a reserve call.
196    unsafe {
197        RingBuffer::submit(u64::from(data), flags);
198    }
199    0.into()
200}
201
202fn bpf_ringbuf_discard<C: EbpfProgramContext>(
203    _context: &mut C::RunContext<'_>,
204    data: BpfValue,
205    flags: BpfValue,
206    _: BpfValue,
207    _: BpfValue,
208    _: BpfValue,
209) -> BpfValue {
210    let flags = RingBufferWakeupPolicy::from(flags);
211
212    // SAFETY: The safety of the operation is ensured by the bpf verifier. The data has to come from
213    // the result of a reserve call.
214    unsafe {
215        RingBuffer::discard(u64::from(data), flags);
216    }
217    0.into()
218}
219
220fn bpf_ktime_get_boot_ns<C: EbpfProgramContext>(
221    _context: &mut C::RunContext<'_>,
222    _: BpfValue,
223    _: BpfValue,
224    _: BpfValue,
225    _: BpfValue,
226    _: BpfValue,
227) -> BpfValue {
228    track_stub!(TODO("https://fxbug.dev/287120494"), "bpf_ktime_get_boot_ns");
229    0.into()
230}
231
232fn bpf_probe_read_user<C: EbpfProgramContext>(
233    _context: &mut C::RunContext<'_>,
234    _: BpfValue,
235    _: BpfValue,
236    _: BpfValue,
237    _: BpfValue,
238    _: BpfValue,
239) -> BpfValue {
240    track_stub!(TODO("https://fxbug.dev/287120494"), "bpf_probe_read_user");
241    0.into()
242}
243
244fn bpf_probe_read_user_str<C: EbpfProgramContext>(
245    _context: &mut C::RunContext<'_>,
246    _: BpfValue,
247    _: BpfValue,
248    _: BpfValue,
249    _: BpfValue,
250    _: BpfValue,
251) -> BpfValue {
252    track_stub!(TODO("https://fxbug.dev/287120494"), "bpf_probe_read_user_str");
253    0.into()
254}
255
256fn bpf_ktime_get_coarse_ns<C: EbpfProgramContext>(
257    _context: &mut C::RunContext<'_>,
258    _: BpfValue,
259    _: BpfValue,
260    _: BpfValue,
261    _: BpfValue,
262    _: BpfValue,
263) -> BpfValue {
264    track_stub!(TODO("https://fxbug.dev/287120494"), "bpf_ktime_get_coarse_ns");
265    0.into()
266}
267
268fn bpf_probe_read_str<C: EbpfProgramContext>(
269    _context: &mut C::RunContext<'_>,
270    _: BpfValue,
271    _: BpfValue,
272    _: BpfValue,
273    _: BpfValue,
274    _: BpfValue,
275) -> BpfValue {
276    track_stub!(TODO("https://fxbug.dev/287120494"), "bpf_probe_read_str");
277    0.into()
278}
279
280fn bpf_get_smp_processor_id<C: EbpfProgramContext>(
281    _context: &mut C::RunContext<'_>,
282    _: BpfValue,
283    _: BpfValue,
284    _: BpfValue,
285    _: BpfValue,
286    _: BpfValue,
287) -> BpfValue {
288    track_stub!(TODO("https://fxbug.dev/287120494"), "bpf_get_smp_processor_id");
289    0.into()
290}
291
292pub trait CurrentTaskContext {
293    fn get_uid_gid(&self) -> (uid_t, gid_t);
294    fn get_tid_tgid(&self) -> (pid_t, pid_t);
295}
296
297pub trait CurrentTaskProgramContext: EbpfProgramContext {
298    fn get_uid_gid<'a>(context: &mut Self::RunContext<'a>) -> (uid_t, gid_t);
299    fn get_tid_tgid<'a>(context: &mut Self::RunContext<'a>) -> (pid_t, pid_t);
300}
301
302impl<C: EbpfProgramContext> CurrentTaskProgramContext for C
303where
304    for<'a> C::RunContext<'a>: CurrentTaskContext,
305{
306    fn get_uid_gid<'a>(context: &mut Self::RunContext<'a>) -> (uid_t, gid_t) {
307        context.get_uid_gid()
308    }
309    fn get_tid_tgid<'a>(context: &mut Self::RunContext<'a>) -> (pid_t, pid_t) {
310        context.get_tid_tgid()
311    }
312}
313
314fn bpf_get_current_uid_gid<C: CurrentTaskProgramContext>(
315    context: &mut C::RunContext<'_>,
316    _: BpfValue,
317    _: BpfValue,
318    _: BpfValue,
319    _: BpfValue,
320    _: BpfValue,
321) -> BpfValue {
322    let (uid, gid) = C::get_uid_gid(context);
323    (uid as u64 | (gid as u64) << 32).into()
324}
325
326fn bpf_get_current_pid_tgid<C: CurrentTaskProgramContext>(
327    context: &mut C::RunContext<'_>,
328    _: BpfValue,
329    _: BpfValue,
330    _: BpfValue,
331    _: BpfValue,
332    _: BpfValue,
333) -> BpfValue {
334    let (pid, tgid) = C::get_tid_tgid(context);
335    (pid as u64 | (tgid as u64) << 32).into()
336}
337
338// Trait for `EbpfProgramContext` where the first argument is a `SocketRef`,
339// i.e. it references a socket.
340pub trait Arg1IsSocketProgramContext: EbpfProgramContext {
341    type Arg1AsSocket<'a>: FromBpfValue<Self::RunContext<'a>> + SocketRef;
342}
343
344impl<C> Arg1IsSocketProgramContext for C
345where
346    C: EbpfProgramContext,
347    for<'a> Self::Arg1<'a>: FromBpfValue<Self::RunContext<'a>> + SocketRef,
348{
349    type Arg1AsSocket<'a> = Self::Arg1<'a>;
350}
351
352// Marker trait for `EbpfProgramContext` that supports `bpf_get_socket_uid`.
353pub trait SocketCookieProgramContext: Arg1IsSocketProgramContext {}
354impl<C> SocketCookieProgramContext for C where C: Arg1IsSocketProgramContext {}
355
356fn bpf_get_socket_cookie<'a, C: SocketCookieProgramContext>(
357    context: &mut C::RunContext<'a>,
358    arg1: BpfValue,
359    _: BpfValue,
360    _: BpfValue,
361    _: BpfValue,
362    _: BpfValue,
363) -> BpfValue {
364    // SAFETY: Verifier checks that the argument points at the value that
365    // that's passed as the first argument.
366    let arg1_as_socket = unsafe { C::Arg1AsSocket::from_bpf_value(context, arg1) };
367    arg1_as_socket.get_socket_cookie().unwrap_or(0).into()
368}
369
370pub trait SocketRef {
371    fn get_socket_cookie(&self) -> Option<u64>;
372    fn get_socket_uid(&self) -> Option<uid_t>;
373}
374
375// A trait for eBPF run context with `bpf_sock` pointers.
376pub trait BpfSockContext: Sized {
377    type BpfSockRef: SocketRef + FromBpfValue<Self>;
378}
379
380pub trait SkStorageProgramContext: EbpfProgramContext {
381    type BpfSockRef<'a>: SocketRef + FromBpfValue<Self::RunContext<'a>>;
382}
383
384impl<C> SkStorageProgramContext for C
385where
386    C: EbpfProgramContext,
387    for<'a> C::RunContext<'a>: BpfSockContext,
388{
389    type BpfSockRef<'a> = <C::RunContext<'a> as BpfSockContext>::BpfSockRef;
390}
391
392#[derive(Copy, Clone, Debug, PartialEq, Eq)]
393pub enum LoadBytesBase {
394    MacHeader,
395    NetworkHeader,
396}
397
398// Marker trait for `EbpfProgramContext` that supports `bpf_get_socket_uid`.
399pub trait SocketUidProgramContext: Arg1IsSocketProgramContext {}
400impl<C> SocketUidProgramContext for C where C: Arg1IsSocketProgramContext {}
401
402fn bpf_get_socket_uid<'a, C: SocketUidProgramContext>(
403    context: &mut C::RunContext<'a>,
404    sk_buf: BpfValue,
405    _: BpfValue,
406    _: BpfValue,
407    _: BpfValue,
408    _: BpfValue,
409) -> BpfValue {
410    const OVERFLOW_UID: uid_t = 65534;
411    // SAFETY: Verifier checks that the first argument points at a `__sk_buff`.
412    let sk_buf = unsafe { C::Arg1AsSocket::from_bpf_value(context, sk_buf) };
413    sk_buf.get_socket_uid().unwrap_or(OVERFLOW_UID).into()
414}
415
416// Trait for packets that support `bpf_load_bytes_relative`.
417pub trait PacketWithLoadBytes {
418    fn load_bytes_relative(
419        &self,
420        base: LoadBytesBase,
421        offset: usize,
422        buf: EbpfBufferPtr<'_>,
423    ) -> i64;
424}
425
426// Trait for `EbpfProgramContext` that supports `bpf_load_bytes_relative`.
427pub trait SkbLoadBytesProgramContext: EbpfProgramContext {
428    fn skb_load_bytes_relative<'a>(
429        context: &mut Self::RunContext<'a>,
430        sk_buf: BpfValue,
431        base: LoadBytesBase,
432        offset: usize,
433        buf: EbpfBufferPtr<'_>,
434    ) -> i64;
435}
436
437impl<C: EbpfProgramContext> SkbLoadBytesProgramContext for C
438where
439    for<'b> C::Arg1<'b>: FromBpfValue<C::RunContext<'b>>,
440    for<'b> C::Arg1<'b>: PacketWithLoadBytes,
441{
442    fn skb_load_bytes_relative<'a>(
443        context: &mut Self::RunContext<'a>,
444        sk_buf: BpfValue,
445        base: LoadBytesBase,
446        offset: usize,
447        buf: EbpfBufferPtr<'_>,
448    ) -> i64 {
449        // SAFETY: Verifier checks that the argument points at the same value
450        // that was passed to the program as the first argument.
451        let sk_buf = unsafe { C::Arg1::from_bpf_value(context, sk_buf) };
452        sk_buf.load_bytes_relative(base, offset, buf)
453    }
454}
455
456fn bpf_skb_load_bytes<'a, C: SkbLoadBytesProgramContext>(
457    context: &mut C::RunContext<'a>,
458    sk_buf: BpfValue,
459    offset: BpfValue,
460    to: BpfValue,
461    len: BpfValue,
462    _: BpfValue,
463) -> BpfValue {
464    let base = LoadBytesBase::NetworkHeader;
465
466    let Ok(offset) = offset.as_u64().try_into() else {
467        return u64::MAX.into();
468    };
469
470    // SAFETY: The verifier ensures that `to` points to a valid buffer of at
471    // least `len` bytes that the eBPF program has permission to access.
472    let buf = unsafe { EbpfBufferPtr::new(to.as_ptr::<u8>(), len.as_u64() as usize) };
473
474    C::skb_load_bytes_relative(context, sk_buf, base, offset, buf).into()
475}
476
477fn bpf_skb_load_bytes_relative<'a, C: SkbLoadBytesProgramContext>(
478    context: &mut C::RunContext<'a>,
479    sk_buf: BpfValue,
480    offset: BpfValue,
481    to: BpfValue,
482    len: BpfValue,
483    start_header: BpfValue,
484) -> BpfValue {
485    let base = match start_header.as_u64() {
486        0 => LoadBytesBase::MacHeader,
487        1 => LoadBytesBase::NetworkHeader,
488        _ => return u64::MAX.into(),
489    };
490
491    let Ok(offset) = offset.as_u64().try_into() else {
492        return u64::MAX.into();
493    };
494
495    // SAFETY: The verifier ensures that `to` points to a valid buffer of at
496    // least `len` bytes that the eBPF program has permission to access.
497    let buf = unsafe { EbpfBufferPtr::new(to.as_ptr::<u8>(), len.as_u64() as usize) };
498
499    C::skb_load_bytes_relative(context, sk_buf, base, offset, buf).into()
500}
501
502fn bpf_sk_storage_get<'a, C: SkStorageProgramContext + MapsProgramContext>(
503    context: &mut C::RunContext<'a>,
504    map: BpfValue,
505    sk: BpfValue,
506    value: BpfValue,
507    flags: BpfValue,
508    _: BpfValue,
509) -> BpfValue {
510    if sk.is_zero() {
511        return BpfValue::default();
512    }
513
514    // SAFETY: Verifier ensures that `sk` is either null or a pointer to
515    // `bpf_sock`. The null case is checked above.
516    let bpf_sock = unsafe { C::BpfSockRef::from_bpf_value(context, sk) };
517
518    // Use socket cookie to identify the socket in the map.
519    let Some(socket_id) = bpf_sock.get_socket_cookie() else {
520        return BpfValue::default();
521    };
522
523    let key = socket_id.as_bytes();
524
525    // SAFETY: The `map` must be a reference to a `Map` object kept alive by the program itself.
526    let map: &Map = unsafe { &*map.as_ptr::<Map>() };
527
528    // Checked by the verifier.
529    assert!(map.schema.map_type == bpf_map_type_BPF_MAP_TYPE_SK_STORAGE);
530
531    C::on_map_access(context, map);
532
533    if let Some(value_ref) = map.lookup(key) {
534        let result: BpfValue = value_ref.ptr().raw_ptr().into();
535        C::add_value_ref(context, value_ref);
536        return result;
537    }
538
539    if flags.as_u32() & BPF_SK_STORAGE_GET_F_CREATE != 0 {
540        let mut vec;
541        let init_val = if value.as_u64() == 0 {
542            vec = SmallVec::<[u8; 128]>::new();
543            vec.resize(map.schema.value_size as usize, 0);
544            (&mut vec[..]).into()
545        } else {
546            // SAFETY: The verifier ensures that `value` points to a valid buffer.
547            unsafe { EbpfBufferPtr::new(value.as_ptr::<u8>(), map.schema.value_size as usize) }
548        };
549
550        let r = map.update(key, init_val, 0);
551        if r.is_ok() {
552            if let Some(value_ref) = map.lookup(key) {
553                let result: BpfValue = value_ref.ptr().raw_ptr().into();
554                C::add_value_ref(context, value_ref);
555                return result;
556            }
557        }
558    }
559
560    BpfValue::default()
561}
562
563fn bpf_sk_fullsock<C: EbpfProgramContext>(
564    _context: &mut C::RunContext<'_>,
565    _: BpfValue,
566    _: BpfValue,
567    _: BpfValue,
568    _: BpfValue,
569    _: BpfValue,
570) -> BpfValue {
571    track_stub!(TODO("https://fxbug.dev/287120494"), "bpf_sk_fullsock");
572    0.into()
573}
574
575pub trait ReturnValueContext {
576    fn set_retval(&mut self, value: i32) -> i32;
577    fn get_retval(&self) -> i32;
578}
579
580pub trait ReturnValueProgramContext: EbpfProgramContext {
581    fn set_retval<'a>(context: &mut Self::RunContext<'a>, value: i32) -> i32;
582    fn get_retval<'a>(context: &mut Self::RunContext<'a>) -> i32;
583}
584
585impl<C: EbpfProgramContext> ReturnValueProgramContext for C
586where
587    for<'a> C::RunContext<'a>: ReturnValueContext,
588{
589    fn set_retval<'a>(context: &mut Self::RunContext<'a>, value: i32) -> i32 {
590        context.set_retval(value)
591    }
592    fn get_retval<'a>(context: &mut Self::RunContext<'a>) -> i32 {
593        context.get_retval()
594    }
595}
596
597fn bpf_set_retval<C: ReturnValueProgramContext>(
598    context: &mut C::RunContext<'_>,
599    value: BpfValue,
600    _: BpfValue,
601    _: BpfValue,
602    _: BpfValue,
603    _: BpfValue,
604) -> BpfValue {
605    C::set_retval(context, value.as_i32()).into()
606}
607
608fn bpf_get_retval<C: ReturnValueProgramContext>(
609    context: &mut C::RunContext<'_>,
610    _: BpfValue,
611    _: BpfValue,
612    _: BpfValue,
613    _: BpfValue,
614    _: BpfValue,
615) -> BpfValue {
616    C::get_retval(context).into()
617}
618
619fn bpf_sk_lookup_tcp<C: EbpfProgramContext>(
620    _context: &mut C::RunContext<'_>,
621    _: BpfValue,
622    _: BpfValue,
623    _: BpfValue,
624    _: BpfValue,
625    _: BpfValue,
626) -> BpfValue {
627    track_stub!(TODO("https://fxbug.dev/287120494"), "bpf_sk_lookup_tcp");
628    0.into()
629}
630
631fn bpf_sk_lookup_udp<C: EbpfProgramContext>(
632    _context: &mut C::RunContext<'_>,
633    _: BpfValue,
634    _: BpfValue,
635    _: BpfValue,
636    _: BpfValue,
637    _: BpfValue,
638) -> BpfValue {
639    track_stub!(TODO("https://fxbug.dev/287120494"), "bpf_sk_lookup_udp");
640    0.into()
641}
642
643fn bpf_sk_release<C: EbpfProgramContext>(
644    _context: &mut C::RunContext<'_>,
645    _: BpfValue,
646    _: BpfValue,
647    _: BpfValue,
648    _: BpfValue,
649    _: BpfValue,
650) -> BpfValue {
651    track_stub!(TODO("https://fxbug.dev/287120494"), "bpf_sk_release");
652    0.into()
653}
654
655fn bpf_get_netns_cookie<C: EbpfProgramContext>(
656    _context: &mut C::RunContext<'_>,
657    _: BpfValue,
658    _: BpfValue,
659    _: BpfValue,
660    _: BpfValue,
661    _: BpfValue,
662) -> BpfValue {
663    track_stub!(TODO("https://fxbug.dev/287120494"), "bpf_get_netns_cookie");
664    const DEFAULT_NETWORK_NAMESPACE_COOKIE: u64 = 1;
665    DEFAULT_NETWORK_NAMESPACE_COOKIE.into()
666}
667
668fn get_common_helpers<C: MapsProgramContext>() -> impl Iterator<Item = (u32, EbpfHelperImpl<C>)> {
669    [
670        (bpf_func_id_BPF_FUNC_ktime_get_boot_ns, EbpfHelperImpl(bpf_ktime_get_boot_ns)),
671        (bpf_func_id_BPF_FUNC_ktime_get_coarse_ns, EbpfHelperImpl(bpf_ktime_get_coarse_ns)),
672        (bpf_func_id_BPF_FUNC_ktime_get_ns, EbpfHelperImpl(bpf_ktime_get_ns)),
673        (bpf_func_id_BPF_FUNC_map_delete_elem, EbpfHelperImpl(bpf_map_delete_elem)),
674        (bpf_func_id_BPF_FUNC_map_lookup_elem, EbpfHelperImpl(bpf_map_lookup_elem)),
675        (bpf_func_id_BPF_FUNC_map_update_elem, EbpfHelperImpl(bpf_map_update_elem)),
676        (bpf_func_id_BPF_FUNC_probe_read_str, EbpfHelperImpl(bpf_probe_read_str)),
677        (bpf_func_id_BPF_FUNC_probe_read_user, EbpfHelperImpl(bpf_probe_read_user)),
678        (bpf_func_id_BPF_FUNC_probe_read_user_str, EbpfHelperImpl(bpf_probe_read_user_str)),
679        (bpf_func_id_BPF_FUNC_ringbuf_discard, EbpfHelperImpl(bpf_ringbuf_discard)),
680        (bpf_func_id_BPF_FUNC_ringbuf_reserve, EbpfHelperImpl(bpf_ringbuf_reserve)),
681        (bpf_func_id_BPF_FUNC_ringbuf_submit, EbpfHelperImpl(bpf_ringbuf_submit)),
682        (bpf_func_id_BPF_FUNC_trace_printk, EbpfHelperImpl(bpf_trace_printk)),
683        (bpf_func_id_BPF_FUNC_get_smp_processor_id, EbpfHelperImpl(bpf_get_smp_processor_id)),
684    ]
685    .into_iter()
686}
687
688/// Returns helper implementations that depend on `CurrentTask`.
689fn get_current_task_helpers<C: CurrentTaskProgramContext>()
690-> impl Iterator<Item = (u32, EbpfHelperImpl<C>)> {
691    [
692        (bpf_func_id_BPF_FUNC_get_current_uid_gid, EbpfHelperImpl(bpf_get_current_uid_gid)),
693        (bpf_func_id_BPF_FUNC_get_current_pid_tgid, EbpfHelperImpl(bpf_get_current_pid_tgid)),
694    ]
695    .into_iter()
696}
697
698// Trait for `EbpfProgramContext` implementations that are used for
699// `BPF_PROG_TYPE_CGROUP_SOCK` programs.
700pub trait CgroupSockProgramContext:
701    MapsProgramContext
702    + SocketCookieProgramContext
703    + CurrentTaskProgramContext
704    + SkStorageProgramContext
705{
706    fn get_helpers() -> HelperSet<Self> {
707        [
708            (bpf_func_id_BPF_FUNC_get_netns_cookie, EbpfHelperImpl(bpf_get_netns_cookie)),
709            (bpf_func_id_BPF_FUNC_get_socket_cookie, EbpfHelperImpl(bpf_get_socket_cookie)),
710            (bpf_func_id_BPF_FUNC_sk_storage_get, EbpfHelperImpl(bpf_sk_storage_get)),
711            (bpf_func_id_BPF_FUNC_sk_lookup_tcp, EbpfHelperImpl(bpf_sk_lookup_tcp)),
712            (bpf_func_id_BPF_FUNC_sk_lookup_udp, EbpfHelperImpl(bpf_sk_lookup_udp)),
713            (bpf_func_id_BPF_FUNC_sk_release, EbpfHelperImpl(bpf_sk_release)),
714        ]
715        .into_iter()
716        .chain(get_common_helpers())
717        .chain(get_current_task_helpers())
718        .collect()
719    }
720}
721
722// Trait for `EbpfProgramContext` implementations that are used for
723// `BPF_PROG_TYPE_CGROUP_SOCKADDR` programs.
724pub trait CgroupSockAddrProgramContext:
725    MapsProgramContext
726    + SocketCookieProgramContext
727    + CurrentTaskProgramContext
728    + SkStorageProgramContext
729{
730    fn get_helpers() -> HelperSet<Self> {
731        [
732            (bpf_func_id_BPF_FUNC_get_netns_cookie, EbpfHelperImpl(bpf_get_netns_cookie)),
733            (bpf_func_id_BPF_FUNC_get_socket_cookie, EbpfHelperImpl(bpf_get_socket_cookie)),
734            (bpf_func_id_BPF_FUNC_sk_storage_get, EbpfHelperImpl(bpf_sk_storage_get)),
735            (bpf_func_id_BPF_FUNC_sk_lookup_tcp, EbpfHelperImpl(bpf_sk_lookup_tcp)),
736            (bpf_func_id_BPF_FUNC_sk_lookup_udp, EbpfHelperImpl(bpf_sk_lookup_udp)),
737            (bpf_func_id_BPF_FUNC_sk_release, EbpfHelperImpl(bpf_sk_release)),
738        ]
739        .into_iter()
740        .chain(get_common_helpers())
741        .chain(get_current_task_helpers())
742        .collect()
743    }
744}
745
746// Trait for `EbpfProgramContext` implementations that are used for
747// `BPF_PROG_TYPE_CGROUP_SOCKOPT` programs.
748pub trait CgroupSockOptProgramContext:
749    MapsProgramContext + CurrentTaskProgramContext + ReturnValueProgramContext + SkStorageProgramContext
750{
751    fn get_helpers() -> HelperSet<Self> {
752        [
753            (bpf_func_id_BPF_FUNC_get_netns_cookie, EbpfHelperImpl(bpf_get_netns_cookie)),
754            (bpf_func_id_BPF_FUNC_set_retval, EbpfHelperImpl(bpf_set_retval)),
755            (bpf_func_id_BPF_FUNC_get_retval, EbpfHelperImpl(bpf_get_retval)),
756            (bpf_func_id_BPF_FUNC_sk_storage_get, EbpfHelperImpl(bpf_sk_storage_get)),
757            (bpf_func_id_BPF_FUNC_sk_lookup_tcp, EbpfHelperImpl(bpf_sk_lookup_tcp)),
758            (bpf_func_id_BPF_FUNC_sk_lookup_udp, EbpfHelperImpl(bpf_sk_lookup_udp)),
759            (bpf_func_id_BPF_FUNC_sk_release, EbpfHelperImpl(bpf_sk_release)),
760        ]
761        .into_iter()
762        .chain(get_common_helpers())
763        .chain(get_current_task_helpers())
764        .collect()
765    }
766}
767
768// Trait for `EbpfProgramContext` implementations that are used for
769// `BPF_PROG_TYPE_SOCKET_FILTER` programs.
770pub trait SocketFilterProgramContext:
771    MapsProgramContext
772    + SocketUidProgramContext
773    + SocketCookieProgramContext
774    + SkbLoadBytesProgramContext
775{
776    fn get_helpers() -> HelperSet<Self> {
777        vec![
778            (bpf_func_id_BPF_FUNC_get_netns_cookie, EbpfHelperImpl(bpf_get_netns_cookie)),
779            (bpf_func_id_BPF_FUNC_get_socket_uid, EbpfHelperImpl(bpf_get_socket_uid)),
780            (bpf_func_id_BPF_FUNC_get_socket_cookie, EbpfHelperImpl(bpf_get_socket_cookie)),
781            (bpf_func_id_BPF_FUNC_skb_load_bytes, EbpfHelperImpl(bpf_skb_load_bytes)),
782            (
783                bpf_func_id_BPF_FUNC_skb_load_bytes_relative,
784                EbpfHelperImpl(bpf_skb_load_bytes_relative),
785            ),
786        ]
787        .into_iter()
788        .chain(get_common_helpers())
789        .collect()
790    }
791}
792
793// Trait for `EbpfProgramContext` implementations that are used for
794// `BPF_PROG_TYPE_CGROUP_SKB` programs.
795pub trait CgroupSkbProgramContext:
796    MapsProgramContext
797    + SocketUidProgramContext
798    + SocketCookieProgramContext
799    + SkbLoadBytesProgramContext
800    + SkStorageProgramContext
801{
802    fn get_helpers() -> HelperSet<Self> {
803        vec![
804            (bpf_func_id_BPF_FUNC_get_netns_cookie, EbpfHelperImpl(bpf_get_netns_cookie)),
805            (bpf_func_id_BPF_FUNC_get_socket_uid, EbpfHelperImpl(bpf_get_socket_uid)),
806            (bpf_func_id_BPF_FUNC_get_socket_cookie, EbpfHelperImpl(bpf_get_socket_cookie)),
807            (bpf_func_id_BPF_FUNC_skb_load_bytes, EbpfHelperImpl(bpf_skb_load_bytes)),
808            (
809                bpf_func_id_BPF_FUNC_skb_load_bytes_relative,
810                EbpfHelperImpl(bpf_skb_load_bytes_relative),
811            ),
812            (bpf_func_id_BPF_FUNC_sk_storage_get, EbpfHelperImpl(bpf_sk_storage_get)),
813            (bpf_func_id_BPF_FUNC_sk_lookup_tcp, EbpfHelperImpl(bpf_sk_lookup_tcp)),
814            (bpf_func_id_BPF_FUNC_sk_lookup_udp, EbpfHelperImpl(bpf_sk_lookup_udp)),
815            (bpf_func_id_BPF_FUNC_sk_release, EbpfHelperImpl(bpf_sk_release)),
816            (bpf_func_id_BPF_FUNC_sk_fullsock, EbpfHelperImpl(bpf_sk_fullsock)),
817        ]
818        .into_iter()
819        .chain(get_common_helpers())
820        .collect()
821    }
822}
823
824/// Macro used to declare program type for a `EbpfProgramContext` implementation.
825/// Implements `StaticHelperSet` trait for the context type.
826///
827/// # Example
828///
829/// The following example declares that `MyEbpfProgramContext` is used to run
830/// socket filter programs:
831///
832/// ```
833/// ebpf_program_context_type!(MyEbpfProgramContext, SocketFilterProgramContext);
834/// ```
835#[macro_export]
836macro_rules! ebpf_program_context_type {
837    ($context:ty, $subtrait:ty) => {
838        impl $subtrait for $context {}
839        ebpf::static_helper_set!($context, <$context as $subtrait>::get_helpers());
840    };
841}
842
843#[cfg(test)]
844mod tests {
845    use super::*;
846    use crate::maps::{Map, PinnedMap};
847    use ebpf::{BpfValue, EbpfProgramContext, FromBpfValue, MapFlags, MapSchema};
848    use linux_uapi::{BPF_SK_STORAGE_GET_F_CREATE, bpf_map_type_BPF_MAP_TYPE_SK_STORAGE};
849
850    struct MockSocket {
851        cookie: u64,
852    }
853    impl SocketRef for MockSocket {
854        fn get_socket_cookie(&self) -> Option<u64> {
855            Some(self.cookie)
856        }
857        fn get_socket_uid(&self) -> Option<uid_t> {
858            Some(0)
859        }
860    }
861    impl<'a> FromBpfValue<TestRunContext<'a>> for MockSocket {
862        unsafe fn from_bpf_value(_context: &mut TestRunContext<'a>, value: BpfValue) -> Self {
863            Self { cookie: value.as_u64() }
864        }
865    }
866
867    struct TestRunContext<'a> {
868        map_refs: Vec<MapValueRef<'a>>,
869    }
870    impl<'a> BpfSockContext for TestRunContext<'a> {
871        type BpfSockRef = MockSocket;
872    }
873    impl<'a> MapsContext<'a> for TestRunContext<'a> {
874        fn on_map_access(&mut self, _map: &Map) {}
875        fn add_value_ref(&mut self, map_ref: MapValueRef<'a>) {
876            self.map_refs.push(map_ref);
877        }
878    }
879
880    struct TestContext;
881    impl EbpfProgramContext for TestContext {
882        type RunContext<'a> = TestRunContext<'a>;
883        type Packet<'a> = ();
884        type Arg1<'a> = ();
885        type Arg2<'a> = ();
886        type Arg3<'a> = ();
887        type Arg4<'a> = ();
888        type Arg5<'a> = ();
889        type Map = PinnedMap;
890    }
891
892    #[fuchsia::test]
893    fn test_sk_storage_get_uaf() {
894        let schema = MapSchema {
895            map_type: bpf_map_type_BPF_MAP_TYPE_SK_STORAGE,
896            key_size: 4,
897            value_size: 8,
898            max_entries: 0,
899            flags: MapFlags::NoPrealloc,
900        };
901        let map = Map::new(schema, "test").unwrap();
902        let map_value = BpfValue::from(&*map as *const Map);
903
904        let mut context = TestRunContext { map_refs: vec![] };
905
906        // 1. Create entry for socket 42
907        let sk_value1 = BpfValue::from(42u64);
908        let init_value1 = [0x11u8; 8];
909        let init_value_ptr1 = BpfValue::from(init_value1.as_ptr());
910        let flags = BpfValue::from(BPF_SK_STORAGE_GET_F_CREATE as u64);
911
912        let ptr1 = bpf_sk_storage_get::<TestContext>(
913            &mut context,
914            map_value,
915            sk_value1,
916            init_value_ptr1,
917            flags,
918            BpfValue::default(),
919        );
920        assert!(!ptr1.is_zero());
921
922        // Verify initial value
923        // SAFETY: ptr1 is a valid pointer to the map value.
924        unsafe {
925            assert_eq!(*(ptr1.as_ptr::<u64>()), 0x1111111111111111);
926        }
927
928        // 2. Delete entry for socket 42 from map
929        let key_bytes = 42u64.to_ne_bytes();
930        map.delete(&key_bytes).unwrap();
931
932        // 3. Create entry for socket 43
933        // If UAF exists, this should reuse the same memory block because it was freed.
934        let sk_value2 = BpfValue::from(43u64);
935        let init_value2 = [0x22u8; 8];
936        let init_value_ptr2 = BpfValue::from(init_value2.as_ptr());
937
938        let ptr2 = bpf_sk_storage_get::<TestContext>(
939            &mut context,
940            map_value,
941            sk_value2,
942            init_value_ptr2,
943            flags,
944            BpfValue::default(),
945        );
946        assert!(!ptr2.is_zero());
947
948        // We want to assert that the value at ptr1 has NOT changed, which means it was not reused.
949        // This assertion will FAIL without the fix (UAF occurs,
950        // ptr1's memory is overwritten with ptr2's init value),
951        // and PASS with the fix (ptr1's memory is kept alive).
952        // SAFETY: ptr1 points to memory that is kept alive by the reference in `context`.
953        unsafe {
954            assert_eq!(*(ptr1.as_ptr::<u64>()), 0x1111111111111111);
955        }
956    }
957
958    #[fuchsia::test]
959    fn test_sk_storage_get_uaf_query() {
960        let schema = MapSchema {
961            map_type: bpf_map_type_BPF_MAP_TYPE_SK_STORAGE,
962            key_size: 4,
963            value_size: 8,
964            max_entries: 0,
965            flags: MapFlags::NoPrealloc,
966        };
967        let map = Map::new(schema, "test").unwrap();
968        let map_value = BpfValue::from(&*map as *const Map);
969
970        let mut context = TestRunContext { map_refs: vec![] };
971
972        // 1. Create entry for socket 42
973        let sk_value1 = BpfValue::from(42u64);
974        let init_value1 = [0x11u8; 8];
975        let init_value_ptr1 = BpfValue::from(init_value1.as_ptr());
976        let flags = BpfValue::from(BPF_SK_STORAGE_GET_F_CREATE as u64);
977
978        let ptr1 = bpf_sk_storage_get::<TestContext>(
979            &mut context,
980            map_value,
981            sk_value1,
982            init_value_ptr1,
983            flags,
984            BpfValue::default(),
985        );
986        assert!(!ptr1.is_zero());
987
988        // Clear context to simulate that we don't hold the creation
989        // reference anymore. The map still holds the reference.
990        context.map_refs.clear();
991
992        // 2. Query entry for socket 42 (without CREATE flag)
993        let ptr1_query = bpf_sk_storage_get::<TestContext>(
994            &mut context,
995            map_value,
996            sk_value1,
997            BpfValue::default(),
998            BpfValue::default(),
999            BpfValue::default(),
1000        );
1001        assert_eq!(ptr1.as_u64(), ptr1_query.as_u64());
1002
1003        // 3. Delete entry for socket 42 from map
1004        let key_bytes = 42u64.to_ne_bytes();
1005        map.delete(&key_bytes).unwrap();
1006
1007        // 4. Create entry for socket 43
1008        // If UAF exists, this should reuse the same memory block
1009        // because it was freed.
1010        let sk_value2 = BpfValue::from(43u64);
1011        let init_value2 = [0x22u8; 8];
1012        let init_value_ptr2 = BpfValue::from(init_value2.as_ptr());
1013
1014        let ptr2 = bpf_sk_storage_get::<TestContext>(
1015            &mut context,
1016            map_value,
1017            sk_value2,
1018            init_value_ptr2,
1019            flags,
1020            BpfValue::default(),
1021        );
1022        assert!(!ptr2.is_zero());
1023
1024        // We want to assert that the value at ptr1_query has NOT
1025        // changed.
1026        // SAFETY: ptr1_query points to memory that is kept alive by
1027        // the reference in `context` (from the query).
1028        unsafe {
1029            assert_eq!(*(ptr1_query.as_ptr::<u64>()), 0x1111111111111111);
1030        }
1031    }
1032}