feat: generic protocol detector

monolake-services/src/common/detect.rs

@@ -0,0 +1,118 @@
+use std::{future::Future, io, io::Cursor};
+
+use monoio::{
+    buf::IoBufMut,
+    io::{AsyncReadRent, AsyncReadRentExt, PrefixedReadIo},
+};
+use service_async::Service;
+
+/// Detect is a trait for detecting a certain pattern in the input stream.
+///
+/// It accepts an input stream and returns a tuple of the detected pattern and the wrapped input
+/// stream which is usually a `PrefixedReadIo`. The implementation can choose to whether add the
+/// prefix data.
+/// If it fails to detect the pattern, it should represent the error inside the `DetOut`.
+pub trait Detect<IO> {
+    type DetOut;
+    type IOOut;
+
+    fn detect(&self, io: IO) -> impl Future<Output = io::Result<(Self::DetOut, Self::IOOut)>>;
+}
+
+/// DetectService is a service that detects a certain pattern in the input stream and forwards the
+/// detected pattern and the wrapped input stream to the inner service.
+pub struct DetectService<D, S> {
+    pub detector: D,
+    pub inner: S,
+}
+
+#[derive(thiserror::Error, Debug)]
+pub enum DetectError<E> {
+    #[error("service error: {0:?}")]
+    Svc(E),
+    #[error("io error: {0:?}")]
+    Io(std::io::Error),
+}
+
+impl<R, S, D, CX> Service<(R, CX)> for DetectService<D, S>
+where
+    D: Detect<R>,
+    S: Service<(D::DetOut, D::IOOut, CX)>,
+{
+    type Response = S::Response;
+    type Error = DetectError<S::Error>;
+
+    async fn call(&self, (io, cx): (R, CX)) -> Result<Self::Response, Self::Error> {
+        let (det, io) = self.detector.detect(io).await.map_err(DetectError::Io)?;
+        self.inner
+            .call((det, io, cx))
+            .await
+            .map_err(DetectError::Svc)
+    }
+}
+
+/// FixedLengthDetector detects a fixed length of bytes from the input stream.
+pub struct FixedLengthDetector<const N: usize, F>(pub F);
+
+impl<const N: usize, F, IO, DetOut> Detect<IO> for FixedLengthDetector<N, F>
+where
+    F: Fn(&mut [u8]) -> DetOut,
+    IO: AsyncReadRent,
+{
+    type DetOut = DetOut;
+    type IOOut = PrefixedReadIo<IO, Cursor<Vec<u8>>>;
+
+    async fn detect(&self, mut io: IO) -> io::Result<(Self::DetOut, Self::IOOut)> {
+        let buf = Vec::with_capacity(N).slice_mut(..N);
+        let (r, buf) = io.read_exact(buf).await;
+        r?;
+
+        let mut buf = buf.into_inner();
+        let r = (self.0)(&mut buf);
+        Ok((r, PrefixedReadIo::new(io, Cursor::new(buf))))
+    }
+}
+
+/// PrefixDetector detects a certain prefix from the input stream.
+///
+/// If the prefix matches, it returns true and the wrapped input stream with the prefix data.
+/// Otherwise, it returns false and the input stream with the prefix data(the prefix maybe less than
+/// the static str's length).
+pub struct PrefixDetector(pub &'static [u8]);
+
+impl<IO> Detect<IO> for PrefixDetector
+where
+    IO: AsyncReadRent,
+{
+    type DetOut = bool;
+    type IOOut = PrefixedReadIo<IO, Cursor<Vec<u8>>>;
+
+    async fn detect(&self, mut io: IO) -> io::Result<(Self::DetOut, Self::IOOut)> {
+        let l = self.0.len();
+        let mut written = 0;
+        let mut buf: Vec<u8> = Vec::with_capacity(l);
+        let mut eq = true;
+        loop {
+            // # Safety
+            // The buf must have enough capacity to write the data.
+            let buf_slice = unsafe { buf.slice_mut_unchecked(written..l) };
+            let (result, buf_slice) = io.read(buf_slice).await;
+            buf = buf_slice.into_inner();
+            match result? {
+                0 => {
+                    break;
+                }
+                n => {
+                    let curr = written;
+                    written += n;
+                    if self.0[curr..written] != buf[curr..written] {
+                        eq = false;
+                        break;
+                    }
+                }
+            }
+        }
+        let io = PrefixedReadIo::new(io, Cursor::new(buf));
+        Ok((eq && written == l, io))
+    }
+}

monolake-services/src/common/mod.rs

@@ -2,6 +2,7 @@
 mod cancel;
 mod context;
 mod delay;
+mod detect;
 mod erase;
 mod map;
 mod panic;
@@ -10,6 +11,7 @@ mod timeout;
 pub use cancel::{linked_list, Canceller, CancellerDropper, Waiter};
 pub use context::ContextService;
 pub use delay::{Delay, DelayService};
+pub use detect::{Detect, DetectService, FixedLengthDetector, PrefixDetector};
 pub use erase::EraseResp;
 pub use map::{FnSvc, Map, MapErr};
 pub use panic::{CatchPanicError, CatchPanicService};

monolake-services/src/http/core.rs

@@ -16,7 +16,7 @@
 //!
 //! - Support for HTTP/1, HTTP/1.1, and HTTP/2 protocols
 //! - Composable design allowing a stack of `HttpHandler` implementations
-//! - Automatic protocol detection when combined with `HttpVersionDetect`
+//! - Automatic protocol detection when combined with `H2Detect`
 //! - Efficient handling of concurrent requests using asynchronous I/O
 //! - Configurable timeout settings for different stages of request processing
 //! - Integration with `service_async` for easy composition in service stacks
@@ -25,17 +25,17 @@
 //! # Usage
 //!
 //! `HttpCoreService` is typically used as part of a larger service stack, often in combination
-//! with `HttpVersionDetect` for automatic protocol detection. Here's a basic example:
+//! with `H2Detect` for automatic protocol detection. Here's a basic example:
 //!
 //! ```ignore
 //! use service_async::{layer::FactoryLayer, stack::FactoryStack};
 //!
-//! use crate::http::{HttpCoreService, HttpVersionDetect};
+//! use crate::http::{HttpCoreService, H2Detect};
 //!
 //! let config = Config { /* ... */ };
 //! let stack = FactoryStack::new(config)
 //!     .push(HttpCoreService::layer())
-//!     .push(HttpVersionDetect::layer())
+//!     .push(H2Detect::layer())
 //!     // ... other handlers implementing HttpHandler ...
 //!     ;
 //!
@@ -52,7 +52,7 @@
 //!
 //! # Automatic Protocol Detection
 //!
-//! When used in conjunction with `HttpVersionDetect`, `HttpCoreService` can automatically
+//! When used in conjunction with `H2Detect`, `HttpCoreService` can automatically
 //! detect whether an incoming connection is using HTTP/1, HTTP/1.1, or HTTP/2, and handle
 //! it appropriately. This allows for seamless support of multiple HTTP versions without
 //! the need for separate server configurations.

monolake-services/src/http/detect.rs

@@ -6,8 +6,8 @@
 //!
 //! # Key Components
 //!
-//! - [`HttpVersionDetect`]: The main service component responsible for HTTP version detection.
-//! - [`HttpVersionDetectError`]: Error type for version detection operations.
+//! - [`H2Detect`]: The main service component responsible for HTTP version detection.
+//! - [`H2DetectError`]: Error type for version detection operations.
 //!
 //! # Features
 //!
@@ -27,7 +27,7 @@
 //! let config = Config { /* ... */ };
 //! let stack = FactoryStack::new(config)
 //!     .push(HttpCoreService::layer())
-//!     .push(HttpVersionDetect::layer())
+//!     .push(H2Detect::layer())
 //!     // ... other layers ...
 //!     ;
 //!
@@ -39,122 +39,66 @@
 //!
 //! - Uses efficient buffering to minimize I/O operations during version detection
 //! - Implements zero-copy techniques where possible to reduce memory overhead
-use std::io::Cursor;
 
-use monoio::{
-    buf::IoBufMut,
-    io::{AsyncReadRent, AsyncWriteRent, PrefixedReadIo},
-};
-use monolake_core::http::HttpAccept;
 use service_async::{
     layer::{layer_fn, FactoryLayer},
-    AsyncMakeService, MakeService, Service,
+    AsyncMakeService, MakeService,
 };
 
-use crate::tcp::Accept;
+use crate::common::{DetectService, PrefixDetector};
 
 const PREFACE: &[u8; 24] = b"PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n";
 
 /// Service for detecting HTTP version and routing connections accordingly.
 ///
-/// `HttpVersionDetect` examines the initial bytes of an incoming connection to
+/// `H2Detect` examines the initial bytes of an incoming connection to
 /// determine whether it's an HTTP/2 connection (by checking for the HTTP/2 preface)
 /// or an HTTP/1.x connection. It then forwards the connection to the inner service
 /// with appropriate version information.
 /// For implementation details and example usage, see the
 /// [module level documentation](crate::http::detect).
 #[derive(Clone)]
-pub struct HttpVersionDetect<T> {
+pub struct H2Detect<T> {
     inner: T,
 }
 
 #[derive(thiserror::Error, Debug)]
-pub enum HttpVersionDetectError<E> {
+pub enum H2DetectError<E> {
     #[error("inner error: {0:?}")]
     Inner(E),
     #[error("io error: {0:?}")]
     Io(std::io::Error),
 }
 
-impl<F: MakeService> MakeService for HttpVersionDetect<F> {
-    type Service = HttpVersionDetect<F::Service>;
+impl<F: MakeService> MakeService for H2Detect<F> {
+    type Service = DetectService<PrefixDetector, F::Service>;
     type Error = F::Error;
 
     fn make_via_ref(&self, old: Option<&Self::Service>) -> Result<Self::Service, Self::Error> {
-        Ok(HttpVersionDetect {
+        Ok(DetectService {
             inner: self.inner.make_via_ref(old.map(|o| &o.inner))?,
+            detector: PrefixDetector(PREFACE),
         })
     }
 }
 
-impl<F: AsyncMakeService> AsyncMakeService for HttpVersionDetect<F> {
-    type Service = HttpVersionDetect<F::Service>;
+impl<F: AsyncMakeService> AsyncMakeService for H2Detect<F> {
+    type Service = DetectService<PrefixDetector, F::Service>;
     type Error = F::Error;
 
     async fn make_via_ref(
         &self,
         old: Option<&Self::Service>,
     ) -> Result<Self::Service, Self::Error> {
-        Ok(HttpVersionDetect {
+        Ok(DetectService {
             inner: self.inner.make_via_ref(old.map(|o| &o.inner)).await?,
+            detector: PrefixDetector(PREFACE),
         })
     }
 }
 
-impl<F> HttpVersionDetect<F> {
+impl<F> H2Detect<F> {
     pub fn layer<C>() -> impl FactoryLayer<C, F, Factory = Self> {
-        layer_fn(|_: &C, inner| HttpVersionDetect { inner })
-    }
-}
-
-impl<T, Stream, CX> Service<Accept<Stream, CX>> for HttpVersionDetect<T>
-where
-    Stream: AsyncReadRent + AsyncWriteRent,
-    T: Service<HttpAccept<PrefixedReadIo<Stream, Cursor<Vec<u8>>>, CX>>,
-{
-    type Response = T::Response;
-    type Error = HttpVersionDetectError<T::Error>;
-
-    async fn call(
-        &self,
-        incoming_stream: Accept<Stream, CX>,
-    ) -> Result<Self::Response, Self::Error> {
-        let (mut stream, addr) = incoming_stream;
-        let mut buf = vec![0; PREFACE.len()];
-        let mut pos = 0;
-        let mut h2_detect = false;
-
-        loop {
-            let buf_slice = unsafe { buf.slice_mut_unchecked(pos..PREFACE.len()) };
-            let (result, buf_slice) = stream.read(buf_slice).await;
-            buf = buf_slice.into_inner();
-            match result {
-                Ok(0) => {
-                    break;
-                }
-                Ok(n) => {
-                    if PREFACE[pos..pos + n] != buf[pos..pos + n] {
-                        break;
-                    }
-                    pos += n;
-                }
-                Err(e) => {
-                    return Err(HttpVersionDetectError::Io(e));
-                }
-            }
-
-            if pos == PREFACE.len() {
-                h2_detect = true;
-                break;
-            }
-        }
-
-        let preface_buf = std::io::Cursor::new(buf);
-        let rewind_io = monoio::io::PrefixedReadIo::new(stream, preface_buf);
-
-        self.inner
-            .call((h2_detect, rewind_io, addr))
-            .await
-            .map_err(HttpVersionDetectError::Inner)
+        layer_fn(|_: &C, inner| H2Detect { inner })
     }
 }

monolake-services/src/http/handlers/connection_persistence.rs

@@ -25,7 +25,7 @@
 //!     common::ContextService,
 //!     http::{
 //!         core::HttpCoreService,
-//!         detect::HttpVersionDetect,
+//!         detect::H2Detect,
 //!         handlers::{
 //!             route::RouteConfig, ConnectionReuseHandler, ContentHandler, RewriteAndRouteHandler,
 //!             UpstreamHandler,
@@ -60,7 +60,7 @@
 //!     .push(RewriteAndRouteHandler::layer())
 //!     .push(ConnectionReuseHandler::layer())
 //!     .push(HttpCoreService::layer())
-//!     .push(HttpVersionDetect::layer());
+//!     .push(H2Detect::layer());
 //!
 //! // Use the service to handle HTTP requests
 //! ```

monolake-services/src/http/handlers/content_handler.rs

@@ -25,7 +25,7 @@
 //!     common::ContextService,
 //!     http::{
 //!         core::HttpCoreService,
-//!         detect::HttpVersionDetect,
+//!         detect::H2Detect,
 //!         handlers::{
 //!             route::RouteConfig, ConnectionReuseHandler, ContentHandler, RewriteAndRouteHandler,
 //!             UpstreamHandler,
@@ -60,7 +60,7 @@
 //!     .push(RewriteAndRouteHandler::layer())
 //!     .push(ConnectionReuseHandler::layer())
 //!     .push(HttpCoreService::layer())
-//!     .push(HttpVersionDetect::layer());
+//!     .push(H2Detect::layer());
 //!
 //! // Use the service to handle HTTP requests
 //! ```

monolake-services/src/http/handlers/mod.rs

@@ -61,7 +61,7 @@
 //!     common::ContextService,
 //!     http::{
 //!         core::HttpCoreService,
-//!         detect::HttpVersionDetect,
+//!         detect::H2Detect,
 //!         handlers::{
 //!             route::RouteConfig, ConnectionReuseHandler, ContentHandler, RewriteAndRouteHandler,
 //!             UpstreamHandler,
@@ -96,7 +96,7 @@
 //!     .push(RewriteAndRouteHandler::layer())
 //!     .push(ConnectionReuseHandler::layer())
 //!     .push(HttpCoreService::layer())
-//!     .push(HttpVersionDetect::layer());
+//!     .push(H2Detect::layer());
 //!
 //! // Use the service to handle HTTP requests
 //! ```