From 9bb64b13f7ea4c4e8cde21a519dad23d29573508 Mon Sep 17 00:00:00 2001
From: Dave Tucker <dave@dtucker.co.uk>
Date: Tue, 1 Aug 2023 17:23:15 +0100
Subject: [PATCH] aya: Set BPF_F_SLEEPABLE for sleepable programs

Signed-off-by: Dave Tucker <dave@dtucker.co.uk>
---
 aya/src/bpf.rs | 52 ++++++++++++++++++++++++++++++++++++--------------
 1 file changed, 38 insertions(+), 14 deletions(-)

diff --git a/aya/src/bpf.rs b/aya/src/bpf.rs
index eed006d4..1edf9574 100644
--- a/aya/src/bpf.rs
+++ b/aya/src/bpf.rs
@@ -562,14 +562,28 @@ impl<'a> BpfLoader<'a> {
                             data: ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level),
                             kind: ProbeKind::KRetProbe,
                         }),
-                        ProgramSection::UProbe { .. } => Program::UProbe(UProbe {
-                            data: ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level),
-                            kind: ProbeKind::UProbe,
-                        }),
-                        ProgramSection::URetProbe { .. } => Program::UProbe(UProbe {
-                            data: ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level),
-                            kind: ProbeKind::URetProbe,
-                        }),
+                        ProgramSection::UProbe { sleepable, .. } => {
+                            let mut data =
+                                ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level);
+                            if *sleepable {
+                                data.flags = BPF_F_SLEEPABLE;
+                            }
+                            Program::UProbe(UProbe {
+                                data,
+                                kind: ProbeKind::UProbe,
+                            })
+                        }
+                        ProgramSection::URetProbe { sleepable, .. } => {
+                            let mut data =
+                                ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level);
+                            if *sleepable {
+                                data.flags = BPF_F_SLEEPABLE;
+                            }
+                            Program::UProbe(UProbe {
+                                data,
+                                kind: ProbeKind::URetProbe,
+                            })
+                        }
                         ProgramSection::TracePoint { .. } => Program::TracePoint(TracePoint {
                             data: ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level),
                         }),
@@ -662,12 +676,22 @@ impl<'a> BpfLoader<'a> {
                                 data: ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level),
                             })
                         }
-                        ProgramSection::FEntry { .. } => Program::FEntry(FEntry {
-                            data: ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level),
-                        }),
-                        ProgramSection::FExit { .. } => Program::FExit(FExit {
-                            data: ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level),
-                        }),
+                        ProgramSection::FEntry { sleepable, .. } => {
+                            let mut data =
+                                ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level);
+                            if *sleepable {
+                                data.flags = BPF_F_SLEEPABLE;
+                            }
+                            Program::FEntry(FEntry { data })
+                        }
+                        ProgramSection::FExit { sleepable, .. } => {
+                            let mut data =
+                                ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level);
+                            if *sleepable {
+                                data.flags = BPF_F_SLEEPABLE;
+                            }
+                            Program::FExit(FExit { data })
+                        }
                         ProgramSection::Extension { .. } => Program::Extension(Extension {
                             data: ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level),
                         }),