use std::convert::Infallible;
use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};
use axum::body::Body as AxumBody;
use axum::http::Request as AxumRequest;
use axum::response::Response as AxumResponse;
use http_body_util::BodyExt;
use tower::{Layer, Service, ServiceExt};
use vercel_runtime::axum::StreamingUtils;
use vercel_runtime::{AppState, Error as VercelError, Request as VercelRequest, ResponseBody};
#[derive(Clone, Default)]
pub struct StreamingVercelLayer;
impl StreamingVercelLayer {
pub fn new() -> Self {
Self
}
}
impl<S> Layer<S> for StreamingVercelLayer {
type Service = StreamingVercelService<S>;
fn layer(&self, inner: S) -> Self::Service {
StreamingVercelService { inner }
}
}
#[derive(Clone)]
pub struct StreamingVercelService<S> {
inner: S,
}
impl<S> Service<(AppState, VercelRequest)> for StreamingVercelService<S>
where
S: Service<AxumRequest<AxumBody>, Response = AxumResponse<AxumBody>, Error = Infallible>
+ Send
+ Clone
+ 'static,
S::Future: Send + 'static,
{
type Response = hyper::Response<ResponseBody>;
type Error = VercelError;
type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
match self.inner.poll_ready(cx) {
Poll::Ready(Ok(())) => Poll::Ready(Ok(())),
Poll::Ready(Err(_)) => Poll::Ready(Err("Inner service error".into())),
Poll::Pending => Poll::Pending,
}
}
fn call(&mut self, (_state, req): (AppState, VercelRequest)) -> Self::Future {
let mut service = self.inner.clone();
Box::pin(async move {
let (parts, body) = req.into_parts();
let body_bytes = BodyExt::collect(body)
.await
.map_err(|e| Box::new(e) as VercelError)?
.to_bytes();
let axum_req = AxumRequest::from_parts(parts, AxumBody::from(body_bytes));
let ready = ServiceExt::ready(&mut service)
.await
.map_err(|_| "Service not ready")?;
let axum_resp = Service::call(ready, axum_req)
.await
.map_err(|_| "Service error")?;
let (resp_parts, resp_body) = axum_resp.into_parts();
let stream_body = StreamingUtils::create_stream_body(resp_body).await?;
Ok(hyper::Response::from_parts(resp_parts, stream_body))
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use axum::Router;
use axum::routing::get;
#[test]
fn layer_composes_with_axum_router() {
let _: StreamingVercelService<Router> =
StreamingVercelLayer::new().layer(Router::new().route("/", get(|| async { "ok" })));
}
}