fix: fix potential UB caused by unsafe and future movement (#154)

* fix: fix potential UB caused by unsafe and future movement

* fix: fix clippy

Author: ihciah

monolake-services/src/common/cancel/mod.rs

@@ -86,7 +86,7 @@ impl Waiter {
     pub fn cancelled(&self) -> bool {
         self.handler
             .upgrade()
-            .map_or(true, |handler| unsafe { &*handler.get() }.cancelled)
+            .is_none_or(|handler| unsafe { &*handler.get() }.cancelled)
     }
 }
 

monolake-services/src/http/core.rs

@@ -63,7 +63,7 @@
 //! - Implements connection keep-alive for HTTP/1.1 to reduce connection overhead
 //! - Supports HTTP/2 multiplexing for efficient handling of concurrent requests
 //! - Automatic protocol detection allows for optimized handling based on the client's capabilities
-use std::{convert::Infallible, fmt::Debug, pin::Pin, time::Duration};
+use std::{convert::Infallible, fmt::Debug, time::Duration};
 
 use bytes::Bytes;
 use certain_map::{Attach, Fork};
@@ -92,7 +92,7 @@ use service_async::{
 };
 use tracing::{error, info, warn};
 
-use super::{generate_response, util::AccompanyPair};
+use super::{generate_response, util::AccompanyPairBase};
 
 /// Core HTTP service handler supporting both HTTP/1.1 and HTTP/2 protocols.
 ///
@@ -172,51 +172,44 @@ impl<H> HttpCoreService<H> {
             // fork ctx
             let (mut store, state) = ctx.fork();
             let forked_ctx = unsafe { state.attach(&mut store) };
+            let mut fut_base = std::pin::pin!(AccompanyPairBase::new(decoder.fill_payload()));
 
             // handle request and reply response
             // 1. do these things simultaneously: read body and send + handle request
-            let mut acc_fut = AccompanyPair::new(
-                self.handler_chain.handle(req, forked_ctx),
-                decoder.fill_payload(),
-            );
-            let res = unsafe { Pin::new_unchecked(&mut acc_fut) }.await;
-            match res {
+            let s1 = fut_base
+                .as_mut()
+                .stage1(self.handler_chain.handle(req, forked_ctx));
+            match s1.await {
                 Ok((resp, should_cont)) => {
                     // 2. do these things simultaneously: read body and send + handle response
-                    let mut f = acc_fut.replace(encoder.send_and_flush(resp));
+                    let s2 = fut_base.as_mut().stage2(encoder.send_and_flush(resp));
                     match self.http_timeout.read_body_timeout {
                         None => {
-                            if let Err(e) = unsafe { Pin::new_unchecked(&mut f) }.await {
+                            if let Err(e) = s2.await {
                                 warn!("error when encode and write response: {e}");
                                 break;
                             }
                         }
-                        Some(body_timeout) => {
-                            match monoio::time::timeout(body_timeout, unsafe {
-                                Pin::new_unchecked(&mut f)
-                            })
-                            .await
-                            {
-                                Err(_) => {
-                                    info!(
-                                        "Connection {:?} write timed out",
-                                        ParamRef::<PeerAddr>::param_ref(&ctx),
-                                    );
-                                    break;
-                                }
-                                Ok(Err(e)) => {
-                                    warn!("error when encode and write response: {e}");
-                                    break;
-                                }
-                                _ => (),
+                        Some(body_timeout) => match monoio::time::timeout(body_timeout, s2).await {
+                            Err(_) => {
+                                info!(
+                                    "Connection {:?} write timed out",
+                                    ParamRef::<PeerAddr>::param_ref(&ctx),
+                                );
+                                break;
                             }
-                        }
+                            Ok(Err(e)) => {
+                                warn!("error when encode and write response: {e}");
+                                break;
+                            }
+                            _ => (),
+                        },
                     }
 
                     if !should_cont {
                         break;
                     }
-                    if let Err(e) = f.into_accompany().await {
+                    if let Err(e) = fut_base.as_mut().stage3().await {
                         warn!("error when decode request body: {e}");
                         break;
                     }

monolake-services/src/http/util.rs

@@ -1,90 +1,152 @@
-use std::{future::Future, task::Poll};
+use std::{
+    future::Future,
+    mem::{ManuallyDrop, MaybeUninit},
+    ops::DerefMut,
+    pin::Pin,
+    task::Poll,
+};
 
 use http::{HeaderValue, Request, Response, StatusCode};
 use monoio_http::common::body::FixedBody;
 use monolake_core::http::{HttpError, HttpHandler, ResponseWithContinue};
 use service_async::Service;
 
-pin_project_lite::pin_project! {
-    /// AccompanyPair for http decoder and processor.
-    /// We have to fill payload when process request
-    /// since inner logic may read chunked body; also
-    /// fill payload when process response since we
-    /// may use the request body stream in response
-    /// body stream.
-    pub(crate) struct AccompanyPair<FMAIN, FACC, T> {
-        #[pin]
-        main: FMAIN,
-        #[pin]
-        accompany: FACC,
-        accompany_slot: Option<T>
+union Union<A, B> {
+    a: ManuallyDrop<A>,
+    b: ManuallyDrop<B>,
+}
+
+/// AccompanyPair for http decoder and processor.
+/// We have to fill payload when process request
+/// since inner logic may read chunked body; also
+/// fill payload when process response since we
+/// may use the request body stream in response
+/// body stream.
+pub(crate) struct AccompanyPairBase<F1, F2, FACC, T> {
+    main: MaybeUninit<Union<F1, F2>>,
+    accompany: FACC,
+    accompany_slot: Option<T>,
+}
+
+impl<F1, F2, FACC, T> AccompanyPairBase<F1, F2, FACC, T> {
+    pub(crate) const fn new(accompany: FACC) -> Self {
+        Self {
+            main: MaybeUninit::uninit(),
+            accompany,
+            accompany_slot: None,
+        }
+    }
+
+    pub(crate) fn stage1(
+        mut self: Pin<&mut Self>,
+        future1: F1,
+    ) -> AccompanyPairS1<F1, F2, FACC, T> {
+        unsafe {
+            self.as_mut().get_unchecked_mut().main.assume_init_mut().a = ManuallyDrop::new(future1);
+        }
+        AccompanyPairS1(self)
+    }
+
+    pub(crate) fn stage2(
+        mut self: Pin<&mut Self>,
+        future2: F2,
+    ) -> AccompanyPairS2<F1, F2, FACC, T> {
+        unsafe {
+            self.as_mut().get_unchecked_mut().main.assume_init_mut().b = ManuallyDrop::new(future2);
+        }
+        AccompanyPairS2(self)
+    }
+
+    pub(crate) const fn stage3(self: Pin<&mut Self>) -> AccompanyPairS3<F1, F2, FACC, T> {
+        AccompanyPairS3(self)
     }
 }
 
-pin_project_lite::pin_project! {
-    /// Accompany for http decoder and processor.
-    pub(crate) struct Accompany<FACC, T> {
-        #[pin]
-        accompany: FACC,
-        accompany_slot: Option<T>
+#[repr(transparent)]
+pub(crate) struct AccompanyPairS1<'a, F1, F2, FACC, T>(
+    Pin<&'a mut AccompanyPairBase<F1, F2, FACC, T>>,
+);
+
+impl<F1, F2, FACC, T> Drop for AccompanyPairS1<'_, F1, F2, FACC, T> {
+    fn drop(&mut self) {
+        unsafe {
+            ManuallyDrop::drop(&mut self.0.as_mut().get_unchecked_mut().main.assume_init_mut().a);
+        }
     }
 }
 
-impl<FMAIN, FACC, T> Future for AccompanyPair<FMAIN, FACC, T>
+impl<F1, F2, FACC, T> Future for AccompanyPairS1<'_, F1, F2, FACC, T>
 where
-    FMAIN: Future,
+    F1: Future,
     FACC: Future<Output = T>,
 {
-    type Output = FMAIN::Output;
+    type Output = F1::Output;
 
-    fn poll(self: std::pin::Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> Poll<Self::Output> {
-        let this = self.project();
+    fn poll(
+        mut self: std::pin::Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> Poll<Self::Output> {
+        let this = unsafe { self.0.as_mut().get_unchecked_mut() };
         if this.accompany_slot.is_none()
-            && let Poll::Ready(t) = this.accompany.poll(cx)
+            && let Poll::Ready(t) = unsafe { Pin::new_unchecked(&mut this.accompany) }.poll(cx)
         {
-            *this.accompany_slot = Some(t);
+            this.accompany_slot = Some(t);
         }
-        this.main.poll(cx)
+        unsafe { Pin::new_unchecked(this.main.assume_init_mut().a.deref_mut()).poll(cx) }
     }
 }
 
-impl<FACC, T> Future for Accompany<FACC, T>
+#[repr(transparent)]
+pub(crate) struct AccompanyPairS2<'a, F1, F2, FACC, T>(
+    Pin<&'a mut AccompanyPairBase<F1, F2, FACC, T>>,
+);
+
+impl<F1, F2, FACC, T> Drop for AccompanyPairS2<'_, F1, F2, FACC, T> {
+    fn drop(&mut self) {
+        unsafe {
+            ManuallyDrop::drop(&mut self.0.as_mut().get_unchecked_mut().main.assume_init_mut().b);
+        }
+    }
+}
+
+impl<F1, F2, FACC, T> Future for AccompanyPairS2<'_, F1, F2, FACC, T>
 where
+    F2: Future,
     FACC: Future<Output = T>,
 {
-    type Output = T;
+    type Output = F2::Output;
 
-    fn poll(self: std::pin::Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> Poll<Self::Output> {
-        let this = self.project();
-        if let Some(t) = this.accompany_slot.take() {
-            return Poll::Ready(t);
+    fn poll(
+        mut self: std::pin::Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> Poll<Self::Output> {
+        let this = unsafe { self.0.as_mut().get_unchecked_mut() };
+        if this.accompany_slot.is_none()
+            && let Poll::Ready(t) = unsafe { Pin::new_unchecked(&mut this.accompany) }.poll(cx)
+        {
+            this.accompany_slot = Some(t);
         }
-        this.accompany.poll(cx)
+        unsafe { Pin::new_unchecked(this.main.assume_init_mut().b.deref_mut()).poll(cx) }
     }
 }
 
-impl<FMAIN, FACC, T> AccompanyPair<FMAIN, FACC, T> {
-    pub(crate) fn new(main: FMAIN, accompany: FACC) -> Self {
-        Self {
-            main,
-            accompany,
-            accompany_slot: None,
-        }
-    }
+#[repr(transparent)]
+pub(crate) struct AccompanyPairS3<'a, F1, F2, FACC, T>(
+    Pin<&'a mut AccompanyPairBase<F1, F2, FACC, T>>,
+);
 
-    pub(crate) fn replace<FMAIN2>(self, main: FMAIN2) -> AccompanyPair<FMAIN2, FACC, T> {
-        AccompanyPair {
-            main,
-            accompany: self.accompany,
-            accompany_slot: self.accompany_slot,
-        }
-    }
+impl<F1, F2, FACC: Future<Output = T>, T> Future for AccompanyPairS3<'_, F1, F2, FACC, T> {
+    type Output = FACC::Output;
 
-    pub(crate) fn into_accompany(self) -> Accompany<FACC, T> {
-        Accompany {
-            accompany: self.accompany,
-            accompany_slot: self.accompany_slot,
+    fn poll(
+        mut self: std::pin::Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> Poll<Self::Output> {
+        let this = unsafe { self.0.as_mut().get_unchecked_mut() };
+        if let Some(t) = this.accompany_slot.take() {
+            return Poll::Ready(t);
         }
+        unsafe { Pin::new_unchecked(&mut this.accompany) }.poll(cx)
     }
 }