aws_smithy_http_server/schema/routing/
service.rs

1/*
2 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3 * SPDX-License-Identifier: Apache-2.0
4 */
5
6//! The multi-protocol routing service: dispatch state, the claim walk, and its futures.
7
8use super::MultiProtocolRoutingServiceBuilder;
9use crate::config::RequestBodyCollectionConfig;
10use crate::routing::SyncRoute;
11use crate::schema::routing::RoutingError;
12use crate::schema::{OperationSchema, ServiceSchema};
13use crate::{
14    body::BoxBody,
15    error::BoxError,
16    schema::{SelectedOperation, SharedServerProtocol},
17};
18use bytes::Bytes;
19use http::{Request, Response};
20
21use crate::schema::routing::{
22    BodyRouteClaim, CollectedBody, MetadataProtocolRouter, OperationTarget, RouteClaim, SharedProtocolRouter,
23};
24use std::{
25    convert::Infallible,
26    fmt,
27    future::Future,
28    pin::Pin,
29    sync::Arc,
30    task::{Context, Poll},
31};
32use tower::Service;
33
34pub(super) struct BoundHandler {
35    pub(super) operation: &'static OperationSchema<'static>,
36    pub(super) collection_config: RequestBodyCollectionConfig,
37    pub(super) handler_route: SyncRoute<crate::body::RequestBody>,
38}
39impl Clone for BoundHandler {
40    fn clone(&self) -> Self {
41        Self {
42            operation: self.operation,
43            collection_config: self.collection_config,
44            handler_route: self.handler_route.clone(),
45        }
46    }
47}
48impl fmt::Debug for BoundHandler {
49    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
50        f.debug_struct("BoundHandler")
51            .field("operation", &self.operation.shape_id())
52            .finish()
53    }
54}
55
56/// A served protocol and the router it built.
57#[derive(Debug)]
58pub(super) struct ProtocolAndRouter {
59    pub(super) router: SharedProtocolRouter,
60    pub(super) protocol: SharedServerProtocol,
61}
62
63/// A service routing normalized requests to a shared handler collection through the protocols the
64/// service declares.
65///
66/// Hyper incoming bodies flow to handlers without boxing. Other compatible request bodies
67/// are accepted and erased into a boxed state on entry; normalized request bodies retain
68/// their existing representation.
69///
70/// Clones share all routing state through one `Arc`. Routing clones only the
71/// selected handler, so cloning the service for each request stays cheap.
72pub struct MultiProtocolRoutingService {
73    pub(super) state: Arc<RoutingState>,
74}
75
76/// Routers, handlers, and configuration shared by service clones and in-flight requests.
77pub(super) struct RoutingState {
78    /// The served protocols in priority order. Every configuration checks claims;
79    /// a single metadata router can do so without an async routing loop.
80    pub(super) protocols: Box<[ProtocolAndRouter]>,
81    pub(super) handlers: Box<[BoundHandler]>,
82    /// Indices into `protocols` of metadata routers checked for streaming inputs before body
83    /// collection. `None` when the service has no streaming-input operations; otherwise includes
84    /// all metadata routers. Recognized streaming inputs skip claims that require body I/O.
85    pub(super) metadata_routers: Option<Box<[usize]>>,
86    /// The provisional allowance the service collects under for body-first routing.
87    pub(super) body_collection_config: RequestBodyCollectionConfig,
88}
89impl Clone for MultiProtocolRoutingService {
90    fn clone(&self) -> Self {
91        Self {
92            state: self.state.clone(),
93        }
94    }
95}
96impl fmt::Debug for MultiProtocolRoutingService {
97    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
98        f.debug_struct("MultiProtocolRoutingService")
99            .field("protocols", &self.state.protocols)
100            .field("handlers", &self.state.handlers)
101            .finish()
102    }
103}
104
105/// The response when no protocol claims the request or offers a deferred rejection.
106///
107/// No protocol owns the request, so no protocol frames the response. Return `404` with
108/// the fixed XML body `<UnknownOperationException/>\n` and no `Content-Type` header,
109/// regardless of the protocols the service serves.
110fn unclaimed() -> Response<BoxBody> {
111    tracing::debug!(error = "no protocol claimed the request", "failed to route");
112    Response::builder()
113        .status(http::StatusCode::NOT_FOUND)
114        .body(crate::body::to_boxed("<UnknownOperationException/>\n"))
115        .expect("a status and static body response is valid")
116}
117
118impl MultiProtocolRoutingService {
119    /// Hands the routed request to its handler, recording the selection for downstream consumers.
120    fn handle(
121        &self,
122        selected: OperationTarget,
123        protocol: usize,
124        mut request: Request<crate::body::RequestBody>,
125    ) -> crate::routing::route::SyncRouteFuture<crate::body::RequestBody> {
126        let binding = &self.state.handlers[selected.index()];
127        debug_assert!(
128            std::ptr::eq(binding.operation, selected.operation()),
129            "router index belongs to a different operation"
130        );
131        SelectedOperation::new(
132            self.state.protocols[protocol].protocol.clone(),
133            binding.operation,
134            binding.collection_config,
135        )
136        .insert_into(request.extensions_mut());
137        binding.handler_route.clone().call_owned(request)
138    }
139
140    fn route_request(&self, request: Request<crate::body::RequestBody>) -> MultiProtocolRoutingFuture {
141        match &self.state.protocols[..] {
142            // A single metadata router can check its claim synchronously.
143            [ProtocolAndRouter {
144                router: SharedProtocolRouter::Metadata(router),
145                ..
146            }] => MultiProtocolRoutingFuture {
147                inner: self.route(router.as_ref(), request),
148                routing: None,
149            },
150            _ => {
151                let (parts, body) = request.into_parts();
152                MultiProtocolRoutingFuture {
153                    inner: State::Routing,
154                    routing: Some(ProtocolRouting {
155                        service: self.clone(),
156                        probe: Request::from_parts(parts, ()),
157                        body: Some(body),
158                        collected: None,
159                        index: 0,
160                        streaming: None,
161                        deferred: None,
162                        deferred_claim: None,
163                    }),
164                }
165            }
166        }
167    }
168
169    /// Checks the service's only metadata protocol without allocating a routing future.
170    fn route(&self, router: &dyn MetadataProtocolRouter, request: Request<crate::body::RequestBody>) -> State {
171        // Probe with the head only: the parts move over and back, nothing is cloned,
172        // and the router stays free of the transport body type.
173        let (parts, body) = request.into_parts();
174        let probe = Request::from_parts(parts, ());
175        let selected = match router.claim(&probe) {
176            RouteClaim::ClaimedWithRoute(selected) | RouteClaim::DeferredClaim(selected) => Ok(selected),
177            RouteClaim::Claimed => router.route(&probe),
178            RouteClaim::DeferredRejection(error) => Err(error),
179            RouteClaim::NoClaim => {
180                return State::Rejected {
181                    response: Some(unclaimed()),
182                }
183            }
184        };
185        match selected {
186            Ok(selected) => {
187                let (parts, ()) = probe.into_parts();
188                State::Handling {
189                    future: self.handle(selected, 0, Request::from_parts(parts, body)),
190                }
191            }
192            Err(err) => State::Rejected {
193                response: Some(self.reject(0, err)),
194            },
195        }
196    }
197
198    /// Frames a routing error with the rejecting protocol's serialization.
199    fn reject(&self, protocol: usize, error: RoutingError) -> Response<BoxBody> {
200        let protocol = &self.state.protocols[protocol].protocol;
201        tracing::debug!(
202            %error,
203            source = ?std::error::Error::source(&error),
204            protocol = %protocol.protocol_id(),
205            "failed to route"
206        );
207        protocol.serialize_routing_error(&error)
208    }
209}
210
211/// The claim walk remains owned until the response completes, including handler readiness.
212struct ProtocolRouting {
213    service: MultiProtocolRoutingService,
214    probe: Request<()>,
215    body: Option<crate::body::RequestBody>,
216    collected: Option<Bytes>,
217    index: usize,
218    streaming: Option<bool>,
219    deferred: Option<(usize, RoutingError)>,
220    deferred_claim: Option<(usize, OperationTarget)>,
221}
222
223impl ProtocolRouting {
224    /// Continues synchronously until body I/O or handler readiness must be polled.
225    fn advance(&mut self) -> State {
226        while self.index < self.service.state.protocols.len() {
227            let index = self.index;
228            let protocol = &self.service.state.protocols[index];
229            let selected = match &protocol.router {
230                SharedProtocolRouter::Metadata(router) => match router.claim(&self.probe) {
231                    RouteClaim::ClaimedWithRoute(selected) => Some(Ok(selected)),
232                    RouteClaim::DeferredClaim(selected) => {
233                        self.deferred_claim.get_or_insert((index, selected));
234                        None
235                    }
236                    RouteClaim::Claimed => Some(router.route(&self.probe)),
237                    RouteClaim::DeferredRejection(error) => {
238                        self.deferred.get_or_insert((index, error));
239                        None
240                    }
241                    RouteClaim::NoClaim => None,
242                },
243                SharedProtocolRouter::Body(router) => {
244                    let claim = router.claim(&self.probe);
245                    match claim {
246                        BodyRouteClaim::ClaimedWithRoute(selected) => Some(Ok(selected)),
247                        BodyRouteClaim::DeferredClaim(selected) => {
248                            self.deferred_claim.get_or_insert((index, selected));
249                            None
250                        }
251                        BodyRouteClaim::NoClaim => None,
252                        BodyRouteClaim::DeferredRejection(error) => {
253                            self.deferred.get_or_insert((index, error));
254                            None
255                        }
256                        BodyRouteClaim::NeedsBodyToClaim | BodyRouteClaim::Claimed => {
257                            let needs_body_to_claim = matches!(claim, BodyRouteClaim::NeedsBodyToClaim);
258                            if needs_body_to_claim
259                                && *self.streaming.get_or_insert_with(|| {
260                                    self.service.state.metadata_routers.as_ref().is_some_and(|indices| {
261                                        indices.iter().any(|index| {
262                                            let SharedProtocolRouter::Metadata(router) =
263                                                &self.service.state.protocols[*index].router
264                                            else {
265                                                unreachable!("recognizers are metadata routers");
266                                            };
267                                            router.recognizes_streaming_input(&self.probe)
268                                        })
269                                    })
270                                })
271                            {
272                                tracing::trace!(protocol = %protocol.protocol.protocol_id(), "skipping protocol that needs body bytes to claim a request recognized as streaming");
273                                None
274                            } else if self.collected.is_none() {
275                                return State::Collecting {
276                                    future: crate::schema::protocol::collect_request_body_with_trailers(
277                                        self.body.take().expect("body is collected once"),
278                                        &self.service.state.body_collection_config,
279                                    ),
280                                    needs_body_to_claim,
281                                };
282                            } else {
283                                self.select_body(needs_body_to_claim)
284                            }
285                        }
286                    }
287                }
288            };
289            if let Some(selected) = selected {
290                return self.dispatch(selected);
291            }
292            self.index += 1;
293        }
294        if let Some((index, selected)) = self.deferred_claim.take() {
295            self.index = index;
296            return self.dispatch(Ok(selected));
297        }
298        State::Rejected {
299            response: Some(match self.deferred.take() {
300                Some((index, error)) => self.service.reject(index, error),
301                None => unclaimed(),
302            }),
303        }
304    }
305
306    /// Resumes the same body claim after collection, without probing the head again.
307    fn select_body(&mut self, needs_body_to_claim: bool) -> Option<Result<OperationTarget, RoutingError>> {
308        let SharedProtocolRouter::Body(router) = &self.service.state.protocols[self.index].router else {
309            unreachable!("only body routers collect for routing");
310        };
311        let probe = std::mem::replace(&mut self.probe, Request::new(()));
312        let request = probe.map(|()| CollectedBody {
313            bytes: self.collected.as_ref().expect("body was collected").clone(),
314        });
315        let claim = if needs_body_to_claim {
316            router.claim_with_body(&request)
317        } else {
318            RouteClaim::Claimed
319        };
320        let selected = match claim {
321            RouteClaim::ClaimedWithRoute(selected) => Some(Ok(selected)),
322            RouteClaim::DeferredClaim(selected) => {
323                self.deferred_claim.get_or_insert((self.index, selected));
324                None
325            }
326            RouteClaim::Claimed => Some(router.route_with_body(&request)),
327            RouteClaim::NoClaim => None,
328            RouteClaim::DeferredRejection(error) => {
329                self.deferred.get_or_insert((self.index, error));
330                None
331            }
332        };
333        self.probe = request.map(|_| ());
334        selected
335    }
336
337    fn dispatch(&mut self, selected: Result<OperationTarget, RoutingError>) -> State {
338        match selected {
339            Ok(selected) => {
340                let probe = std::mem::replace(&mut self.probe, Request::new(()));
341                State::Handling {
342                    future: self.service.handle(
343                        selected,
344                        self.index,
345                        probe.map(|()| self.body.take().expect("body remains until dispatch")),
346                    ),
347                }
348            }
349            Err(error) => State::Rejected {
350                response: Some(self.service.reject(self.index, error)),
351            },
352        }
353    }
354}
355
356pin_project_lite::pin_project! {
357    #[project = StateProj]
358    enum State {
359        Routing,
360        Collecting {
361            #[pin]
362            future: crate::schema::protocol::RequestBodyCollectionFuture<crate::body::RequestBody>,
363            needs_body_to_claim: bool,
364        },
365        Handling {
366            #[pin]
367            future: crate::routing::route::SyncRouteFuture<crate::body::RequestBody>,
368        },
369        Rejected {
370            response: Option<Response<BoxBody>>,
371        },
372        Done,
373    }
374}
375
376pin_project_lite::pin_project! {
377    /// Response future for schema routing.
378    pub struct MultiProtocolRoutingFuture {
379        #[pin]
380        inner: State,
381        routing: Option<ProtocolRouting>,
382    }
383}
384impl Future for MultiProtocolRoutingFuture {
385    type Output = Result<Response<BoxBody>, Infallible>;
386    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
387        let mut this = self.project();
388        loop {
389            let next = match this.inner.as_mut().project() {
390                StateProj::Routing => this.routing.as_mut().expect("claim walk exists").advance(),
391                StateProj::Collecting {
392                    future,
393                    needs_body_to_claim,
394                } => {
395                    let routing = this.routing.as_mut().expect("claim walk exists");
396                    match std::task::ready!(future.poll(cx)) {
397                        Ok((bytes, trailers)) => {
398                            routing.body = Some(crate::body::RequestBody::buffered(bytes.clone(), trailers));
399                            routing.collected = Some(bytes);
400                            match routing.select_body(*needs_body_to_claim) {
401                                Some(selected) => routing.dispatch(selected),
402                                None => {
403                                    routing.index += 1;
404                                    State::Routing
405                                }
406                            }
407                        }
408                        Err(error) => State::Rejected {
409                            response: Some(crate::schema::body_collection_rejection(
410                                &*routing.service.state.protocols[routing.index].protocol,
411                                error,
412                            )),
413                        },
414                    }
415                }
416                StateProj::Handling { future } => {
417                    let response = std::task::ready!(future.poll(cx));
418                    this.inner.set(State::Done);
419                    *this.routing = None;
420                    return Poll::Ready(response);
421                }
422                StateProj::Rejected { response } => {
423                    let response = response.take().expect("polled after completion");
424                    this.inner.set(State::Done);
425                    *this.routing = None;
426                    return Poll::Ready(Ok(response));
427                }
428                StateProj::Done => panic!("polled after completion"),
429            };
430            this.inner.set(next);
431        }
432    }
433}
434
435impl MultiProtocolRoutingService {
436    /// Starts configuring routing for the supplied service schema.
437    pub fn builder(service: &'static ServiceSchema<'static>) -> MultiProtocolRoutingServiceBuilder {
438        MultiProtocolRoutingServiceBuilder::new(service)
439    }
440}
441/// Any compatible body enters. A hyper incoming body and an already-normalized
442/// [`RequestBody`](crate::body::RequestBody) stay unerased; any other body — tests, adapters,
443/// upgrade layers — is erased into a boxed state (see [`RequestBody::new`](crate::body::RequestBody::new)).
444impl<RB> Service<Request<RB>> for MultiProtocolRoutingService
445where
446    RB: http_body::Body<Data = Bytes> + Send + Sync + 'static,
447    RB::Error: Into<BoxError>,
448{
449    type Response = Response<BoxBody>;
450    type Error = Infallible;
451    type Future = MultiProtocolRoutingFuture;
452    fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll<Result<(), Infallible>> {
453        Poll::Ready(Ok(()))
454    }
455    fn call(&mut self, request: Request<RB>) -> Self::Future {
456        self.route_request(request.map(crate::body::RequestBody::new))
457    }
458}
459
460#[cfg(test)]
461mod tests {
462    use super::super::test_helpers::{binding, FIRST, REST_JSON, SECOND};
463    use super::*;
464    use crate::body::Body;
465    use http::{HeaderValue, StatusCode};
466    use tower::ServiceExt;
467
468    #[tokio::test]
469    async fn normalized_and_boxed_bodies_preserve_shared_routing() {
470        use crate::body::RequestBody;
471        use aws_smithy_schema::shape_id;
472        use http_body_util::BodyExt;
473
474        static MULTI: ServiceSchema<'static> = ServiceSchema::new(
475            shape_id!("test", "Service"),
476            None,
477            &[
478                shape_id!("aws.protocols", "restJson1"),
479                shape_id!("aws.protocols", "restXml"),
480            ],
481            &[&FIRST, &SECOND],
482        );
483
484        fn route() -> SyncRoute<RequestBody> {
485            SyncRoute::new(tower::service_fn(|request: Request<RequestBody>| async {
486                assert!(SelectedOperation::get_from_request(&request).is_some());
487                let bytes = request.into_body().collect().await.unwrap().to_bytes();
488                Ok::<_, Infallible>(Response::new(crate::body::to_boxed(bytes)))
489            }))
490        }
491
492        // Normalized bodies and shared handler services retain their thread-safety guarantees.
493        crate::test_helpers::assert_send::<MultiProtocolRoutingService>();
494        crate::test_helpers::assert_sync::<MultiProtocolRoutingService>();
495        crate::test_helpers::assert_sync::<SyncRoute<RequestBody>>();
496
497        for schema in [&REST_JSON, &MULTI] {
498            let service = MultiProtocolRoutingServiceBuilder::new(schema)
499                .operation_handler_bindings([(&FIRST, route()), (&SECOND, route())])
500                .build()
501                .unwrap();
502            let cloned = service.clone();
503            assert!(Arc::ptr_eq(&service.state, &cloned.state));
504
505            let request = Request::builder()
506                .method("POST")
507                .uri("/first")
508                .header("content-type", "application/json")
509                .body(http_body_util::Full::new(Bytes::from_static(b"payload")))
510                .unwrap();
511            let normalized = Request::builder()
512                .method("POST")
513                .uri("/second")
514                .header("content-type", "application/json")
515                .body(RequestBody::from_bytes(Bytes::from_static(b"normalized")))
516                .unwrap();
517
518            // Exercise public ingress and the buffered normalized body concurrently,
519            // through both the single-protocol shortcut and the async claim walk.
520            let ingress = tokio::spawn(cloned.oneshot(request));
521            let routing = tokio::spawn(service.route_request(normalized));
522            let response = ingress.await.unwrap().unwrap();
523            assert_eq!(response.into_body().collect().await.unwrap().to_bytes(), "payload");
524            let response = routing.await.unwrap().unwrap();
525            assert_eq!(response.into_body().collect().await.unwrap().to_bytes(), "normalized");
526        }
527    }
528
529    #[tokio::test]
530    async fn builder_defers_layers_until_build_and_preserves_stack_order() {
531        use std::sync::atomic::{AtomicUsize, Ordering};
532        let applications = Arc::new(AtomicUsize::new(0));
533        let counter = applications.clone();
534
535        let layer = tower::layer::layer_fn(move |inner: SyncRoute<Body>| {
536            counter.fetch_add(1, Ordering::Release);
537
538            tower::service_fn(move |mut request: Request<Body>| {
539                assert!(SelectedOperation::get_from_request(&request).is_some());
540                assert_eq!(request.headers()["x-layer-order"], "outer");
541                request
542                    .headers_mut()
543                    .insert("x-layer-order", HeaderValue::from_static("inner"));
544                inner.clone().oneshot(request)
545            })
546        });
547
548        let outer = tower::util::MapRequestLayer::new(|mut request: Request<Body>| {
549            assert!(!request.headers().contains_key("x-layer-order"));
550            request
551                .headers_mut()
552                .insert("x-layer-order", HeaderValue::from_static("outer"));
553            request
554        });
555
556        let stack = tower::layer::util::Stack::new(layer, outer);
557        let builder = MultiProtocolRoutingService::builder(&REST_JSON)
558            .operation_handler_bindings([binding(&FIRST), binding(&SECOND)])
559            .layer(stack);
560        assert_eq!(applications.load(Ordering::Acquire), 0);
561
562        let app = builder.build().unwrap();
563        assert_eq!(applications.load(Ordering::Acquire), 2);
564
565        let response = app
566            .oneshot(
567                Request::builder()
568                    .method("POST")
569                    .uri("/first")
570                    .header("content-type", "application/json")
571                    .body(Body::empty())
572                    .unwrap(),
573            )
574            .await
575            .unwrap();
576        assert_eq!(response.status(), StatusCode::OK);
577    }
578}
aws_smithy_http_server/schema/routing/service.rs

aws_smithy_http_server/schema/routing/
service.rs
