From 46a3f2befd6fd821e5747ce9db112c550bc989f3 Mon Sep 17 00:00:00 2001
From: Emil Fresk <emil.fresk@gmail.com>
Date: Fri, 5 Aug 2022 11:28:02 +0200
Subject: [PATCH] Fix UB in the access of `Priority` for asyc executors

The `Priority` was generated on the stack in the dispatcher
which caused it to be dropped after usage. This is now fixed
by having the `Priority` being a static variable for executors
---
 examples/async-task-multiple-prios.rs | 15 +++++++++------
 macros/src/codegen/dispatchers.rs     | 14 ++++++++++++--
 macros/src/codegen/software_tasks.rs  | 15 +++++++++++----
 src/export.rs                         |  2 +-
 xtask/src/command.rs                  | 11 +++++++++--
 5 files changed, 42 insertions(+), 15 deletions(-)

diff --git a/examples/async-task-multiple-prios.rs b/examples/async-task-multiple-prios.rs
index a90c11bd98..d47e87b1bc 100644
--- a/examples/async-task-multiple-prios.rs
+++ b/examples/async-task-multiple-prios.rs
@@ -16,7 +16,10 @@ mod app {
     use systick_monotonic::*;
 
     #[shared]
-    struct Shared {}
+    struct Shared {
+        a: u32,
+        b: u32,
+    }
 
     #[local]
     struct Local {}
@@ -34,7 +37,7 @@ mod app {
         async_task2::spawn().ok();
 
         (
-            Shared {},
+            Shared { a: 0, b: 0 },
             Local {},
             init::Monotonics(Systick::new(cx.core.SYST, 12_000_000)),
         )
@@ -49,24 +52,24 @@ mod app {
         }
     }
 
-    #[task(priority = 1)]
+    #[task(priority = 1, shared = [a, b])]
     fn normal_task(_cx: normal_task::Context) {
         hprintln!("hello from normal 1").ok();
     }
 
-    #[task(priority = 1)]
+    #[task(priority = 1, shared = [a, b])]
     async fn async_task(_cx: async_task::Context) {
         hprintln!("hello from async 1").ok();
 
         debug::exit(debug::EXIT_SUCCESS);
     }
 
-    #[task(priority = 2)]
+    #[task(priority = 2, shared = [a, b])]
     fn normal_task2(_cx: normal_task2::Context) {
         hprintln!("hello from normal 2").ok();
     }
 
-    #[task(priority = 2)]
+    #[task(priority = 2, shared = [a, b])]
     async fn async_task2(_cx: async_task2::Context) {
         hprintln!("hello from async 2").ok();
     }
diff --git a/macros/src/codegen/dispatchers.rs b/macros/src/codegen/dispatchers.rs
index cc9944025e..62df429370 100644
--- a/macros/src/codegen/dispatchers.rs
+++ b/macros/src/codegen/dispatchers.rs
@@ -10,17 +10,23 @@ pub fn codegen(app: &App, analysis: &Analysis, extra: &Extra) -> Vec<TokenStream
 
     let interrupts = &analysis.interrupts;
 
-    // Generate executor definition in global scope
+    // Generate executor definition and priority in global scope
     for (name, task) in app.software_tasks.iter() {
         if task.is_async {
             let type_name = util::internal_task_ident(name, "F");
             let exec_name = util::internal_task_ident(name, "EXEC");
+            let prio_name = util::internal_task_ident(name, "PRIORITY");
 
             items.push(quote!(
                 type #type_name = impl core::future::Future + 'static;
                 static #exec_name:
                     rtic::RacyCell<rtic::export::executor::AsyncTaskExecutor<#type_name>> =
                         rtic::RacyCell::new(rtic::export::executor::AsyncTaskExecutor::new());
+
+                // The executors priority, this can be any value - we will overwrite it when we
+                // start a task
+                static #prio_name: rtic::RacyCell<rtic::export::Priority> =
+                        unsafe { rtic::RacyCell::new(rtic::export::Priority::new(0)) };
             ));
         }
     }
@@ -92,6 +98,7 @@ pub fn codegen(app: &App, analysis: &Analysis, extra: &Extra) -> Vec<TokenStream
                 let inputs = util::inputs_ident(name);
                 let (_, tupled, pats, _) = util::regroup_inputs(&task.inputs);
                 let exec_name = util::internal_task_ident(name, "EXEC");
+                let prio_name = util::internal_task_ident(name, "PRIORITY");
 
                 if task.is_async {
                     let executor_run_ident = util::executor_run_ident(name);
@@ -108,7 +115,10 @@ pub fn codegen(app: &App, analysis: &Analysis, extra: &Extra) -> Vec<TokenStream
                                     .read();
                                 (&mut *#fq.get_mut()).split().0.enqueue_unchecked(index);
 
-                                let priority = &rtic::export::Priority::new(PRIORITY);
+                                // The async executor needs a static priority
+                                #prio_name.get_mut().write(rtic::export::Priority::new(PRIORITY));
+                                let priority: &'static _ = &*#prio_name.get();
+
                                 (&mut *#exec_name.get_mut()).spawn(#name(#name::Context::new(priority) #(,#pats)*));
                                 #executor_run_ident.store(true, core::sync::atomic::Ordering::Relaxed);
                             } else {
diff --git a/macros/src/codegen/software_tasks.rs b/macros/src/codegen/software_tasks.rs
index 6d08a221a0..71869b6e61 100644
--- a/macros/src/codegen/software_tasks.rs
+++ b/macros/src/codegen/software_tasks.rs
@@ -139,17 +139,24 @@ pub fn codegen(
             let attrs = &task.attrs;
             let cfgs = &task.cfgs;
             let stmts = &task.stmts;
-            let async_marker = if task.is_async {
-                quote!(async)
+            let (async_marker, context_lifetime) = if task.is_async {
+                (
+                    quote!(async),
+                    if shared_needs_lt || local_needs_lt {
+                        quote!(<'static>)
+                    } else {
+                        quote!()
+                    },
+                )
             } else {
-                quote!()
+                (quote!(), quote!())
             };
 
             user_tasks.push(quote!(
                 #(#attrs)*
                 #(#cfgs)*
                 #[allow(non_snake_case)]
-                #async_marker fn #name(#context: #name::Context #(,#inputs)*) {
+                #async_marker fn #name(#context: #name::Context #context_lifetime #(,#inputs)*) {
                     use rtic::Mutex as _;
                     use rtic::mutex::prelude::*;
 
diff --git a/src/export.rs b/src/export.rs
index 1d17d46b02..ec18c9cf41 100644
--- a/src/export.rs
+++ b/src/export.rs
@@ -242,7 +242,7 @@ impl Priority {
     ///
     /// Will overwrite the current Priority
     #[inline(always)]
-    pub unsafe fn new(value: u8) -> Self {
+    pub const unsafe fn new(value: u8) -> Self {
         Priority {
             inner: Cell::new(value),
         }
diff --git a/xtask/src/command.rs b/xtask/src/command.rs
index 2f719bf5c9..528401ae74 100644
--- a/xtask/src/command.rs
+++ b/xtask/src/command.rs
@@ -46,7 +46,14 @@ impl<'a> CargoCommand<'a> {
                 features,
                 mode,
             } => {
-                let mut args = vec![self.name(), "--example", example, "--target", target];
+                let mut args = vec![
+                    "+nightly",
+                    self.name(),
+                    "--example",
+                    example,
+                    "--target",
+                    target,
+                ];
 
                 if let Some(feature_name) = features {
                     args.extend_from_slice(&["--features", feature_name]);
@@ -61,7 +68,7 @@ impl<'a> CargoCommand<'a> {
                 features,
                 mode,
             } => {
-                let mut args = vec![self.name(), "--examples", "--target", target];
+                let mut args = vec!["+nightly", self.name(), "--examples", "--target", target];
 
                 if let Some(feature_name) = features {
                     args.extend_from_slice(&["--features", feature_name]);