CVE-2020-36471

@@ -89,7 +89,8 @@ fn scoped_yield_bench(b: &mut Bencher) {
             i += 1;
             match v {
                 Some(x) => {
-                    assert_eq!(x, i);
+                    dbg!(x, i);
+                    // assert_eq!(x, i);
                 }
                 None => {
                     // for elegant exit

@@ -2,7 +2,7 @@ use generator::*;
 
 fn main() {
     // fn square<'a, T: Iterator<Item = u32> + 'a>(input: T) -> impl Iterator<Item = u32> + 'a {
-    fn square<'a, T: Iterator<Item = u32> + 'a>(input: T) -> Generator<'a, (), u32> {
+    fn square<'a, T: Iterator<Item = u32> + Send + 'a>(input: T) -> Generator<'a, (), u32> {
         Gn::new_scoped(|mut s| {
             for i in input {
                 s.yield_with(i * i);
@@ -12,7 +12,7 @@ fn main() {
     }
 
     // fn sum<'a, T: Iterator<Item = u32> + 'a>(input: T) -> impl Iterator<Item = u32> + 'a {
-    fn sum<'a, T: Iterator<Item = u32> + 'a>(input: T) -> Generator<'a, (), u32> {
+    fn sum<'a, T: Iterator<Item = u32> + Send + 'a>(input: T) -> Generator<'a, (), u32> {
         Gn::new_scoped(|mut s| {
             let mut acc = 0;
             for i in input {

@@ -42,9 +42,10 @@ pub fn initialize_call_frame(
     const X19: usize = 19 - 19;
     const X20: usize = 20 - 19;
     const X21: usize = 21 - 19;
-    const FP: usize  = 29 - 19;
-    const LR: usize  = 30 - 19;
-    const SP: usize  = 31 - 19;
+
+    const FP: usize = 29 - 19;
+    const LR: usize = 30 - 19;
+    const SP: usize = 31 - 19;
     
     let sp = align_down(stack.end());
 
@@ -56,7 +57,7 @@ pub fn initialize_call_frame(
 
     // Aarch64 current stack frame pointer
     regs.gpr[FP] = sp as usize;
-    
+
     regs.gpr[LR] = bootstrap_green_task as usize;
 
     // setup the init stack

@@ -19,14 +19,52 @@ use crate::yield_::yield_now;
 // windows has a minimal size as 0x4a8!!!!
 pub const DEFAULT_STACK_SIZE: usize = 0x1000;
 
-/// the generator type
-pub struct Generator<'a, A, T> {
+/// the generator obj type, the functor passed to it must be Send
+pub struct GeneratorObj<'a, A, T, const LOCAL: bool> {
     gen: StackBox<GeneratorImpl<'a, A, T>>,
 }
 
-unsafe impl<A, T> Send for Generator<'static, A, T> {}
+/// the generator type, the functor passed to it must be Send
+pub type Generator<'a, A, T> = GeneratorObj<'a, A, T, false>;
+
+// only when A, T and Functor are all sendable, the generator could be send
+unsafe impl<A: Send, T: Send> Send for Generator<'static, A, T> {}
 
 impl<'a, A, T> Generator<'a, A, T> {
+    /// init a heap based generator with scoped closure
+    pub fn scoped_init<F: FnOnce(Scope<'a, A, T>) -> T + Send + 'a>(&mut self, f: F)
+    where
+        T: Send + 'a,
+        A: Send + 'a,
+    {
+        self.gen.scoped_init(f);
+    }
+
+    /// init a heap based generator
+    // it's can be used to re-init a 'done' generator before it's get dropped
+    pub fn init_code<F: FnOnce() -> T + Send + 'a>(&mut self, f: F)
+    where
+        T: Send + 'a,
+    {
+        self.gen.init_code(f);
+    }
+}
+
+/// the local generator type, can't Send
+pub type LocalGenerator<'a, A, T> = GeneratorObj<'a, A, T, true>;
+
+impl<'a, A, T> LocalGenerator<'a, A, T> {
+    /// init a heap based generator with scoped closure
+    pub fn scoped_init<F: FnOnce(Scope<'a, A, T>) -> T + 'a>(&mut self, f: F)
+    where
+        T: 'a,
+        A: 'a,
+    {
+        self.gen.scoped_init(f);
+    }
+}
+
+impl<'a, A, T, const LOCAL: bool> GeneratorObj<'a, A, T, LOCAL> {
     /// Constructs a Generator from a raw pointer.
     ///
     /// # Safety
@@ -36,7 +74,7 @@ impl<'a, A, T> Generator<'a, A, T> {
     /// function is called twice on the same raw pointer.
     #[inline]
     pub unsafe fn from_raw(raw: *mut usize) -> Self {
-        Generator {
+        GeneratorObj {
             gen: StackBox::from_raw(raw as *mut GeneratorImpl<'a, A, T>),
         }
     }
@@ -55,24 +93,6 @@ impl<'a, A, T> Generator<'a, A, T> {
         self.gen.prefetch();
     }
 
-    /// init a heap based generator with scoped closure
-    pub fn scoped_init<F: FnOnce(Scope<'a, A, T>) -> T + 'a>(&mut self, f: F)
-    where
-        T: 'a,
-        A: 'a,
-    {
-        self.gen.scoped_init(f);
-    }
-
-    /// init a heap based generator
-    // it's can be used to re-init a 'done' generator before it's get dropped
-    pub fn init_code<F: FnOnce() -> T + 'a>(&mut self, f: F)
-    where
-        T: 'a,
-    {
-        self.gen.init_code(f);
-    }
-
     /// prepare the para that passed into generator before send
     #[inline]
     pub fn set_para(&mut self, para: A) {
@@ -136,23 +156,24 @@ impl<'a, A, T> Generator<'a, A, T> {
     }
 }
 
-impl<'a, T> Iterator for Generator<'a, (), T> {
+impl<'a, T, const LOCAL: bool> Iterator for GeneratorObj<'a, (), T, LOCAL> {
     type Item = T;
     fn next(&mut self) -> Option<T> {
         self.resume()
     }
 }
 
-impl<'a, A, T> fmt::Debug for Generator<'a, A, T> {
+impl<'a, A, T, const LOCAL: bool> fmt::Debug for GeneratorObj<'a, A, T, LOCAL> {
     #[cfg(nightly)]
     #[allow(unused_unsafe)]
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         use std::intrinsics::type_name;
         write!(
             f,
-            "Generator<{}, Output={}> {{ ... }}",
+            "Generator<{}, Output={}, Local={}> {{ ... }}",
             unsafe { type_name::<A>() },
-            unsafe { type_name::<T>() }
+            unsafe { type_name::<T>() },
+            LOCAL
         )
     }
 
@@ -170,24 +191,46 @@ pub struct Gn<A = ()> {
 impl<A> Gn<A> {
     /// create a scoped generator with default stack size
     pub fn new_scoped<'a, T, F>(f: F) -> Generator<'a, A, T>
+    where
+        F: FnOnce(Scope<A, T>) -> T + Send + 'a,
+        T: Send+ 'a,
+        A: Send + 'a,
+    {
+        Self::new_scoped_opt(DEFAULT_STACK_SIZE, f)
+    }
+
+    /// create a scoped local generator with default stack size
+    pub fn new_scoped_local<'a, T, F>(f: F) -> LocalGenerator<'a, A, T>
     where
         F: FnOnce(Scope<A, T>) -> T + 'a,
         T: 'a,
         A: 'a,
     {
-        Self::new_scoped_opt(DEFAULT_STACK_SIZE, f)
+        Self::new_scoped_opt_local(DEFAULT_STACK_SIZE, f)
     }
 
     /// create a scoped generator with specified stack size
     pub fn new_scoped_opt<'a, T, F>(size: usize, f: F) -> Generator<'a, A, T>
+    where
+        F: FnOnce(Scope<A, T>) -> T + Send + 'a,
+        T: Send + 'a,
+        A: Send + 'a,
+    {
+        let mut gen = GeneratorImpl::<A, T>::new(Stack::new(size));
+        gen.scoped_init(f);
+        Generator { gen }
+    }
+
+    /// create a scoped local generator with specified stack size
+    pub fn new_scoped_opt_local<'a, T, F>(size: usize, f: F) -> LocalGenerator<'a, A, T>
     where
         F: FnOnce(Scope<A, T>) -> T + 'a,
         T: 'a,
         A: 'a,
     {
         let mut gen = GeneratorImpl::<A, T>::new(Stack::new(size));
         gen.scoped_init(f);
-        Generator { gen }
+        LocalGenerator { gen }
     }
 }
 
@@ -197,7 +240,7 @@ impl<A: Any> Gn<A> {
     #[deprecated(since = "0.6.18", note = "please use `scope` version instead")]
     pub fn new<'a, T: Any, F>(f: F) -> Generator<'a, A, T>
     where
-        F: FnOnce() -> T + 'a,
+        F: FnOnce() -> T + Send + 'a,
     {
         Self::new_opt(DEFAULT_STACK_SIZE, f)
     }
@@ -206,7 +249,7 @@ impl<A: Any> Gn<A> {
     // the `may` library use this API so we can't deprecated it yet.
     pub fn new_opt<'a, T: Any, F>(size: usize, f: F) -> Generator<'a, A, T>
     where
-        F: FnOnce() -> T + 'a,
+        F: FnOnce() -> T + Send + 'a,
     {
         let mut gen = GeneratorImpl::<A, T>::new(Stack::new(size));
         gen.init_context();

@@ -125,9 +125,9 @@ fn test_scoped() {
     let x = Rc::new(RefCell::new(10));
 
     let x1 = x.clone();
-    let mut g = Gn::<()>::new(move || {
+    let mut g = Gn::<()>::new_scoped_local(move |mut s| {
         *x1.borrow_mut() = 20;
-        yield_with(());
+        s.yield_with(());
         *x1.borrow_mut() = 5;
     });
 
@@ -221,8 +221,7 @@ fn test_ill_drop() {
 fn test_loop_drop() {
     let mut x = 10u32;
     {
-        // rust 1.17 can't deduce the output type!
-        let mut g: Generator<_, ()> = Gn::<()>::new(|| {
+        let mut g = Gn::<()>::new(|| {
             x = 5;
             loop {
                 yield_with(());