v0.24.5

src/proxy/body.rs

@@ -0,0 +1,439 @@
+//! Passes frames through unchanged, including trailers. No body collection.
+use crate::{
+    logging::{now_ms, LogEntry, Logs},
+    metrics::Metrics,
+};
+use axum::body::Body;
+use bytes::Bytes;
+use http_body::{Body as HttpBody, Frame, SizeHint};
+use std::{
+    future::Future,
+    pin::Pin,
+    sync::{
+        atomic::{AtomicBool, AtomicU16, AtomicU64, Ordering},
+        Arc, Mutex,
+    },
+    task::{Context, Poll},
+    time::{Duration, Instant},
+};
+use tokio::{
+    sync::OwnedSemaphorePermit,
+    time::{sleep, Sleep},
+};
+
+pub struct Exchange {
+    pub vhost_metrics: Option<Arc<crate::metrics::VhostMetrics>>,
+    pub vhost: Option<String>,
+    pub route: Mutex<Option<String>>,
+    pub cache_effective_enabled: Mutex<Option<bool>>,
+    pub transform_effective_enabled: Mutex<Option<bool>>,
+    pub backend: Mutex<Option<Arc<crate::balance::Lease>>>,
+    pub id: u64,
+    pub method: String,
+    pub host: String,
+    pub path: String,
+    pub http_version: &'static str,
+    pub transport: &'static str,
+    pub websocket_bootstrap: Option<&'static str>,
+    pub websocket_phase: Mutex<Option<&'static str>>,
+    pub websocket_error: Mutex<Option<&'static str>>,
+    pub h3_reset_code: Mutex<Option<&'static str>>,
+    pub upstream_status: AtomicU16,
+    pub upstream_error: Mutex<Option<String>>,
+    pub client_ip: Mutex<Option<String>>,
+    pub started: Instant,
+    pub status: AtomicU16,
+    pub cache_status: AtomicU16,
+    pub cache_reason: Mutex<&'static str>,
+    pub failed: AtomicBool,
+    pub aborted: AtomicBool,
+    pub request_bytes: AtomicU64,
+    pub response_bytes: AtomicU64,
+    pub metrics: Arc<Metrics>,
+    pub logs: Arc<Logs>,
+    pub _permit: OwnedSemaphorePermit,
+}
+impl Exchange {
+    pub fn websocket_phase(&self, phase: &'static str) {
+        if self.websocket_bootstrap.is_some() {
+            *self.websocket_phase.lock().unwrap_or_else(|e| e.into_inner()) = Some(phase);
+        }
+    }
+    pub fn websocket_error(&self, phase: &'static str, error: &'static str) {
+        if self.websocket_bootstrap.is_some() {
+            *self.websocket_phase.lock().unwrap_or_else(|e| e.into_inner()) = Some(phase);
+            *self.websocket_error.lock().unwrap_or_else(|e| e.into_inner()) = Some(error);
+        }
+    }
+    pub fn h3_reset(&self, code: &'static str) {
+        *self.h3_reset_code.lock().unwrap_or_else(|e| e.into_inner()) = Some(code);
+    }
+    pub fn backend_bytes(&self, bytes: u64, incoming: bool) {
+        if let Some(backend) = self
+            .backend
+            .lock()
+            .unwrap_or_else(|e| e.into_inner())
+            .as_ref()
+        {
+            backend.bytes(bytes, incoming);
+        }
+    }
+    fn backend_error(&self, reason: &'static str) {
+        if let Some(backend) = self
+            .backend
+            .lock()
+            .unwrap_or_else(|e| e.into_inner())
+            .as_ref()
+        {
+            backend.error(reason, self.started.elapsed());
+        }
+    }
+}
+impl Drop for Exchange {
+    fn drop(&mut self) {
+        let status = self.status.load(Ordering::Relaxed);
+        let failed = self.failed.load(Ordering::Relaxed);
+        let aborted = self.aborted.load(Ordering::Relaxed);
+        let duration = self.started.elapsed();
+        if let Some(m) = &self.vhost_metrics {
+            m.complete(
+                status,
+                failed || aborted,
+                duration,
+                self.request_bytes.load(Ordering::Relaxed),
+                self.response_bytes.load(Ordering::Relaxed),
+            );
+        }
+        self.metrics.active.fetch_sub(1, Ordering::Relaxed);
+        self.metrics.completed.fetch_add(1, Ordering::Relaxed);
+        self.metrics
+            .duration_us
+            .fetch_add(duration.as_micros() as u64, Ordering::Relaxed);
+        if failed || aborted || status >= 500 {
+            self.metrics.errors.fetch_add(1, Ordering::Relaxed);
+        }
+        match status / 100 {
+            1 => &self.metrics.status_1xx,
+            2 => &self.metrics.status_2xx,
+            3 => &self.metrics.status_3xx,
+            4 => &self.metrics.status_4xx,
+            _ => &self.metrics.status_5xx,
+        }
+        .fetch_add(1, Ordering::Relaxed);
+        let backend = self
+            .backend
+            .lock()
+            .unwrap_or_else(|e| e.into_inner())
+            .as_ref()
+            .map(|l| l.member.backend.id.clone());
+        let websocket_phase = *self.websocket_phase.lock().unwrap_or_else(|e| e.into_inner());
+        self.logs.push(LogEntry {
+            cache_reason: *self.cache_reason.lock().unwrap_or_else(|e| e.into_inner()),
+            vhost: self.vhost.clone(),
+            route: self.route.lock().unwrap_or_else(|e| e.into_inner()).clone(),
+            cache_effective_enabled: *self.cache_effective_enabled.lock().unwrap_or_else(|e| e.into_inner()),
+            transform_effective_enabled: *self.transform_effective_enabled.lock().unwrap_or_else(|e| e.into_inner()),
+            backend,
+            id: self.id,
+            time_ms: now_ms(),
+            method: self.method.clone(),
+            host: self.host.clone(),
+            path: self.path.clone(),
+            http_version: self.http_version,
+            transport: self.transport,
+            websocket_bootstrap: self.websocket_bootstrap,
+            websocket_phase,
+            websocket_error: *self.websocket_error.lock().unwrap_or_else(|e| e.into_inner()),
+            h3_reset_code: *self.h3_reset_code.lock().unwrap_or_else(|e| e.into_inner()),
+            upstream_status: match self.upstream_status.load(Ordering::Relaxed) {
+                0 => None,
+                value => Some(value),
+            },
+            upstream_error: self.upstream_error.lock().unwrap_or_else(|e| e.into_inner()).clone(),
+            client_ip: self.client_ip.lock().unwrap_or_else(|e| e.into_inner()).clone(),
+            status,
+            duration_ms: duration.as_millis() as u64,
+            cache: match self.cache_status.load(Ordering::Relaxed) {
+                1 => "HIT",
+                2 => "MISS",
+                3 => "BYPASS",
+                _ => "-",
+            },
+            request_bytes: self.request_bytes.load(Ordering::Relaxed),
+            response_bytes: self.response_bytes.load(Ordering::Relaxed),
+            outcome: if failed {
+                "body_error"
+            } else if aborted {
+                "cancelled"
+            } else if status >= 500 {
+                "proxy_or_upstream_error"
+            } else if status == 101 || websocket_phase == Some("tunnel_closed") {
+                "websocket_closed"
+            } else {
+                "complete"
+            },
+        });
+    }
+}
+#[derive(Clone, Copy)]
+pub enum Direction {
+    Request,
+    Response,
+}
+pub struct MeteredBody {
+    inner: Body,
+    exchange: Option<Arc<Exchange>>,
+    direction: Direction,
+    timer: Pin<Box<Sleep>>,
+    idle: Duration,
+    finished: bool,
+}
+impl MeteredBody {
+    pub fn new(inner: Body, exchange: Arc<Exchange>, direction: Direction, idle: Duration) -> Self {
+        let finished = inner.is_end_stream();
+        Self {
+            inner,
+            exchange: Some(exchange),
+            direction,
+            timer: Box::pin(sleep(idle)),
+            idle,
+            finished,
+        }
+    }
+    fn finish(&mut self) {
+        self.finished = true;
+        self.exchange.take();
+    }
+}
+impl HttpBody for MeteredBody {
+    type Data = Bytes;
+    type Error = axum::Error;
+    fn poll_frame(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Option<Result<Frame<Bytes>, Self::Error>>> {
+        let this = self.get_mut();
+        if this.finished {
+            this.exchange.take();
+            return Poll::Ready(None);
+        }
+        match Pin::new(&mut this.inner).poll_frame(cx) {
+            Poll::Ready(Some(Ok(frame))) => {
+                if let (Some(data), Some(exchange)) = (frame.data_ref(), this.exchange.as_ref()) {
+                    let count = data.len() as u64;
+                    exchange.backend_bytes(count, matches!(this.direction, Direction::Request));
+                    match this.direction {
+                        Direction::Request => {
+                            exchange.request_bytes.fetch_add(count, Ordering::Relaxed);
+                            exchange
+                                .metrics
+                                .bytes_in
+                                .fetch_add(count, Ordering::Relaxed);
+                        }
+                        Direction::Response => {
+                            exchange.response_bytes.fetch_add(count, Ordering::Relaxed);
+                            exchange
+                                .metrics
+                                .bytes_out
+                                .fetch_add(count, Ordering::Relaxed);
+                        }
+                    }
+                }
+                this.timer
+                    .as_mut()
+                    .reset(tokio::time::Instant::now() + this.idle);
+                if this.inner.is_end_stream() {
+                    this.finish();
+                }
+                Poll::Ready(Some(Ok(frame)))
+            }
+            Poll::Ready(Some(Err(error))) => {
+                if let Some(exchange) = &this.exchange {
+                    exchange.failed.store(true, Ordering::Relaxed);
+                    exchange.backend_error(match this.direction {
+                        Direction::Request => "send_error",
+                        Direction::Response => "read_error",
+                    });
+                }
+                this.finish();
+                Poll::Ready(Some(Err(error)))
+            }
+            Poll::Ready(None) => {
+                this.finish();
+                Poll::Ready(None)
+            }
+            Poll::Pending => {
+                if this.timer.as_mut().poll(cx).is_ready() {
+                    if let Some(exchange) = &this.exchange {
+                        exchange.failed.store(true, Ordering::Relaxed);
+                        exchange.backend_error(match this.direction {
+                            Direction::Request => "send_timeout",
+                            Direction::Response => "read_timeout",
+                        });
+                    }
+                    this.finish();
+                    Poll::Ready(Some(Err(axum::Error::new(std::io::Error::new(
+                        std::io::ErrorKind::TimedOut,
+                        "body idle timeout",
+                    )))))
+                } else {
+                    Poll::Pending
+                }
+            }
+        }
+    }
+    fn is_end_stream(&self) -> bool {
+        self.finished || self.inner.is_end_stream()
+    }
+    fn size_hint(&self) -> SizeHint {
+        self.inner.size_hint()
+    }
+}
+impl Drop for MeteredBody {
+    fn drop(&mut self) {
+        if !self.finished
+            && !self.inner.is_end_stream()
+            && matches!(self.direction, Direction::Response)
+        {
+            if let Some(exchange) = &self.exchange {
+                exchange.aborted.store(true, Ordering::Relaxed);
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use http_body_util::BodyExt;
+    use tokio::sync::Semaphore;
+    async fn observation() -> (Arc<Exchange>, Arc<Semaphore>) {
+        let slots = Arc::new(Semaphore::new(1));
+        let permit = slots.clone().acquire_owned().await.unwrap();
+        let metrics = Arc::new(Metrics::default());
+        metrics.active.store(1, Ordering::Relaxed);
+        let exchange = Arc::new(Exchange {
+            vhost_metrics: None,
+            vhost: None,
+            route: Mutex::new(None),
+            cache_effective_enabled: Mutex::new(None),
+            transform_effective_enabled: Mutex::new(None),
+            backend: Mutex::new(None),
+            id: 1,
+            method: "POST".into(),
+            host: "localhost".into(),
+            path: "/echo".into(),
+            http_version: "1.1",
+            transport: "h1",
+            websocket_bootstrap: None,
+            websocket_phase: Mutex::new(None),
+            websocket_error: Mutex::new(None),
+            h3_reset_code: Mutex::new(None),
+            upstream_status: AtomicU16::new(0),
+            upstream_error: Mutex::new(None),
+            client_ip: Mutex::new(Some("127.0.0.1".into())),
+            started: Instant::now(),
+            status: AtomicU16::new(200),
+            cache_status: AtomicU16::new(0),
+            cache_reason: Mutex::new("not_proxy_response"),
+            failed: AtomicBool::new(false),
+            aborted: AtomicBool::new(false),
+            request_bytes: AtomicU64::new(0),
+            response_bytes: AtomicU64::new(0),
+            metrics,
+            logs: Arc::new(Logs::default()),
+            _permit: permit,
+        });
+        (exchange, slots)
+    }
+    #[tokio::test]
+    async fn permit_and_metrics_cover_both_bodies() {
+        let (exchange, slots) = observation().await;
+        let metrics = exchange.metrics.clone();
+        let logs = exchange.logs.clone();
+        let input = MeteredBody::new(
+            Body::from("in"),
+            exchange.clone(),
+            Direction::Request,
+            Duration::from_secs(1),
+        );
+        let output = MeteredBody::new(
+            Body::from("out"),
+            exchange.clone(),
+            Direction::Response,
+            Duration::from_secs(1),
+        );
+        drop(exchange);
+        assert_eq!(
+            input.collect().await.unwrap().to_bytes(),
+            Bytes::from_static(b"in")
+        );
+        assert_eq!(slots.available_permits(), 0);
+        assert_eq!(
+            output.collect().await.unwrap().to_bytes(),
+            Bytes::from_static(b"out")
+        );
+        assert_eq!(slots.available_permits(), 1);
+        assert_eq!(metrics.active.load(Ordering::Relaxed), 0);
+        assert_eq!(metrics.bytes_in.load(Ordering::Relaxed), 2);
+        assert_eq!(metrics.bytes_out.load(Ordering::Relaxed), 3);
+        assert_eq!(metrics.completed.load(Ordering::Relaxed), 1);
+        assert_eq!(logs.recent().len(), 1);
+    }
+    #[tokio::test]
+    async fn idle_body_fails_and_releases_capacity() {
+        let (exchange, slots) = observation().await;
+        let metrics = exchange.metrics.clone();
+        let pending = futures_util::stream::pending::<Result<Bytes, std::io::Error>>();
+        let output = MeteredBody::new(
+            Body::from_stream(pending),
+            exchange,
+            Direction::Response,
+            Duration::from_millis(10),
+        );
+        assert!(output.collect().await.is_err());
+        assert_eq!(slots.available_permits(), 1);
+        assert_eq!(metrics.errors.load(Ordering::Relaxed), 1);
+    }
+}
+
+/// Retains the first inspected frame without collecting the rest of a request.
+/// A trailer/empty DATA frame still prevents storing a request-body variant.
+pub struct ReplayBody {
+    head: Option<http_body::Frame<bytes::Bytes>>,
+    tail: axum::body::Body,
+}
+impl ReplayBody {
+    pub fn new(head: http_body::Frame<bytes::Bytes>, tail: axum::body::Body) -> Self {
+        Self {
+            head: Some(head),
+            tail,
+        }
+    }
+}
+impl http_body::Body for ReplayBody {
+    type Data = bytes::Bytes;
+    type Error = axum::Error;
+    fn poll_frame(
+        mut self: std::pin::Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<Option<Result<http_body::Frame<Self::Data>, Self::Error>>> {
+        if let Some(head) = self.head.take() {
+            return std::task::Poll::Ready(Some(Ok(head)));
+        }
+        std::pin::Pin::new(&mut self.tail).poll_frame(cx)
+    }
+    fn is_end_stream(&self) -> bool {
+        self.head.is_none() && self.tail.is_end_stream()
+    }
+    fn size_hint(&self) -> http_body::SizeHint {
+        let mut hint = http_body::SizeHint::new();
+        if self.head.is_none() {
+            return self.tail.size_hint();
+        }
+        if let Some(data) = self.head.as_ref().and_then(|f| f.data_ref()) {
+            hint.set_lower(data.len() as u64);
+        }
+        hint
+    }
+}