itsi-server 0.2.15 → 0.2.17

data/ext/itsi_server/src/server/binds/listener.rs

@@ -9,7 +9,7 @@ use super::bind_protocol::BindProtocol;
 use super::tls::ItsiTlsAcceptor;
 use itsi_error::{ItsiError, Result};
 use itsi_tracing::info;
-use socket2::{Domain, Protocol, Socket, Type};
+use socket2::{Domain, Protocol, SockRef, Socket, Type};
 use std::fmt::Display;
 use std::net::{IpAddr, SocketAddr, TcpListener};
 use std::os::fd::{AsRawFd, FromRawFd, RawFd};
@@ -111,7 +111,7 @@ impl TokioListener {
         Self::to_tokio_io(Stream::TcpStream(tcp_stream), None).await
     }
 
-    pub async fn spawn_state_task(&self, mut shutdown_receiver: Receiver<RunningPhase>) {
+    pub async fn spawn_acme_event_task(&self, mut shutdown_receiver: Receiver<RunningPhase>) {
         if let TokioListener::TcpTls(
             _,
             ItsiTlsAcceptor::Automatic(_acme_acceptor, state, _server_config),
@@ -274,15 +274,53 @@ impl Display for Listener {
 }
 
 impl Listener {
-    pub fn into_tokio_listener(self) -> TokioListener {
+    pub fn rebind_listener(listener: TcpListener) -> TcpListener {
+        let sock = SockRef::from(&listener);
+        let (reuse_address, reuse_port) = (
+            sock.reuse_address().unwrap_or(true),
+            sock.reuse_port().unwrap_or(true),
+        );
+
+        if !reuse_address || !reuse_port {
+            return listener;
+        }
+
+        let (ip, port) = sock
+            .local_addr()
+            .unwrap()
+            .as_socket()
+            .map(|addr| (addr.ip(), addr.port()))
+            .unwrap();
+
+        let socket_opts = SocketOpts {
+            reuse_address: sock.reuse_address().unwrap_or(true), // default: true
+            reuse_port: sock.reuse_port().unwrap_or(false),      // default: false
+            nodelay: sock.nodelay().unwrap_or(false),            // default: false
+            recv_buffer_size: sock.recv_buffer_size().unwrap_or(0),
+            send_buffer_size: sock.send_buffer_size().unwrap_or(0),
+            listen_backlog: 1024, // cannot query – pick sane default
+        };
+
+        connect_tcp_socket(ip, port, &socket_opts).unwrap()
+    }
+
+    pub fn into_tokio_listener(self, no_rebind: bool) -> TokioListener {
         match self {
-            Listener::Tcp(listener) => {
+            Listener::Tcp(mut listener) => {
+                if cfg!(target_os = "linux") && !no_rebind {
+                    listener = Listener::rebind_listener(listener);
+                }
                 TokioListener::Tcp(TokioTcpListener::from_std(listener).unwrap())
             }
-            Listener::TcpTls((listener, acceptor)) => TokioListener::TcpTls(
-                TokioTcpListener::from_std(listener).unwrap(),
-                acceptor.clone(),
-            ),
+            Listener::TcpTls((mut listener, acceptor)) => {
+                if cfg!(target_os = "linux") && !no_rebind {
+                    listener = Listener::rebind_listener(listener);
+                }
+                TokioListener::TcpTls(
+                    TokioTcpListener::from_std(listener).unwrap(),
+                    acceptor.clone(),
+                )
+            }
             Listener::Unix(listener) => {
                 TokioListener::Unix(TokioUnixListener::from_std(listener).unwrap())
             }
@@ -414,6 +452,9 @@ fn connect_tcp_socket(addr: IpAddr, port: u16, socket_opts: &SocketOpts) -> Resu
     socket.set_reuse_port(socket_opts.reuse_port).ok();
     socket.set_nonblocking(true).ok();
     socket.set_nodelay(socket_opts.nodelay).ok();
+    socket
+        .set_send_buffer_size(socket_opts.send_buffer_size)
+        .ok();
     socket
         .set_recv_buffer_size(socket_opts.recv_buffer_size)
         .ok();

data/ext/itsi_server/src/server/frame_stream.rs

@@ -0,0 +1,142 @@
+use bytes::{Bytes, BytesMut};
+use futures::Stream;
+use std::convert::Infallible;
+use std::future::Future;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+use std::time::Duration;
+use tokio::sync::mpsc::Receiver;
+use tokio::sync::watch;
+use tokio::time::{sleep, Sleep};
+
+use super::serve_strategy::single_mode::RunningPhase;
+
+#[derive(Debug)]
+pub struct FrameStream {
+    receiver: Receiver<Bytes>,
+    shutdown_rx: watch::Receiver<RunningPhase>,
+    drained: bool,
+}
+
+impl FrameStream {
+    pub fn new(receiver: Receiver<Bytes>, shutdown_rx: watch::Receiver<RunningPhase>) -> Self {
+        Self {
+            receiver,
+            shutdown_rx,
+            drained: false,
+        }
+    }
+}
+
+impl Stream for FrameStream {
+    type Item = Result<Bytes, Infallible>;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        let this = self.get_mut();
+
+        if this.drained {
+            return Poll::Ready(None);
+        }
+
+        match Pin::new(&mut this.receiver).poll_recv(cx) {
+            Poll::Ready(Some(bytes)) => Poll::Ready(Some(Ok(bytes))),
+            Poll::Ready(None) => {
+                this.drained = true;
+                Poll::Ready(None)
+            }
+            Poll::Pending => {
+                if this.shutdown_rx.has_changed().unwrap_or(false)
+                    && *this.shutdown_rx.borrow() == RunningPhase::ShutdownPending
+                {
+                    while let Ok(bytes) = this.receiver.try_recv() {
+                        return Poll::Ready(Some(Ok(bytes)));
+                    }
+                    this.drained = true;
+                    return Poll::Ready(None);
+                }
+
+                Poll::Pending
+            }
+        }
+    }
+}
+
+/// BufferedStream wraps a stream of Bytes and coalesces chunks into a larger buffer,
+/// flushing either after `max_flush_bytes` is reached or `max_flush_interval` elapses.
+pub struct BufferedStream<S> {
+    inner: S,
+    buffer: BytesMut,
+    max_flush_bytes: usize,
+    max_flush_interval: Duration,
+    flush_deadline: Option<Pin<Box<Sleep>>>,
+}
+
+impl<S> BufferedStream<S> {
+    pub fn new(inner: S, max_flush_bytes: usize, max_flush_interval: Duration) -> Self {
+        Self {
+            inner,
+            buffer: BytesMut::with_capacity(max_flush_bytes),
+            max_flush_bytes,
+            max_flush_interval,
+            flush_deadline: None,
+        }
+    }
+}
+
+impl<S> Stream for BufferedStream<S>
+where
+    S: Stream<Item = Result<Bytes, Infallible>> + Unpin,
+{
+    type Item = Result<Bytes, Infallible>;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        let this = self.get_mut();
+
+        loop {
+            // Flush on timer if needed
+            if let Some(deadline) = &mut this.flush_deadline {
+                if Pin::new(deadline).poll(cx).is_ready() && !this.buffer.is_empty() {
+                    let flushed = this.buffer.split().freeze();
+                    this.flush_deadline = None;
+                    return Poll::Ready(Some(Ok(flushed)));
+                }
+            }
+
+            match Pin::new(&mut this.inner).poll_next(cx) {
+                Poll::Ready(Some(Ok(bytes))) => {
+                    this.buffer.extend_from_slice(&bytes);
+
+                    if bytes.is_empty() || this.buffer.len() >= this.max_flush_bytes {
+                        let flushed = this.buffer.split().freeze();
+                        this.flush_deadline = None;
+                        return Poll::Ready(Some(Ok(flushed)));
+                    }
+
+                    if this.flush_deadline.is_none() {
+                        this.flush_deadline = Some(Box::pin(sleep(this.max_flush_interval)));
+                    }
+                }
+                Poll::Ready(None) => {
+                    if this.buffer.is_empty() {
+                        return Poll::Ready(None);
+                    } else {
+                        let flushed = this.buffer.split().freeze();
+                        this.flush_deadline = None;
+                        return Poll::Ready(Some(Ok(flushed)));
+                    }
+                }
+                Poll::Pending => {
+                    if let Some(deadline) = &mut this.flush_deadline {
+                        let deadline = deadline.as_mut();
+                        if deadline.poll(cx).is_ready() && !this.buffer.is_empty() {
+                            let flushed = this.buffer.split().freeze();
+                            this.flush_deadline = None;
+                            return Poll::Ready(Some(Ok(flushed)));
+                        }
+                    }
+                    return Poll::Pending;
+                }
+            }
+        }
+    }
+}

data/ext/itsi_server/src/server/http_message_types.rs

@@ -1,13 +1,146 @@
-use std::convert::Infallible;
-
 use bytes::Bytes;
-use http::{Request, Response};
-use http_body_util::combinators::BoxBody;
-use hyper::body::Incoming;
+use core::fmt;
+use futures::Stream;
+use futures_util::TryStreamExt;
+use http::Request;
+use http_body_util::{combinators::WithTrailers, BodyExt, Either, Empty, Full, StreamBody};
+use hyper::body::{Body, Frame, Incoming, SizeHint};
+use std::{
+    convert::Infallible,
+    pin::Pin,
+    task::{Context, Poll},
+};
 
 use super::size_limited_incoming::SizeLimitedIncoming;
 
-pub type HttpResponse = Response<BoxBody<Bytes, Infallible>>;
+type Inner = Either<Full<Bytes>, Empty<Bytes>>;
+
+type BoxStream =
+    Pin<Box<dyn Stream<Item = Result<Frame<Bytes>, Infallible>> + Send + Sync + 'static>>;
+
+pub struct PlainBody(Either<StreamBody<BoxStream>, Inner>);
+
+impl fmt::Debug for PlainBody {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match &self.0 {
+            Either::Left(_) => f.write_str("PlainBody::Stream(..)"),
+            Either::Right(inner) => match inner {
+                Either::Left(full) => f.debug_tuple("PlainBody::Full").field(full).finish(),
+                Either::Right(_) => f.write_str("PlainBody::Empty"),
+            },
+        }
+    }
+}
+type DynErr = Box<dyn std::error::Error + Send + Sync>;
+
+impl Body for PlainBody {
+    type Data = Bytes;
+    type Error = DynErr;
+
+    fn poll_frame(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
+        unsafe { self.map_unchecked_mut(|s| &mut s.0) }.poll_frame(cx)
+    }
+
+    fn size_hint(&self) -> SizeHint {
+        self.0.size_hint()
+    }
+}
+
+impl PlainBody {
+    fn stream<S>(s: S) -> Self
+    where
+        S: Stream<Item = Result<Bytes, Infallible>> + Send + Sync + 'static,
+    {
+        let boxed: BoxStream = Box::pin(s.map_ok(Frame::data));
+        Self(Either::Left(StreamBody::new(boxed)))
+    }
+
+    fn full(bytes: Bytes) -> Self {
+        Self(Either::Right(Either::Left(Full::new(bytes))))
+    }
+
+    fn empty() -> Self {
+        Self(Either::Right(Either::Right(Empty::new())))
+    }
+}
+
+type BoxTrailers = Pin<
+    Box<dyn std::future::Future<Output = Option<Result<http::HeaderMap, DynErr>>> + Send + Sync>,
+>;
+
+pub enum HttpBody {
+    Plain(PlainBody),
+    WithT(WithTrailers<PlainBody, BoxTrailers>),
+}
+
+impl fmt::Debug for HttpBody {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            HttpBody::Plain(b) => f.debug_tuple("HttpBody::Plain").field(b).finish(),
+            HttpBody::WithT(_) => f.write_str("HttpBody::WithT(..)"),
+        }
+    }
+}
+
+impl Body for HttpBody {
+    type Data = Bytes;
+    type Error = DynErr;
+
+    fn poll_frame(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
+        unsafe {
+            match self.get_unchecked_mut() {
+                HttpBody::Plain(b) => Pin::new_unchecked(b).poll_frame(cx),
+                HttpBody::WithT(b) => Pin::new_unchecked(b).poll_frame(cx),
+            }
+        }
+    }
+
+    fn size_hint(&self) -> SizeHint {
+        match self {
+            HttpBody::Plain(b) => b.size_hint(),
+            HttpBody::WithT(b) => b.size_hint(),
+        }
+    }
+}
+
+impl HttpBody {
+    pub fn stream<S>(s: S) -> Self
+    where
+        S: Stream<Item = Result<Bytes, Infallible>> + Send + Sync + 'static,
+    {
+        HttpBody::Plain(PlainBody::stream(s))
+    }
+
+    pub fn full(bytes: Bytes) -> Self {
+        HttpBody::Plain(PlainBody::full(bytes))
+    }
+
+    pub fn empty() -> Self {
+        HttpBody::Plain(PlainBody::empty())
+    }
+
+    pub fn with_trailers<Fut>(self, fut: Fut) -> Self
+    where
+        Fut: std::future::Future<Output = Option<Result<http::HeaderMap, DynErr>>>
+            + Send
+            + Sync
+            + 'static,
+    {
+        let boxed: BoxTrailers = Box::pin(fut);
+        match self {
+            HttpBody::Plain(p) => HttpBody::WithT(p.with_trailers(boxed)),
+            already @ HttpBody::WithT(_) => already,
+        }
+    }
+}
+
+pub type HttpResponse = http::Response<HttpBody>;
 pub type HttpRequest = Request<SizeLimitedIncoming<Incoming>>;
 
 pub trait ConversionExt {
@@ -32,7 +165,7 @@ pub trait PathExt {
     fn no_trailing_slash(&self) -> &str;
 }
 
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, Copy)]
 pub enum ResponseFormat {
     JSON,
     HTML,
@@ -64,7 +197,7 @@ impl PathExt for str {
         if self == "/" {
             self
         } else {
-            self.trim_end_matches("/")
+            self.trim_end_matches('/')
         }
     }
 }
@@ -91,7 +224,7 @@ impl RequestExt for HttpRequest {
     fn query_param(&self, query_name: &str) -> Option<&str> {
         self.uri()
             .query()
-            .and_then(|query| query.split('&').find(|param| param.starts_with(query_name)))
-            .map(|param| param.split('=').nth(1).unwrap_or(""))
+            .and_then(|q| q.split('&').find(|p| p.starts_with(query_name)))
+            .map(|p| p.split('=').nth(1).unwrap_or(""))
     }
 }

data/ext/itsi_server/src/server/io_stream.rs

@@ -2,6 +2,7 @@ use pin_project::pin_project;
 use tokio::net::{TcpStream, UnixStream};
 use tokio_rustls::server::TlsStream;
 
+use std::io::{self, IoSlice};
 use std::os::unix::io::{AsRawFd, RawFd};
 use std::pin::Pin;
 use std::task::{Context, Poll};
@@ -34,12 +35,12 @@ pub enum IoStream {
 }
 
 impl IoStream {
-    pub fn addr(&self) -> SockAddr {
+    pub fn addr(&self) -> String {
         match self {
-            IoStream::Tcp { addr, .. } => addr.clone(),
-            IoStream::TcpTls { addr, .. } => addr.clone(),
-            IoStream::Unix { addr, .. } => addr.clone(),
-            IoStream::UnixTls { addr, .. } => addr.clone(),
+            IoStream::Tcp { addr, .. } => addr.to_string(),
+            IoStream::TcpTls { addr, .. } => addr.to_string(),
+            IoStream::Unix { addr, .. } => addr.to_string(),
+            IoStream::UnixTls { addr, .. } => addr.to_string(),
         }
     }
 }
@@ -90,6 +91,28 @@ impl AsyncWrite for IoStream {
             IoStreamEnumProj::UnixTls { stream, .. } => stream.poll_shutdown(cx),
         }
     }
+
+    fn poll_write_vectored(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        bufs: &[IoSlice<'_>],
+    ) -> Poll<io::Result<usize>> {
+        match self.project() {
+            IoStreamEnumProj::Tcp { stream, .. } => stream.poll_write_vectored(cx, bufs),
+            IoStreamEnumProj::TcpTls { stream, .. } => stream.poll_write_vectored(cx, bufs),
+            IoStreamEnumProj::Unix { stream, .. } => stream.poll_write_vectored(cx, bufs),
+            IoStreamEnumProj::UnixTls { stream, .. } => stream.poll_write_vectored(cx, bufs),
+        }
+    }
+
+    fn is_write_vectored(&self) -> bool {
+        match self {
+            IoStream::Tcp { stream, .. } => stream.is_write_vectored(),
+            IoStream::TcpTls { stream, .. } => stream.is_write_vectored(),
+            IoStream::Unix { stream, .. } => stream.is_write_vectored(),
+            IoStream::UnixTls { stream, .. } => stream.is_write_vectored(),
+        }
+    }
 }
 
 impl AsRawFd for IoStream {

data/ext/itsi_server/src/server/lifecycle_event.rs

@@ -1,4 +1,4 @@
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq)]
 pub enum LifecycleEvent {
     Start,
     Shutdown,

data/ext/itsi_server/src/server/middleware_stack/middlewares/auth_basic.rs

@@ -3,7 +3,6 @@ use base64::{engine::general_purpose, Engine};
 use bytes::Bytes;
 use either::Either;
 use http::{Response, StatusCode};
-use http_body_util::{combinators::BoxBody, Full};
 use magnus::error::Result;
 use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
@@ -11,7 +10,7 @@ use std::str;
 use tracing::debug;
 
 use crate::{
-    server::http_message_types::{HttpRequest, HttpResponse, RequestExt},
+    server::http_message_types::{HttpBody, HttpRequest, HttpResponse, RequestExt},
     services::{itsi_http_service::HttpRequestContext, password_hasher::verify_password_hash},
 };
 
@@ -34,7 +33,7 @@ impl AuthBasic {
                 "WWW-Authenticate",
                 format!("Basic realm=\"{}\"", self.realm),
             )
-            .body(BoxBody::new(Full::new(Bytes::from("Unauthorized"))))
+            .body(HttpBody::full(Bytes::from("Unauthorized")))
             .unwrap()
     }
 }

data/ext/itsi_server/src/server/middleware_stack/middlewares/compression.rs

@@ -1,5 +1,5 @@
 use crate::{
-    server::http_message_types::{HttpRequest, HttpResponse},
+    server::http_message_types::{HttpBody, HttpRequest, HttpResponse},
     services::itsi_http_service::HttpRequestContext,
 };
 
@@ -20,8 +20,8 @@ use http::{
     header::{GetAll, CONTENT_ENCODING, CONTENT_LENGTH, CONTENT_TYPE},
     HeaderValue, Response,
 };
-use http_body_util::{combinators::BoxBody, BodyExt, Full, StreamBody};
-use hyper::body::{Body, Frame};
+use http_body_util::{BodyExt, StreamBody};
+use hyper::body::Body;
 use magnus::error::Result;
 use serde::{Deserialize, Serialize};
 use std::convert::Infallible;
@@ -126,15 +126,13 @@ impl MimeType {
     }
 }
 
-fn stream_encode<R>(encoder: R) -> BoxBody<Bytes, Infallible>
+fn stream_encode<R>(encoder: R) -> HttpBody
 where
     R: AsyncRead + Unpin + Sync + Send + 'static,
 {
-    let encoded_stream = ReaderStream::new(encoder).map(|res| {
-        res.map(Frame::data)
-            .map_err(|_| -> Infallible { unreachable!("We handle IO errors above") })
-    });
-    BoxBody::new(StreamBody::new(encoded_stream))
+    let encoded_stream = ReaderStream::new(encoder)
+        .map(|res| res.map_err(|_| -> Infallible { unreachable!("We handle IO errors above") }));
+    HttpBody::stream(StreamBody::new(encoded_stream))
 }
 
 fn update_content_encoding(parts: &mut http::response::Parts, new_encoding: HeaderValue) {
@@ -293,7 +291,7 @@ impl MiddlewareLayer for Compression {
                 }
                 CompressionAlgorithm::Identity => unreachable!(),
             };
-            BoxBody::new(Full::new(Bytes::from(compressed_bytes)))
+            HttpBody::full(Bytes::from(compressed_bytes))
         } else {
             let stream = body
                 .into_data_stream()

data/ext/itsi_server/src/server/middleware_stack/middlewares/cors.rs

@@ -1,12 +1,11 @@
 use super::{FromValue, MiddlewareLayer};
 use crate::{
-    server::http_message_types::{HttpRequest, HttpResponse, RequestExt},
+    server::http_message_types::{HttpBody, HttpRequest, HttpResponse, RequestExt},
     services::itsi_http_service::HttpRequestContext,
 };
 
 use async_trait::async_trait;
 use http::{HeaderMap, Method, Response};
-use http_body_util::{combinators::BoxBody, Empty};
 use itsi_error::ItsiError;
 use magnus::error::Result;
 use serde::Deserialize;
@@ -273,7 +272,7 @@ impl MiddlewareLayer for Cors {
             let mut response_builder = Response::builder().status(204);
             *response_builder.headers_mut().unwrap() = headers;
             let response = response_builder
-                .body(BoxBody::new(Empty::new()))
+                .body(HttpBody::empty())
                 .map_err(ItsiError::new)?;
             return Ok(either::Either::Right(response));
         }

data/ext/itsi_server/src/server/middleware_stack/middlewares/csp.rs

@@ -1,6 +1,6 @@
 use super::FromValue;
 use crate::{
-    server::http_message_types::{HttpRequest, HttpResponse},
+    server::http_message_types::{HttpBody, HttpRequest, HttpResponse},
     services::itsi_http_service::HttpRequestContext,
 };
 use async_trait::async_trait;
@@ -8,7 +8,7 @@ use bytes::{Bytes, BytesMut};
 use either::Either;
 use futures::TryStreamExt;
 use http::{HeaderValue, StatusCode};
-use http_body_util::{combinators::BoxBody, BodyExt, Empty};
+use http_body_util::BodyExt;
 use itsi_error::ItsiError;
 use serde::{Deserialize, Serialize};
 use std::sync::Arc;
@@ -164,7 +164,7 @@ impl super::MiddlewareLayer for Csp {
                 }
             }
 
-            let mut resp = HttpResponse::new(BoxBody::new(Empty::new()));
+            let mut resp = HttpResponse::new(HttpBody::empty());
             *resp.status_mut() = StatusCode::NO_CONTENT;
             return Ok(Either::Right(resp));
         }