1 files changed, 41 insertions, 56 deletions

diff --git a/rust/fatcat-openapi/src/context.rs b/rust/fatcat-openapi/src/context.rs
index 57d11be..d782855 100644
--- a/rust/fatcat-openapi/src/context.rs
+++ b/rust/fatcat-openapi/src/context.rs
@@ -1,13 +1,12 @@
 use crate::Api;
-use futures::Future;
-use hyper;
+use futures::future::BoxFuture;
 use hyper::header::HeaderName;
-use hyper::{body::Payload, service::Service, Error, Request, Response, StatusCode};
+use hyper::{service::Service, Error, Request, Response, StatusCode};
 use std::default::Default;
 use std::io;
 use std::marker::PhantomData;
+use std::task::{Context, Poll};
 use swagger::auth::{AuthData, Authorization, Bearer, Scopes};
-use swagger::context::ContextualPayload;
 use swagger::{EmptyContext, Has, Pop, Push, XSpanIdString};
 use url::form_urlencoded;
 
@@ -31,55 +30,49 @@ where
 }
 
 // Make a service that adds context.
-impl<'a, T, SC, A, B, C, D, E, ME, S, OB, F> hyper::service::MakeService<&'a SC>
-    for MakeAddContext<T, A>
+impl<Target, T, A, B, C, D> Service<Target> for MakeAddContext<T, A>
 where
-    A: Default + Push<XSpanIdString, Result = B>,
+    Target: Send,
+    A: Default + Push<XSpanIdString, Result = B> + Send,
     B: Push<Option<AuthData>, Result = C>,
     C: Push<Option<Authorization>, Result = D>,
     D: Send + 'static,
-    T: hyper::service::MakeService<
-        &'a SC,
-        Error = E,
-        MakeError = ME,
-        Service = S,
-        ReqBody = ContextualPayload<hyper::Body, D>,
-        ResBody = OB,
-        Future = F,
-    >,
-    S: Service<Error = E, ReqBody = ContextualPayload<hyper::Body, D>, ResBody = OB> + 'static,
-    ME: swagger::ErrorBound,
-    E: swagger::ErrorBound,
-    F: Future<Item = S, Error = ME> + Send + 'static,
-    S::Future: Send,
-    OB: Payload,
+    T: Service<Target> + Send,
+    T::Future: Send + 'static,
 {
-    type ReqBody = hyper::Body;
-    type ResBody = OB;
-    type Error = E;
-    type MakeError = ME;
-    type Service = AddContext<S, A>;
-    type Future = Box<dyn Future<Item = Self::Service, Error = ME> + Send + 'static>;
+    type Error = T::Error;
+    type Response = AddContext<T::Response, A, B, C, D>;
+    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
 
-    fn make_service(&mut self, ctx: &'a SC) -> Self::Future {
-        Box::new(self.inner.make_service(ctx).map(|s| AddContext::new(s)))
+    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        self.inner.poll_ready(cx)
+    }
+
+    fn call(&mut self, target: Target) -> Self::Future {
+        let service = self.inner.call(target);
+
+        Box::pin(async move { Ok(AddContext::new(service.await?)) })
     }
 }
 
-/// Middleware to extract authentication data from request
-pub struct AddContext<T, A> {
+/// Middleware to add context data from the request
+pub struct AddContext<T, A, B, C, D>
+where
+    A: Default + Push<XSpanIdString, Result = B>,
+    B: Push<Option<AuthData>, Result = C>,
+    C: Push<Option<Authorization>, Result = D>,
+{
     inner: T,
     marker: PhantomData<A>,
 }
 
-impl<T, A, B, C, D> AddContext<T, A>
+impl<T, A, B, C, D> AddContext<T, A, B, C, D>
 where
     A: Default + Push<XSpanIdString, Result = B>,
     B: Push<Option<AuthData>, Result = C>,
     C: Push<Option<Authorization>, Result = D>,
-    T: Service,
 {
-    pub fn new(inner: T) -> AddContext<T, A> {
+    pub fn new(inner: T) -> Self {
         AddContext {
             inner,
             marker: PhantomData,
@@ -87,24 +80,25 @@ where
     }
 }
 
-impl<T, A, B, C, D> Service for AddContext<T, A>
+impl<T, A, B, C, D, ReqBody> Service<Request<ReqBody>> for AddContext<T, A, B, C, D>
 where
     A: Default + Push<XSpanIdString, Result = B>,
     B: Push<Option<AuthData>, Result = C>,
     C: Push<Option<Authorization>, Result = D>,
     D: Send + 'static,
-    T: Service<ReqBody = ContextualPayload<hyper::Body, D>>,
-    T::Future: Future<Item = Response<T::ResBody>, Error = T::Error> + Send + 'static,
+    T: Service<(Request<ReqBody>, D)>,
 {
-    type ReqBody = hyper::Body;
-    type ResBody = T::ResBody;
     type Error = T::Error;
-    type Future = Box<dyn Future<Item = Response<T::ResBody>, Error = T::Error> + Send + 'static>;
+    type Future = T::Future;
+    type Response = T::Response;
 
-    fn call(&mut self, req: Request<Self::ReqBody>) -> Self::Future {
-        let context = A::default().push(XSpanIdString::get_or_generate(&req));
-        let (head, body) = req.into_parts();
-        let headers = head.headers.clone();
+    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        self.inner.poll_ready(cx)
+    }
+
+    fn call(&mut self, request: Request<ReqBody>) -> Self::Future {
+        let context = A::default().push(XSpanIdString::get_or_generate(&request));
+        let headers = request.headers();
 
         {
             use std::ops::Deref;
@@ -114,22 +108,13 @@ where
                 let context = context.push(Some(auth_data));
                 let context = context.push(None::<Authorization>);
 
-                let body = ContextualPayload {
-                    inner: body,
-                    context: context,
-                };
-
-                return Box::new(self.inner.call(hyper::Request::from_parts(head, body)));
+                return self.inner.call((request, context));
             }
         }
 
         let context = context.push(None::<AuthData>);
         let context = context.push(None::<Authorization>);
-        let body = ContextualPayload {
-            inner: body,
-            context: context,
-        };
 
-        Box::new(self.inner.call(hyper::Request::from_parts(head, body)))
+        self.inner.call((request, context))
     }
 }