Add option_layer (#1696)

axum-extra/src/either.rs

@@ -93,12 +93,16 @@
 //! [`BytesRejection`]: axum::extract::rejection::BytesRejection
 //! [`IntoResponse::into_response`]: https://docs.rs/axum/0.5/axum/response/index.html#returning-different-response-types
 
+use std::task::{Context, Poll};
+
 use axum::{
     async_trait,
     extract::FromRequestParts,
     response::{IntoResponse, Response},
 };
 use http::request::Parts;
+use tower_layer::Layer;
+use tower_service::Service;
 
 /// Combines two extractors or responses into a single type.
 ///
@@ -267,3 +271,42 @@ impl_traits_for_either!(Either5 => [E1, E2, E3, E4], E5);
 impl_traits_for_either!(Either6 => [E1, E2, E3, E4, E5], E6);
 impl_traits_for_either!(Either7 => [E1, E2, E3, E4, E5, E6], E7);
 impl_traits_for_either!(Either8 => [E1, E2, E3, E4, E5, E6, E7], E8);
+
+impl<E1, E2, S> Layer<S> for Either<E1, E2>
+where
+    E1: Layer<S>,
+    E2: Layer<S>,
+{
+    type Service = Either<E1::Service, E2::Service>;
+
+    fn layer(&self, inner: S) -> Self::Service {
+        match self {
+            Either::E1(layer) => Either::E1(layer.layer(inner)),
+            Either::E2(layer) => Either::E2(layer.layer(inner)),
+        }
+    }
+}
+
+impl<R, E1, E2> Service<R> for Either<E1, E2>
+where
+    E1: Service<R>,
+    E2: Service<R, Response = E1::Response, Error = E1::Error>,
+{
+    type Response = E1::Response;
+    type Error = E1::Error;
+    type Future = futures_util::future::Either<E1::Future, E2::Future>;
+
+    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        match self {
+            Either::E1(inner) => inner.poll_ready(cx),
+            Either::E2(inner) => inner.poll_ready(cx),
+        }
+    }
+
+    fn call(&mut self, req: R) -> Self::Future {
+        match self {
+            Either::E1(inner) => futures_util::future::Either::Left(inner.call(req)),
+            Either::E2(inner) => futures_util::future::Either::Right(inner.call(req)),
+        }
+    }
+}