embassy-executor: unsafe tasks as unsafe

4 files changed, 41 insertions, 1 deletions

diff --git a/embassy-executor-macros/src/macros/task.rs b/embassy-executor-macros/src/macros/task.rs
index 1c5e3571d..f01cc3b6c 100644
--- a/embassy-executor-macros/src/macros/task.rs
+++ b/embassy-executor-macros/src/macros/task.rs
@@ -120,6 +120,18 @@ pub fn run(args: TokenStream, item: TokenStream) -> TokenStream {
     task_inner.vis = syn::Visibility::Inherited;
     task_inner.sig.ident = task_inner_ident.clone();
 
+    // Forcefully mark the inner task as safe.
+    // SAFETY: We only ever call task_inner in functions
+    //         with the same safety preconditions as task_inner
+    task_inner.sig.unsafety = None;
+    let task_body = task_inner.body;
+    task_inner.body = quote! {
+        #[allow(unused_unsafe, reason = "Not all function bodies may require being in an unsafe block")]
+        unsafe {
+            #task_body
+        }
+    };
+
     // assemble the original input arguments,
     // including any attributes that may have
     // been applied previously
@@ -186,6 +198,7 @@ pub fn run(args: TokenStream, item: TokenStream) -> TokenStream {
     // Copy the generics + where clause to avoid more spurious errors.
     let generics = &f.sig.generics;
     let where_clause = &f.sig.generics.where_clause;
+    let unsafety = &f.sig.unsafety;
 
     let result = quote! {
         // This is the user's task function, renamed.
@@ -196,7 +209,7 @@ pub fn run(args: TokenStream, item: TokenStream) -> TokenStream {
         #task_inner
 
         #(#task_outer_attrs)*
-        #visibility fn #task_ident #generics (#fargs) -> #embassy_executor::SpawnToken<impl Sized> #where_clause{
+        #visibility #unsafety fn #task_ident #generics (#fargs) -> #embassy_executor::SpawnToken<impl Sized> #where_clause{
             #task_outer_body
         }
 
diff --git a/embassy-executor/CHANGELOG.md b/embassy-executor/CHANGELOG.md
index 914863a83..7404961f3 100644
--- a/embassy-executor/CHANGELOG.md
+++ b/embassy-executor/CHANGELOG.md
@@ -21,6 +21,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Added support for `-> impl Future<Output = ()>` in `#[task]`
 - Fixed `Send` unsoundness with `-> impl Future` tasks
 - Marked `Spawner::for_current_executor` as `unsafe`
+- `#[task]` now properly marks the generated function as unsafe if the task is marked unsafe
 
 ## 0.7.0 - 2025-01-02
 
diff --git a/embassy-executor/tests/ui.rs b/embassy-executor/tests/ui.rs
index 7757775ee..8b83cd368 100644
--- a/embassy-executor/tests/ui.rs
+++ b/embassy-executor/tests/ui.rs
@@ -32,4 +32,5 @@ fn ui() {
     t.compile_fail("tests/ui/self.rs");
     t.compile_fail("tests/ui/type_error.rs");
     t.compile_fail("tests/ui/where_clause.rs");
+    t.pass("tests/ui/task_safety_attribute.rs");
 }
diff --git a/embassy-executor/tests/ui/task_safety_attribute.rs b/embassy-executor/tests/ui/task_safety_attribute.rs
new file mode 100644
index 000000000..ab5a2f99f
--- /dev/null
+++ b/embassy-executor/tests/ui/task_safety_attribute.rs
@@ -0,0 +1,25 @@
+#![cfg_attr(feature = "nightly", feature(impl_trait_in_assoc_type))]
+#![deny(unused_unsafe)]
+
+use std::mem;
+
+#[embassy_executor::task]
+async fn safe() {}
+
+#[embassy_executor::task]
+async unsafe fn not_safe() {}
+
+#[export_name = "__pender"]
+fn pender(_: *mut ()) {
+    // The test doesn't link if we don't include this.
+    // We never call this anyway.
+}
+
+fn main() {
+    let _forget_me = safe();
+    // SAFETY: not_safe has not safety preconditions
+    let _forget_me2 = unsafe { not_safe() };
+
+    mem::forget(_forget_me);
+    mem::forget(_forget_me2);
+}