relay-core-lib 0.4.1

use crate::interceptor::{BoxError, HttpBody};
use hyper::body::{Body, Bytes, Frame, SizeHint};
use std::pin::Pin;
use std::task::{Context, Poll};
use std::time::Duration;
use tokio::time::Instant;

/// Wraps a body stream with bandwidth throttling (bytes/sec).
/// Inserts artificial delays between data frames to ensure
/// throughput does not exceed the configured rate.
pub struct ThrottleBody {
    inner: HttpBody,
    bytes_per_sec: u64,
    last_frame_at: Option<Instant>,
}

impl ThrottleBody {
    pub fn new(inner: HttpBody, bytes_per_sec: u64) -> Self {
        Self {
            inner,
            bytes_per_sec,
            last_frame_at: None,
        }
    }
}

impl Body for ThrottleBody {
    type Data = Bytes;
    type Error = BoxError;

    fn poll_frame(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
        let frame = match Pin::new(&mut self.inner).poll_frame(cx) {
            Poll::Ready(Some(Ok(frame))) => frame,
            other => return other,
        };

        // Calculate per-frame delay based on data size
        if let Some(data) = frame.data_ref() {
            let bytes = data.len() as u64;
            if bytes > 0 && self.bytes_per_sec > 0 {
                let frame_dur = Duration::from_micros(bytes * 1_000_000 / self.bytes_per_sec);
                let now = Instant::now();

                if let Some(last) = self.last_frame_at {
                    let elapsed = now.duration_since(last);
                    if elapsed < frame_dur {
                        let remaining = frame_dur - elapsed;
                        // Since we can't .await in poll_frame, schedule a wake
                        let waker = cx.waker().clone();
                        tokio::spawn(async move {
                            tokio::time::sleep(remaining).await;
                            waker.wake();
                        });
                        return Poll::Pending;
                    }
                }
                self.last_frame_at = Some(now);
            }
        }

        Poll::Ready(Some(Ok(frame)))
    }

    fn is_end_stream(&self) -> bool {
        self.inner.is_end_stream()
    }

    fn size_hint(&self) -> SizeHint {
        self.inner.size_hint()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use bytes::Bytes;
    use http_body_util::{BodyExt, Full};

    #[tokio::test]
    async fn test_throttle_body_preserves_data() {
        let data = Bytes::from("test-body-data");
        let body: HttpBody = Full::new(data.clone())
            .map_err(|e| -> BoxError { Box::new(e) })
            .boxed();
        // High rate limit — no effective throttling
        let throttled = ThrottleBody::new(body, 1_000_000);
        let collected = throttled.collect().await.unwrap().to_bytes();
        assert_eq!(collected, data);
    }

    #[tokio::test]
    async fn test_throttle_body_passthrough_empty() {
        let body: HttpBody = Full::new(Bytes::new())
            .map_err(|e| -> BoxError { Box::new(e) })
            .boxed();
        let throttled = ThrottleBody::new(body, 1000);
        let collected = throttled.collect().await.unwrap().to_bytes();
        assert_eq!(collected.len(), 0);
    }
}