aws_smithy_http_server/operation/
upgrade.rs

1/*
2 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3 * SPDX-License-Identifier: Apache-2.0
4 */
5
6use std::{
7    convert::Infallible,
8    future::{Future, Ready},
9    marker::PhantomData,
10    pin::Pin,
11    task::{Context, Poll},
12    time::Duration,
13};
14
15use futures_util::ready;
16use http::{header, StatusCode, Version};
17use pin_project_lite::pin_project;
18use tower::{util::Oneshot, Service, ServiceExt};
19use tracing::{debug, error, trace};
20
21use crate::{
22    body::{BoxBody, __private::HttpBody},
23    plugin::Plugin,
24    request::{FromParts, FromRequest},
25    response::IntoResponse,
26    runtime_error::InternalFailureException,
27    schema::{
28        collect_request_body, BodyDirective, DeserializableShape, DeserializeError, HttpModeledError, SelectedOperation,
29    },
30    service::ServiceShape,
31    settings::RequestBodyReadTimeoutConfig,
32};
33
34use super::{OperationShape, SchemaOperationShape, StreamingOperationShape};
35use aws_smithy_schema::serde::SerializableStruct;
36use aws_smithy_types::body::SdkBody;
37
38/// A [`Plugin`] responsible for taking an operation [`Service`], accepting and returning Smithy
39/// types and converting it into a [`Service`] taking and returning [`http`] types.
40///
41/// See [`Upgrade`].
42#[derive(Debug, Clone)]
43pub struct UpgradePlugin<Extractors> {
44    _extractors: PhantomData<Extractors>,
45    request_body_read_timeout: Option<RequestBodyReadTimeoutConfig>,
46}
47
48/// Schema-driven, protocol-neutral HTTP upgrade plugin for operations without streaming members.
49#[derive(Debug, Clone)]
50pub struct DynUpgradePlugin<Extractors> {
51    _extractors: PhantomData<Extractors>,
52}
53
54impl<Extractors> DynUpgradePlugin<Extractors> {
55    pub fn new() -> Self {
56        Self {
57            _extractors: PhantomData,
58        }
59    }
60}
61
62impl<Extractors> Default for DynUpgradePlugin<Extractors> {
63    fn default() -> Self {
64        Self::new()
65    }
66}
67
68impl<Ser, Op, T, Extractors> Plugin<Ser, Op, T> for DynUpgradePlugin<Extractors>
69where
70    Ser: ServiceShape,
71    Op: SchemaOperationShape,
72{
73    type Output = DynUpgrade<Ser::Protocol, Op, Extractors, T>;
74    fn apply(&self, inner: T) -> Self::Output {
75        DynUpgrade {
76            _protocol: PhantomData,
77            _operation: PhantomData,
78            _extractors: PhantomData,
79            inner,
80        }
81    }
82}
83
84/// Upgrade service for a non-streaming schema operation.
85///
86/// The body is collected under the operation's
87/// [`RequestBodyCollectionConfig`](crate::config::RequestBodyCollectionConfig), carried by
88/// [`SelectedOperation`], when the selected protocol asks for it, the input is read through the erased protocol handle from the request
89/// extensions, and the output or error is serialized through the same handle.
90pub struct DynUpgrade<P, Op, Extractors, S> {
91    /// The service's [`ServiceShape::Protocol`]; request-part extractors resolve against it.
92    _protocol: PhantomData<fn() -> P>,
93    _operation: PhantomData<Op>,
94    _extractors: PhantomData<Extractors>,
95    inner: S,
96}
97
98impl<P, Op, Extractors, S: Clone> Clone for DynUpgrade<P, Op, Extractors, S> {
99    fn clone(&self) -> Self {
100        Self {
101            _protocol: PhantomData,
102            _operation: PhantomData,
103            _extractors: PhantomData,
104            inner: self.inner.clone(),
105        }
106    }
107}
108
109enum SelectedOperationError {
110    Missing,
111    Mismatched(SelectedOperation),
112}
113
114impl SelectedOperationError {
115    fn into_response(self) -> http::Response<BoxBody> {
116        match self {
117            Self::Missing => empty_internal_server_error(),
118            Self::Mismatched(selected) => IntoResponse::<()>::into_response_with_protocol(
119                crate::runtime_error::RuntimeError::from(InternalFailureException),
120                &**selected.protocol(),
121            ),
122        }
123    }
124}
125
126/// Reads the routed operation out of the request extensions and checks it is `Op`.
127fn selected_operation<Op: SchemaOperationShape>(
128    extensions: &http::Extensions,
129) -> Result<SelectedOperation, SelectedOperationError> {
130    let Some(selected) = SelectedOperation::get_from_extensions(extensions).cloned() else {
131        error!(operation = %Op::ID.absolute(), "selected protocol operation missing from request extensions");
132        return Err(SelectedOperationError::Missing);
133    };
134    if !std::ptr::eq(selected.operation(), Op::SCHEMA) {
135        error!(operation = %Op::ID.absolute(), "selected protocol operation is incompatible with the routed operation");
136        return Err(SelectedOperationError::Mismatched(selected));
137    }
138    Ok(selected)
139}
140
141/// Converts the HTTP parts and body into the runtime-api request the protocols read.
142pub(crate) fn convert_request<B>(
143    parts: http::request::Parts,
144    body: B,
145) -> Result<aws_smithy_runtime_api::http::Request<B>, DeserializeError> {
146    aws_smithy_runtime_api::http::Request::try_from(http::Request::from_parts(parts, body))
147        .map_err(|err| DeserializeError::Serde(aws_smithy_schema::serde::SerdeError::custom(err.to_string())))
148}
149
150impl<P, Op, Extractors, B, S> Service<http::Request<B>> for DynUpgrade<P, Op, Extractors, S>
151where
152    Op: SchemaOperationShape,
153    Op::Input: DeserializableShape + Send + 'static,
154    Op::Output: SerializableStruct + Send + 'static,
155    Extractors: FromParts<P> + Send + 'static,
156    <Extractors as FromParts<P>>::Rejection: std::fmt::Display + IntoResponse<P>,
157    B: HttpBody + Send + 'static,
158    B::Data: Send,
159    B::Error: std::error::Error + Send + Sync + 'static,
160    S: Service<(Op::Input, Extractors), Response = Op::Output> + Clone + Send + 'static,
161    S::Error: HttpModeledError,
162    S::Future: Send + 'static,
163{
164    type Response = http::Response<BoxBody>;
165    type Error = Infallible;
166    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
167
168    fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
169        Poll::Ready(Ok(()))
170    }
171
172    fn call(&mut self, req: http::Request<B>) -> Self::Future {
173        let clone = self.inner.clone();
174        let service = std::mem::replace(&mut self.inner, clone);
175        Box::pin(async move {
176            let (mut parts, body) = req.into_parts();
177            let selected = match selected_operation::<Op>(&parts.extensions) {
178                Ok(selected) => selected,
179                Err(error) => return Ok(error.into_response()),
180            };
181            let protocol = selected.protocol();
182            let operation = selected.operation();
183            if operation.input().members().iter().any(|member| member.streaming()) {
184                error!("streaming operation routed through DynUpgrade");
185                return Ok(IntoResponse::<P>::into_response_with_protocol(
186                    crate::runtime_error::RuntimeError::from(InternalFailureException),
187                    &**protocol,
188                ));
189            }
190            let extractors = match Extractors::from_parts(&mut parts) {
191                Ok(value) => value,
192                Err(err) => {
193                    // Typically a missing layer, for example `Extension<State>` without `ExtensionLayer`.
194                    // Match legacy FromRequest: log the rejection with Display, preserving its diagnostic.
195                    error!(error = %err, operation = %operation.shape_id(), protocol = %protocol.protocol_id(),
196                        "additional parameter for the handler function could not be constructed");
197                    return Ok(err.into_response_with_protocol(&**protocol));
198                }
199            };
200            let mut request = match convert_request(parts, body) {
201                Ok(request) => request.map(Some),
202                Err(err) => {
203                    debug!(error = %err, operation = %operation.shape_id(), protocol = %protocol.protocol_id(),
204                        "failed to deserialize request into operation's input");
205                    return Ok(protocol.serialize_rejection(err));
206                }
207            };
208            let body = request.body_mut().take().expect("request body is present");
209            if let Err(err) = protocol.validate_request_headers(operation, request.headers()) {
210                debug!(error = %err, operation = %operation.shape_id(), protocol = %protocol.protocol_id(),
211                    "failed to deserialize request into operation's input");
212                return Ok(protocol.serialize_rejection(err));
213            }
214            let bytes = match protocol.request_body_requirement(operation) {
215                BodyDirective::Collect => match collect_request_body(body, &selected.request_body_config()).await {
216                    Ok(bytes) => bytes,
217                    Err(err) => {
218                        debug!(error = %err, operation = %operation.shape_id(), protocol = %protocol.protocol_id(),
219                            "failed to deserialize request into operation's input");
220                        return Ok(crate::schema::body_collection_rejection(
221                            &**protocol,
222                            err.map_body_error(crate::Error::new),
223                        ));
224                    }
225                },
226                BodyDirective::Skip => bytes::Bytes::new(),
227            };
228            // Request::map requires Fn, so we cannot move captured bytes out of the closure.
229            // Bytes::clone is cheap: it shares the backing buffer without copying the body.
230            let request = request.map(|_| bytes.clone());
231            let input = {
232                let mut deserializer = match protocol.deserialize_request(operation.input(), &request) {
233                    Ok(value) => value,
234                    Err(err) => {
235                        debug!(error = %err, operation = %operation.shape_id(), protocol = %protocol.protocol_id(),
236                            "failed to deserialize request into operation's input");
237                        return Ok(protocol.serialize_rejection(err));
238                    }
239                };
240                match Op::Input::deserialize(&mut *deserializer) {
241                    Ok(value) => value,
242                    Err(err) => {
243                        debug!(error = %err, operation = %operation.shape_id(), protocol = %protocol.protocol_id(),
244                            "failed to deserialize request into operation's input");
245                        return Ok(protocol.serialize_rejection(err));
246                    }
247                }
248            };
249            match service.oneshot((input, extractors)).await {
250                Ok(output) => Ok(protocol.serialize_response(operation.output(), &output)),
251                Err(err) => Ok(protocol.serialize_error(&err)),
252            }
253        })
254    }
255}
256
257/// Schema-driven, protocol-neutral HTTP upgrade plugin for operations with a streaming member.
258///
259/// A streaming payload is handed to the generated [`StreamingOperationShape`] glue as an
260/// [`SdkBody`]. HTTP-bound fields are decoded before the handler receives the live stream.
261/// The live-body conversion needs the body to be
262/// `Sync`, so services with streaming operations run on a `Sync` body such as
263/// [`BoxBodySync`](crate::body::BoxBodySync) or hyper's incoming body.
264#[derive(Debug, Clone)]
265pub struct DynStreamingUpgradePlugin<Extractors> {
266    _extractors: PhantomData<Extractors>,
267}
268
269impl<Extractors> DynStreamingUpgradePlugin<Extractors> {
270    pub fn new() -> Self {
271        Self {
272            _extractors: PhantomData,
273        }
274    }
275}
276
277impl<Extractors> Default for DynStreamingUpgradePlugin<Extractors> {
278    fn default() -> Self {
279        Self::new()
280    }
281}
282
283impl<Ser, Op, T, Extractors> Plugin<Ser, Op, T> for DynStreamingUpgradePlugin<Extractors>
284where
285    Ser: ServiceShape,
286    Op: StreamingOperationShape,
287{
288    type Output = DynStreamingUpgrade<Ser::Protocol, Op, Extractors, T>;
289    fn apply(&self, inner: T) -> Self::Output {
290        DynStreamingUpgrade {
291            _protocol: PhantomData,
292            _operation: PhantomData,
293            _extractors: PhantomData,
294            inner,
295        }
296    }
297}
298
299/// Upgrade service for a schema operation with a streaming input or output.
300///
301/// A streaming input reaches the protocol with an empty [`Request`](aws_smithy_runtime_api::http::Request) body, so the protocol
302/// reads URI and header bindings only; the live body goes to
303/// [`StreamingOperationShape::deserialize_streaming_input`]. A non-streaming input on such an
304/// operation is collected exactly as [`DynUpgrade`] collects it.
305pub struct DynStreamingUpgrade<P, Op, Extractors, S> {
306    /// The service's [`ServiceShape::Protocol`]; request-part extractors resolve against it.
307    _protocol: PhantomData<fn() -> P>,
308    _operation: PhantomData<Op>,
309    _extractors: PhantomData<Extractors>,
310    inner: S,
311}
312
313impl<P, Op, Extractors, S: Clone> Clone for DynStreamingUpgrade<P, Op, Extractors, S> {
314    fn clone(&self) -> Self {
315        Self {
316            _protocol: PhantomData,
317            _operation: PhantomData,
318            _extractors: PhantomData,
319            inner: self.inner.clone(),
320        }
321    }
322}
323
324impl<P, Op, Extractors, B, S> Service<http::Request<B>> for DynStreamingUpgrade<P, Op, Extractors, S>
325where
326    Op: StreamingOperationShape,
327    Op::Input: Send + 'static,
328    Op::Output: Send + 'static,
329    Extractors: FromParts<P> + Send + 'static,
330    <Extractors as FromParts<P>>::Rejection: std::fmt::Display + IntoResponse<P>,
331    B: HttpBody<Data = bytes::Bytes> + Send + Sync + 'static,
332    B::Error: std::error::Error + Send + Sync + 'static,
333    S: Service<(Op::Input, Extractors), Response = Op::Output> + Clone + Send + 'static,
334    S::Error: HttpModeledError,
335    S::Future: Send + 'static,
336{
337    type Response = http::Response<BoxBody>;
338    type Error = Infallible;
339    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
340
341    fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
342        Poll::Ready(Ok(()))
343    }
344
345    fn call(&mut self, req: http::Request<B>) -> Self::Future {
346        let clone = self.inner.clone();
347        let service = std::mem::replace(&mut self.inner, clone);
348        Box::pin(async move {
349            let (mut parts, body) = req.into_parts();
350            let selected = match selected_operation::<Op>(&parts.extensions) {
351                Ok(selected) => selected,
352                Err(error) => return Ok(error.into_response()),
353            };
354            let protocol = selected.protocol();
355            let operation = selected.operation();
356            let has_event_stream = [operation.input(), operation.output()].iter().any(|schema| {
357                schema
358                    .members()
359                    .iter()
360                    .any(|member| member.streaming() && member.shape_type() == aws_smithy_schema::ShapeType::Union)
361            });
362            if has_event_stream && protocol.event_stream_framing().is_none() {
363                trace!(operation = %operation.shape_id(), protocol = %protocol.protocol_id(),
364                    "selected protocol does not support event streams");
365                return Ok(empty_not_implemented());
366            }
367
368            let extractors = match Extractors::from_parts(&mut parts) {
369                Ok(value) => value,
370                Err(err) => {
371                    // Typically a missing layer, for example `Extension<State>` without `ExtensionLayer`.
372                    // Match legacy FromRequest: log the rejection with Display, preserving its diagnostic.
373                    error!(error = %err, operation = %operation.shape_id(), protocol = %protocol.protocol_id(),
374                        "additional parameter for the handler function could not be constructed");
375                    return Ok(err.into_response_with_protocol(&**protocol));
376                }
377            };
378            let request = match convert_request(parts, crate::body::RequestBody::new(body)) {
379                Ok(request) => request,
380                Err(err) => {
381                    debug!(error = %err, operation = %operation.shape_id(), protocol = %protocol.protocol_id(),
382                        "failed to deserialize request into operation's input");
383                    return Ok(protocol.serialize_rejection(err));
384                }
385            };
386            let future = Op::deserialize_streaming_input(request, protocol.clone());
387            let input = match future.await {
388                Ok(value) => value,
389                Err(err) => {
390                    debug!(error = %err, operation = %operation.shape_id(), protocol = %protocol.protocol_id(),
391                        "failed to deserialize request into operation's input");
392                    return Ok(protocol.serialize_rejection(err));
393                }
394            };
395            match service.oneshot((input, extractors)).await {
396                Ok(output) => Ok(Op::serialize_streaming_output(output, protocol)),
397                Err(err) => Ok(protocol.serialize_error(&err)),
398            }
399        })
400    }
401}
402
403/// Prepares collected or streaming inputs from an owned request.
404///
405/// The synchronous callback consumes the borrowed HTTP deserializer before its returned future
406/// is awaited. The live stream and decoded bindings can therefore outlive request preparation.
407#[doc(hidden)]
408pub fn deserialize_streaming_request<I, F>(
409    operation: &'static crate::schema::OperationSchema<'static>,
410    request: aws_smithy_runtime_api::http::Request<crate::body::RequestBody>,
411    protocol: crate::schema::SharedServerProtocol,
412    deserialize: F,
413) -> super::StreamingInputFuture<I>
414where
415    I: Send + 'static,
416    F: FnOnce(
417            &mut dyn aws_smithy_schema::serde::ShapeDeserializer,
418            SdkBody,
419            crate::schema::SharedServerProtocol,
420        ) -> super::StreamingInputFuture<I>
421        + Send
422        + 'static,
423{
424    Box::pin(async move {
425        protocol.validate_request_headers(operation, request.headers())?;
426        let config = SelectedOperation::get_from_runtime_request(&request).map(SelectedOperation::request_body_config);
427        let mut request = request.map(Some);
428        let body = request.body_mut().take().expect("request body is present");
429        let directive = protocol.request_body_requirement(operation);
430        let input_streams = operation.input().members().iter().any(|member| {
431            member.streaming()
432                && (member.shape_type() == aws_smithy_schema::ShapeType::Union || directive != BodyDirective::Collect)
433        });
434        let (bytes, body) = if input_streams {
435            let body = crate::body::__private::wrap_streaming_body(body, config.unwrap_or_default().read_poll_budget);
436            (bytes::Bytes::new(), SdkBody::from_body_1_x(body))
437        } else if directive == BodyDirective::Collect {
438            // TODO(incremental CBOR blob decoding): Support streaming once the handler contract for
439            // trailing fields and validation is settled. Collection preserves generated serde timing.
440            let config = config.ok_or_else(|| {
441                DeserializeError::InternalFailure(crate::Error::new(std::io::Error::other(
442                    "selected operation missing while preparing collected input",
443                )))
444            })?;
445            let bytes = collect_request_body(body, &config)
446                .await
447                .map_err(DeserializeError::from)?;
448            (bytes, SdkBody::empty())
449        } else {
450            (bytes::Bytes::new(), SdkBody::empty())
451        };
452        // Request::map requires Fn, so we cannot move captured bytes out of the closure.
453        // Bytes::clone is cheap: it shares the backing buffer without copying the body.
454        let request = request.map(|_| bytes.clone());
455        let future = {
456            let mut deserializer = protocol.deserialize_request(operation.input(), &request)?;
457            deserialize(&mut *deserializer, body, protocol.clone())
458        };
459        future.await
460    })
461}
462
463/// An empty 501 response for functionality unsupported by the selected protocol.
464pub(crate) fn empty_not_implemented() -> http::Response<BoxBody> {
465    let mut response = http::Response::new(crate::body::empty());
466    *response.status_mut() = http::StatusCode::NOT_IMPLEMENTED;
467    response
468}
469
470/// An empty `500 Internal Server Error`: the answer when generated glue cannot even build a
471/// response, such as a failure to serialize an `initial-response` frame.
472#[doc(hidden)]
473pub fn empty_internal_server_error() -> http::Response<BoxBody> {
474    let mut response = http::Response::new(crate::body::empty());
475    *response.status_mut() = http::StatusCode::INTERNAL_SERVER_ERROR;
476    response
477}
478
479impl<Extractors> Default for UpgradePlugin<Extractors> {
480    fn default() -> Self {
481        Self {
482            _extractors: PhantomData,
483            request_body_read_timeout: None,
484        }
485    }
486}
487
488impl<Extractors> UpgradePlugin<Extractors> {
489    /// Creates a new [`UpgradePlugin`].
490    pub fn new() -> Self {
491        Self::default()
492    }
493
494    /// Selects the runtime body-read default appropriate to the modeled operation.
495    #[doc(hidden)]
496    pub fn with_default_request_body_read_timeout(mut self, payload: bool, operation: &'static str) -> Self {
497        self.request_body_read_timeout = RequestBodyReadTimeoutConfig::default_for_operation(payload, operation);
498        self
499    }
500
501    /// Applies a deadline while the request body is read and the operation input is constructed.
502    ///
503    /// The deadline is removed before the operation handler is invoked. Zero disables it.
504    #[doc(hidden)]
505    pub fn with_request_body_read_timeout(mut self, timeout: Duration, operation: &'static str) -> Self {
506        self.request_body_read_timeout = RequestBodyReadTimeoutConfig::new(timeout, operation);
507        self
508    }
509}
510
511impl<Ser, Op, T, Extractors> Plugin<Ser, Op, T> for UpgradePlugin<Extractors>
512where
513    Ser: ServiceShape,
514    Op: OperationShape,
515{
516    type Output = Upgrade<Ser::Protocol, (Op::Input, Extractors), T>;
517
518    fn apply(&self, inner: T) -> Self::Output {
519        Upgrade {
520            _protocol: PhantomData,
521            _input: PhantomData,
522            inner,
523            request_body_read_timeout: self.request_body_read_timeout,
524        }
525    }
526}
527
528/// A [`Service`] responsible for wrapping an operation [`Service`] accepting and returning Smithy
529/// types, and converting it into a [`Service`] accepting and returning [`http`] types.
530pub struct Upgrade<Protocol, Input, S> {
531    _protocol: PhantomData<Protocol>,
532    _input: PhantomData<Input>,
533    inner: S,
534    request_body_read_timeout: Option<RequestBodyReadTimeoutConfig>,
535}
536
537impl<P, Input, S> Clone for Upgrade<P, Input, S>
538where
539    S: Clone,
540{
541    fn clone(&self) -> Self {
542        Self {
543            _protocol: PhantomData,
544            _input: PhantomData,
545            inner: self.inner.clone(),
546            request_body_read_timeout: self.request_body_read_timeout,
547        }
548    }
549}
550
551struct RequestBodyReadTimeoutState {
552    // Created on the first pending poll rather than per request, so requests whose
553    // input is ready immediately never need a timer. This keeps runtimes without a
554    // tokio time driver working and avoids a per-request allocation.
555    sleep: Option<Pin<Box<tokio::time::Sleep>>>,
556    timeout: Duration,
557    version: Version,
558    operation: &'static str,
559    #[cfg(feature = "request-id")]
560    request_id: Option<String>,
561}
562
563impl RequestBodyReadTimeoutState {
564    fn is_expired(&mut self, cx: &mut Context<'_>) -> bool {
565        self.sleep.as_mut().is_some_and(|sleep| {
566            // The timer driver may not have delivered expiry yet, so also check
567            // the clock. Immediately ready inputs never allocate a timer.
568            tokio::time::Instant::now() >= sleep.deadline() || sleep.as_mut().poll(cx).is_ready()
569        })
570    }
571
572    fn response(&self) -> http::Response<crate::body::BoxBody> {
573        #[cfg(feature = "request-id")]
574        if let Some(request_id) = self.request_id.as_deref() {
575            tracing::debug!(
576                operation = self.operation,
577                request_id,
578                timeout_millis = self.timeout.as_millis(),
579                "request body read timed out"
580            );
581        } else {
582            tracing::debug!(
583                operation = self.operation,
584                timeout_millis = self.timeout.as_millis(),
585                "request body read timed out"
586            );
587        }
588        #[cfg(not(feature = "request-id"))]
589        tracing::debug!(
590            operation = self.operation,
591            timeout_millis = self.timeout.as_millis(),
592            "request body read timed out"
593        );
594
595        let mut response = http::Response::builder().status(StatusCode::REQUEST_TIMEOUT);
596        if matches!(self.version, Version::HTTP_10 | Version::HTTP_11) {
597            response = response.header(header::CONNECTION, "close");
598        }
599        response.body(crate::body::empty()).expect("valid response")
600    }
601}
602
603#[cfg(feature = "request-id")]
604fn request_id<B>(req: &http::Request<B>) -> Option<String> {
605    req.extensions()
606        .get::<crate::request::request_id::ServerRequestId>()
607        .map(ToString::to_string)
608}
609
610pin_project! {
611    #[project = InnerProj]
612    #[project_replace = InnerProjReplace]
613    enum Inner<FromFut, HandlerFut> {
614        FromRequest {
615            #[pin]
616            inner: FromFut
617        },
618        Inner {
619            #[pin]
620            call: HandlerFut
621        }
622    }
623}
624
625type InnerAlias<Input, Protocol, B, S> = Inner<<Input as FromRequest<Protocol, B>>::Future, Oneshot<S, Input>>;
626
627pin_project! {
628    /// The [`Service::Future`] of [`Upgrade`].
629    pub struct UpgradeFuture<Protocol, Input, B, S>
630    where
631        Input: FromRequest<Protocol, B>,
632        S: Service<Input>,
633    {
634        service: Option<S>,
635        #[pin]
636        inner: InnerAlias<Input, Protocol, B, S>,
637        request_body_read_timeout: Option<RequestBodyReadTimeoutState>,
638    }
639}
640
641impl<P, Input, B, S> Future for UpgradeFuture<P, Input, B, S>
642where
643    Input: FromRequest<P, B>,
644    <Input as FromRequest<P, B>>::Rejection: std::fmt::Display,
645    S: Service<Input>,
646    S::Response: IntoResponse<P>,
647    S::Error: IntoResponse<P>,
648{
649    type Output = Result<http::Response<crate::body::BoxBody>, Infallible>;
650
651    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
652        loop {
653            let mut this = self.as_mut().project();
654            let this2 = this.inner.as_mut().project();
655
656            let call = match this2 {
657                InnerProj::FromRequest { inner } => {
658                    if let Some(timeout) = this.request_body_read_timeout {
659                        if timeout.is_expired(cx) {
660                            return Poll::Ready(Ok(timeout.response()));
661                        }
662                    }
663                    let result = inner.poll(cx);
664                    if let Some(timeout) = this.request_body_read_timeout {
665                        if result.is_pending() && timeout.sleep.is_none() {
666                            timeout.sleep = Some(Box::pin(tokio::time::sleep(timeout.timeout)));
667                        }
668                        // Input may become ready on the same poll that crosses an
669                        // armed deadline. Expiry must win before calling the handler.
670                        if timeout.is_expired(cx) {
671                            return Poll::Ready(Ok(timeout.response()));
672                        }
673                    }
674                    let result = match result {
675                        Poll::Ready(result) => result,
676                        Poll::Pending => return Poll::Pending,
677                    };
678                    *this.request_body_read_timeout = None;
679                    match result {
680                        Ok(ok) => this
681                            .service
682                            .take()
683                            .expect("futures cannot be polled after completion")
684                            .oneshot(ok),
685                        Err(err) => {
686                            // The error may arise either from a `FromRequest` failure for any user-defined
687                            // handler's additional input parameters, or from a de-serialization failure
688                            // of an input parameter specific to the operation.
689                            tracing::trace!(error = %err, "parameter for the handler cannot be constructed");
690                            return Poll::Ready(Ok(err.into_response()));
691                        }
692                    }
693                }
694                InnerProj::Inner { call } => {
695                    let result = ready!(call.poll(cx));
696                    let output = match result {
697                        Ok(ok) => ok.into_response(),
698                        Err(err) => err.into_response(),
699                    };
700                    return Poll::Ready(Ok(output));
701                }
702            };
703
704            this.inner.as_mut().project_replace(Inner::Inner { call });
705        }
706    }
707}
708
709impl<P, Input, B, S> Service<http::Request<B>> for Upgrade<P, Input, S>
710where
711    Input: FromRequest<P, B>,
712    <Input as FromRequest<P, B>>::Rejection: std::fmt::Display,
713    S: Service<Input> + Clone,
714    S::Response: IntoResponse<P>,
715    S::Error: IntoResponse<P>,
716{
717    type Response = http::Response<crate::body::BoxBody>;
718    type Error = Infallible;
719    type Future = UpgradeFuture<P, Input, B, S>;
720
721    fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
722        // The check that the inner service is ready is done by `Oneshot` in `UpgradeFuture`'s
723        // implementation.
724        Poll::Ready(Ok(()))
725    }
726
727    fn call(&mut self, req: http::Request<B>) -> Self::Future {
728        let clone = self.inner.clone();
729        let service = std::mem::replace(&mut self.inner, clone);
730        let request_body_read_timeout = self.request_body_read_timeout.map(|config| {
731            debug_assert!(
732                !config.timeout.is_zero(),
733                "disabled request body read timeouts must be represented as None"
734            );
735            RequestBodyReadTimeoutState {
736                sleep: None,
737                timeout: config.timeout,
738                version: req.version(),
739                operation: config.operation,
740                #[cfg(feature = "request-id")]
741                request_id: request_id(&req),
742            }
743        });
744        UpgradeFuture {
745            service: Some(service),
746            inner: Inner::FromRequest {
747                inner: <Input as FromRequest<P, B>>::from_request(req),
748            },
749            request_body_read_timeout,
750        }
751    }
752}
753
754/// A [`Service`] which always returns an internal failure message and logs an error.
755#[derive(Copy)]
756pub struct MissingFailure<P> {
757    _protocol: PhantomData<fn(P)>,
758}
759
760impl<P> Default for MissingFailure<P> {
761    fn default() -> Self {
762        Self { _protocol: PhantomData }
763    }
764}
765
766impl<P> Clone for MissingFailure<P> {
767    fn clone(&self) -> Self {
768        MissingFailure { _protocol: PhantomData }
769    }
770}
771
772impl<R, P> Service<R> for MissingFailure<P>
773where
774    InternalFailureException: IntoResponse<P>,
775{
776    type Response = http::Response<BoxBody>;
777    type Error = Infallible;
778    type Future = Ready<Result<Self::Response, Self::Error>>;
779
780    fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
781        Poll::Ready(Ok(()))
782    }
783
784    fn call(&mut self, _request: R) -> Self::Future {
785        error!("the operation has not been set");
786        std::future::ready(Ok(InternalFailureException.into_response()))
787    }
788}
789
790#[cfg(test)]
791mod tests;
792
793/// Missing-handler fallback using the protocol selected by schema routing.
794///
795/// Returns an empty internal server error if middleware removed the selected protocol context.
796#[derive(Clone, Copy, Debug, Default)]
797pub struct SchemaMissingFailure;
798impl Service<http::Request<crate::body::Body>> for SchemaMissingFailure {
799    type Response = http::Response<BoxBody>;
800    type Error = Infallible;
801    type Future = Ready<Result<Self::Response, Self::Error>>;
802    fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll<Result<(), Infallible>> {
803        Poll::Ready(Ok(()))
804    }
805    fn call(&mut self, request: http::Request<crate::body::Body>) -> Self::Future {
806        error!("the operation has not been set");
807        let Some(selected) = SelectedOperation::get_from_request(&request) else {
808            error!("request reached the missing-handler fallback without a selected protocol operation");
809            return std::future::ready(Ok(empty_internal_server_error()));
810        };
811        let rejection = crate::schema::DeserializeError::InternalFailure(crate::Error::new(String::from(
812            "the operation has not been set",
813        )));
814        std::future::ready(Ok(selected.protocol().serialize_rejection(rejection)))
815    }
816}
817
818#[cfg(test)]
819mod timeout_tests {
820    use super::*;
821    use std::sync::{
822        atomic::{AtomicUsize, Ordering},
823        Arc,
824    };
825    use std::task::Waker;
826    use tokio::sync::oneshot;
827
828    struct TestProtocol;
829    struct TestInput;
830    type InputResult = Result<TestInput, StatusCode>;
831    type ResponseFuture = Pin<Box<dyn Future<Output = Result<http::Response<BoxBody>, Infallible>>>>;
832
833    impl IntoResponse<TestProtocol> for StatusCode {
834        fn into_response(self) -> http::Response<BoxBody> {
835            http::Response::builder()
836                .status(self)
837                .body(crate::body::empty())
838                .unwrap()
839        }
840    }
841
842    impl FromRequest<TestProtocol, oneshot::Receiver<InputResult>> for TestInput {
843        type Rejection = StatusCode;
844        type Future = Pin<Box<dyn Future<Output = InputResult> + Send>>;
845
846        fn from_request(request: http::Request<oneshot::Receiver<InputResult>>) -> Self::Future {
847            Box::pin(async { request.into_body().await.expect("input sender must remain alive") })
848        }
849    }
850
851    fn request(
852        timeout: Option<Duration>,
853        version: Version,
854        handler_delay: Option<Duration>,
855    ) -> (ResponseFuture, oneshot::Sender<InputResult>, Arc<AtomicUsize>) {
856        let calls = Arc::new(AtomicUsize::new(0));
857        let handler_calls = calls.clone();
858        let mut upgrade = Upgrade::<TestProtocol, TestInput, _> {
859            _protocol: PhantomData,
860            _input: PhantomData,
861            inner: tower::service_fn(move |_input| {
862                handler_calls.fetch_add(1, Ordering::SeqCst);
863                async move {
864                    if let Some(delay) = handler_delay {
865                        tokio::time::sleep(delay).await;
866                    }
867                    Ok::<_, Infallible>(StatusCode::OK)
868                }
869            }),
870            request_body_read_timeout: timeout.map(|timeout| RequestBodyReadTimeoutConfig {
871                timeout,
872                operation: "test#Operation",
873            }),
874        };
875        let (sender, receiver) = oneshot::channel();
876        let request = http::Request::builder().version(version).body(receiver).unwrap();
877        (Box::pin(upgrade.call(request)), sender, calls)
878    }
879
880    fn assert_pending(future: &mut ResponseFuture) {
881        let mut cx = Context::from_waker(Waker::noop());
882        assert!(future.as_mut().poll(&mut cx).is_pending());
883    }
884
885    #[tokio::test(start_paused = true)]
886    async fn expired_deadline_rejects_ready_input() {
887        for version in [Version::HTTP_11, Version::HTTP_2] {
888            for succeeds in [true, false] {
889                let (mut response, sender, calls) = request(Some(Duration::from_millis(10)), version, None);
890                assert_pending(&mut response);
891                tokio::time::advance(Duration::from_millis(10)).await;
892                // The input is ready before the next poll, so expiry must win over
893                // both successful deserialization and an input rejection.
894                assert!(sender
895                    .send(if succeeds {
896                        Ok(TestInput)
897                    } else {
898                        Err(StatusCode::BAD_REQUEST)
899                    })
900                    .is_ok());
901                let response = response.await.unwrap();
902                assert_eq!(response.status(), StatusCode::REQUEST_TIMEOUT);
903                assert_eq!(calls.load(Ordering::SeqCst), 0);
904                if version == Version::HTTP_11 {
905                    assert_eq!(response.headers()[header::CONNECTION], "close");
906                } else {
907                    assert!(!response.headers().contains_key(header::CONNECTION));
908                }
909            }
910        }
911    }
912
913    #[tokio::test(start_paused = true)]
914    async fn input_ready_before_deadline_reaches_handler() {
915        let (mut response, sender, calls) = request(Some(Duration::from_millis(10)), Version::HTTP_11, None);
916        assert_pending(&mut response);
917        tokio::time::advance(Duration::from_millis(9)).await;
918        assert!(sender.send(Ok(TestInput)).is_ok());
919        assert_eq!(response.await.unwrap().status(), StatusCode::OK);
920        assert_eq!(calls.load(Ordering::SeqCst), 1);
921    }
922
923    #[tokio::test(start_paused = true)]
924    async fn body_deadline_does_not_time_out_handler() {
925        let (mut response, sender, calls) = request(
926            Some(Duration::from_millis(10)),
927            Version::HTTP_11,
928            Some(Duration::from_millis(20)),
929        );
930        assert_pending(&mut response);
931        tokio::time::advance(Duration::from_millis(9)).await;
932        assert!(sender.send(Ok(TestInput)).is_ok());
933        assert_pending(&mut response);
934        tokio::time::advance(Duration::from_millis(100)).await;
935        assert_eq!(response.await.unwrap().status(), StatusCode::OK);
936        assert_eq!(calls.load(Ordering::SeqCst), 1);
937    }
938
939    #[tokio::test(start_paused = true)]
940    async fn disabled_deadline_allows_late_input() {
941        let (mut response, sender, calls) = request(None, Version::HTTP_11, None);
942        assert_pending(&mut response);
943        tokio::time::advance(Duration::from_secs(60)).await;
944        assert!(sender.send(Ok(TestInput)).is_ok());
945        assert_eq!(response.await.unwrap().status(), StatusCode::OK);
946        assert_eq!(calls.load(Ordering::SeqCst), 1);
947    }
948
949    #[test]
950    fn immediately_ready_input_does_not_require_timer_driver() {
951        let (mut response, sender, calls) = request(Some(Duration::from_millis(10)), Version::HTTP_11, None);
952        assert!(sender.send(Ok(TestInput)).is_ok());
953        let mut cx = Context::from_waker(Waker::noop());
954        match response.as_mut().poll(&mut cx) {
955            Poll::Ready(Ok(response)) => assert_eq!(response.status(), StatusCode::OK),
956            _ => panic!("immediately ready input should complete without a timer driver"),
957        }
958        assert_eq!(calls.load(Ordering::SeqCst), 1);
959    }
960}
aws_smithy_http_server/operation/upgrade.rs

aws_smithy_http_server/operation/
upgrade.rs
