use core::cell::Cell;
use core::future::poll_fn;
use core::ops::{Deref, DerefMut};
use core::pin::Pin;
use core::task::{Poll, Context, Waker};

struct Malicious<'a, T> {
    fixed: Box<T>, 
    movable: Box<T>,
    comm: &'a Cell<Box<T>>,
    which: &'a Cell<bool>,
}

impl<'a, T> Deref for Malicious<'a, T> {
    type Target = T;
    fn deref(&self) -> &T {
        if self.which.get() {
            &*self.movable
        } else {
            &*self.fixed
        }
    }
}
impl<'a, T> DerefMut for Malicious<'a, T> {
    fn deref_mut(&mut self) -> &mut T {
        if self.which.get() {
            Cell::from_mut(&mut self.movable).swap(self.comm);
            &mut *self.movable
        } else {
            &mut *self.fixed
        }
    }
}

async fn future_maker(ready: &Cell<Poll<()>>) {
    let s = String::from("Hello, world!");
    let sref = &s;
    let pf = poll_fn(move |_| ready.get());
    pf.await;
    println!("{}", sref);
}

fn main() {
    let ready = Cell::new(Poll::Pending);
    let comm = Cell::new(Box::new(future_maker(&ready)));
    let which = Cell::new(false);
    let m = Malicious {
        fixed: Box::new(future_maker(&ready)),
        movable: Box::new(future_maker(&ready)),
        comm: &comm,
        which: &which
    };

    // Safety: m is a reference to *m.fixed, and *m.fixed is not
    // moved out of until it is dropped; m's Drop, Deref,
    // DerefMut methods do not move out of the reference target
    let mut pin = unsafe { Pin::new_unchecked(m) };

    // Pin::new_unchecked does not require that the reference
    // always points to the same place; this can be exploited
    // to change the reference to point to *m.movable, then
    // move out of *m.movable, causing a use-after-free
    let mut noop_context = Context::from_waker(Waker::noop());
    which.set(true);
    // as_mut() derefs to what was, before this command, comm
    assert_eq!(pin.as_mut().poll(&mut noop_context), Poll::Pending);
    // move the box that was dereffed back into comm
    pin.as_mut();
    // move the future into a different box
    let mut old_box = comm.replace(Box::new(future_maker(&ready)));
    let future = std::mem::replace(&mut *old_box, future_maker(&ready));
    let new_box = Box::new(future);
    comm.set(new_box);
    // poll the future from its new location (UB because we are
    // polling a moved future)
    ready.set(Poll::Ready(()));
    assert_eq!(pin.as_mut().poll(&mut noop_context), Poll::Ready(()));
}