1use std::{
2 collections::VecDeque,
3 fmt,
4 future::Future,
5 io, mem, net,
6 pin::Pin,
7 rc::Rc,
8 task::{Context, Poll},
9};
10
11use actix_codec::{Framed, FramedParts};
12use actix_rt::time::sleep_until;
13use actix_service::Service;
14use bitflags::bitflags;
15use bytes::{Buf, BytesMut};
16use futures_core::ready;
17use pin_project_lite::pin_project;
18use tokio::io::{AsyncRead, AsyncWrite};
19use tokio_util::codec::{Decoder as _, Encoder as _};
20use tracing::{error, trace};
21
22use super::{
23 codec::Codec,
24 decoder::MAX_BUFFER_SIZE,
25 payload::{Payload, PayloadSender, PayloadStatus},
26 timer::TimerState,
27 Message, MessageType,
28};
29use crate::{
30 body::{BodySize, BoxBody, MessageBody},
31 config::ServiceConfig,
32 error::{DispatchError, ParseError, PayloadError},
33 service::HttpFlow,
34 ConnectionType, Error, Extensions, HttpMessage, OnConnectData, Request, Response, StatusCode,
35};
36
37const LW_BUFFER_SIZE: usize = 1024;
38const HW_BUFFER_SIZE: usize = 1024 * 8;
39const MAX_PIPELINED_MESSAGES: usize = 16;
40
41bitflags! {
42 #[derive(Debug, Clone, Copy)]
43 pub struct Flags: u8 {
44 const STARTED = 0b0000_0001;
46
47 const FINISHED = 0b0000_0010;
49
50 const KEEP_ALIVE = 0b0000_0100;
52
53 const SHUTDOWN = 0b0000_1000;
55
56 const READ_DISCONNECT = 0b0001_0000;
58
59 const WRITE_DISCONNECT = 0b0010_0000;
61
62 const LINGER = 0b0100_0000;
64
65 const DRAINING = 0b1000_0000;
72 }
73}
74
75#[cfg(not(test))]
81pin_project! {
82 pub struct Dispatcher<T, S, B, X, U>
84 where
85 S: Service<Request>,
86 S::Error: Into<Response<BoxBody>>,
87
88 B: MessageBody,
89
90 X: Service<Request, Response = Request>,
91 X::Error: Into<Response<BoxBody>>,
92
93 U: Service<(Request, Framed<T, Codec>), Response = ()>,
94 U::Error: fmt::Display,
95 {
96 #[pin]
97 inner: DispatcherState<T, S, B, X, U>,
98 }
99}
100
101#[cfg(test)]
102pin_project! {
103 pub struct Dispatcher<T, S, B, X, U>
105 where
106 S: Service<Request>,
107 S::Error: Into<Response<BoxBody>>,
108
109 B: MessageBody,
110
111 X: Service<Request, Response = Request>,
112 X::Error: Into<Response<BoxBody>>,
113
114 U: Service<(Request, Framed<T, Codec>), Response = ()>,
115 U::Error: fmt::Display,
116 {
117 #[pin]
118 pub(super) inner: DispatcherState<T, S, B, X, U>,
119
120 pub(super) poll_count: u64,
122 }
123}
124
125pin_project! {
126 #[project = DispatcherStateProj]
127 pub(super) enum DispatcherState<T, S, B, X, U>
128 where
129 S: Service<Request>,
130 S::Error: Into<Response<BoxBody>>,
131
132 B: MessageBody,
133
134 X: Service<Request, Response = Request>,
135 X::Error: Into<Response<BoxBody>>,
136
137 U: Service<(Request, Framed<T, Codec>), Response = ()>,
138 U::Error: fmt::Display,
139 {
140 Normal { #[pin] inner: InnerDispatcher<T, S, B, X, U> },
141 Upgrade { #[pin] fut: U::Future },
142 }
143}
144
145pin_project! {
146 #[project = InnerDispatcherProj]
147 pub(super) struct InnerDispatcher<T, S, B, X, U>
148 where
149 S: Service<Request>,
150 S::Error: Into<Response<BoxBody>>,
151
152 B: MessageBody,
153
154 X: Service<Request, Response = Request>,
155 X::Error: Into<Response<BoxBody>>,
156
157 U: Service<(Request, Framed<T, Codec>), Response = ()>,
158 U::Error: fmt::Display,
159 {
160 flow: Rc<HttpFlow<S, X, U>>,
161 pub(super) flags: Flags,
162 peer_addr: Option<net::SocketAddr>,
163 conn_data: Option<Rc<Extensions>>,
164 config: ServiceConfig,
165 error: Option<DispatchError>,
166
167 #[pin]
168 pub(super) state: State<S, B, X>,
169 payload: Option<PayloadSender>,
171 payload_drainable: bool,
173 messages: VecDeque<DispatcherMessage>,
174
175 head_timer: TimerState,
176 ka_timer: TimerState,
177 shutdown_timer: TimerState,
178 graceful_shutdown: Option<crate::config::GracefulShutdownFuture>,
179
180 pub(super) io: Option<T>,
181 read_buf: BytesMut,
182 write_buf: BytesMut,
183 h1_write_buffer_size: usize,
184 codec: Codec,
185 }
186}
187
188enum DispatcherMessage {
189 Item(Request),
190 Upgrade(Request),
191 Error(Response<()>),
192}
193
194pin_project! {
195 #[project = StateProj]
196 pub(super) enum State<S, B, X>
197 where
198 S: Service<Request>,
199 X: Service<Request, Response = Request>,
200 B: MessageBody,
201 {
202 None,
203 ExpectCall { #[pin] fut: X::Future },
204 ServiceCall { #[pin] fut: S::Future },
205 SendPayload { #[pin] body: B },
206 SendErrorPayload { #[pin] body: BoxBody },
207 }
208}
209
210impl<S, B, X> State<S, B, X>
211where
212 S: Service<Request>,
213 X: Service<Request, Response = Request>,
214 B: MessageBody,
215{
216 pub(super) fn is_none(&self) -> bool {
217 matches!(self, State::None)
218 }
219}
220
221impl<S, B, X> fmt::Debug for State<S, B, X>
222where
223 S: Service<Request>,
224 X: Service<Request, Response = Request>,
225 B: MessageBody,
226{
227 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
228 match self {
229 Self::None => write!(f, "State::None"),
230 Self::ExpectCall { .. } => f.debug_struct("State::ExpectCall").finish_non_exhaustive(),
231 Self::ServiceCall { .. } => {
232 f.debug_struct("State::ServiceCall").finish_non_exhaustive()
233 }
234 Self::SendPayload { .. } => {
235 f.debug_struct("State::SendPayload").finish_non_exhaustive()
236 }
237 Self::SendErrorPayload { .. } => f
238 .debug_struct("State::SendErrorPayload")
239 .finish_non_exhaustive(),
240 }
241 }
242}
243
244#[derive(Debug)]
245enum PollResponse {
246 Upgrade(Request),
247 DoNothing,
248 DrainWriteBuf,
249}
250
251impl<T, S, B, X, U> Dispatcher<T, S, B, X, U>
252where
253 T: AsyncRead + AsyncWrite + Unpin,
254
255 S: Service<Request>,
256 S::Error: Into<Response<BoxBody>>,
257 S::Response: Into<Response<B>>,
258
259 B: MessageBody,
260
261 X: Service<Request, Response = Request>,
262 X::Error: Into<Response<BoxBody>>,
263
264 U: Service<(Request, Framed<T, Codec>), Response = ()>,
265 U::Error: fmt::Display,
266{
267 pub(crate) fn new(
269 io: T,
270 flow: Rc<HttpFlow<S, X, U>>,
271 config: ServiceConfig,
272 peer_addr: Option<net::SocketAddr>,
273 conn_data: OnConnectData,
274 ) -> Self {
275 Dispatcher {
276 inner: DispatcherState::Normal {
277 inner: InnerDispatcher {
278 flow,
279 flags: Flags::empty(),
280 peer_addr,
281 conn_data: conn_data.0.map(Rc::new),
282 config: config.clone(),
283 error: None,
284
285 state: State::None,
286 payload: None,
287 payload_drainable: false,
288 messages: VecDeque::new(),
289
290 head_timer: TimerState::new(config.client_request_deadline().is_some()),
291 ka_timer: TimerState::new(config.keep_alive().enabled()),
292 shutdown_timer: TimerState::new(config.client_disconnect_deadline().is_some()),
293 graceful_shutdown: config.graceful_shutdown(),
294
295 io: Some(io),
296 read_buf: BytesMut::with_capacity(HW_BUFFER_SIZE),
297 write_buf: BytesMut::with_capacity(HW_BUFFER_SIZE),
298 h1_write_buffer_size: config.h1_write_buffer_size(),
299 codec: Codec::new(config),
300 },
301 },
302
303 #[cfg(test)]
304 poll_count: 0,
305 }
306 }
307}
308
309impl<T, S, B, X, U> InnerDispatcher<T, S, B, X, U>
310where
311 T: AsyncRead + AsyncWrite + Unpin,
312
313 S: Service<Request>,
314 S::Error: Into<Response<BoxBody>>,
315 S::Response: Into<Response<B>>,
316
317 B: MessageBody,
318
319 X: Service<Request, Response = Request>,
320 X::Error: Into<Response<BoxBody>>,
321
322 U: Service<(Request, Framed<T, Codec>), Response = ()>,
323 U::Error: fmt::Display,
324{
325 fn can_read(&self, cx: &mut Context<'_>) -> bool {
326 if self.flags.contains(Flags::READ_DISCONNECT) {
327 false
328 } else if let Some(ref info) = self.payload {
329 matches!(
330 info.need_read(cx),
331 PayloadStatus::Read | PayloadStatus::Dropped
332 )
333 } else {
334 true
335 }
336 }
337
338 fn client_disconnected(self: Pin<&mut Self>) {
339 let this = self.project();
340
341 this.flags
342 .insert(Flags::READ_DISCONNECT | Flags::WRITE_DISCONNECT);
343
344 if let Some(mut payload) = this.payload.take() {
345 payload.set_error(PayloadError::Incomplete(None));
346 }
347 }
348
349 fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
350 let InnerDispatcherProj { io, write_buf, .. } = self.project();
351 let mut io = Pin::new(io.as_mut().unwrap());
352
353 let len = write_buf.len();
354 let mut written = 0;
355
356 while written < len {
357 match io.as_mut().poll_write(cx, &write_buf[written..])? {
358 Poll::Ready(0) => {
359 error!("write zero; closing");
360 return Poll::Ready(Err(io::Error::new(io::ErrorKind::WriteZero, "")));
361 }
362
363 Poll::Ready(n) => written += n,
364
365 Poll::Pending => {
366 write_buf.advance(written);
367 return Poll::Pending;
368 }
369 }
370 }
371
372 write_buf.clear();
374
375 io.poll_flush(cx)
377 }
378
379 fn enter_linger(flags: &mut Flags) {
380 flags.remove(Flags::KEEP_ALIVE);
381 flags.insert(Flags::LINGER | Flags::FINISHED);
382 }
383
384 fn ensure_linger_timer(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> bool {
385 let this = self.as_mut().project();
386
387 if matches!(this.shutdown_timer, TimerState::Active { .. }) {
388 return true;
389 }
390
391 if let Some(deadline) = this.config.client_disconnect_deadline() {
392 this.shutdown_timer
393 .set_and_init(cx, sleep_until(deadline.into()), line!());
394 true
395 } else {
396 false
397 }
398 }
399
400 fn poll_linger(
401 mut self: Pin<&mut Self>,
402 cx: &mut Context<'_>,
403 ) -> Result<Poll<()>, DispatchError> {
404 if self.as_mut().poll_flush(cx)?.is_pending() {
405 return Ok(Poll::Pending);
406 }
407
408 if !self.as_mut().ensure_linger_timer(cx) {
409 let this = self.as_mut().project();
410 this.flags.remove(Flags::LINGER);
411 this.flags.insert(Flags::SHUTDOWN);
412 return Ok(Poll::Ready(()));
413 }
414
415 loop {
416 let should_disconnect = self.as_mut().read_available(cx)?;
417 let this = self.as_mut().project();
418 let mut progressed = false;
419
420 if !this.read_buf.is_empty() {
421 this.read_buf.clear();
422 progressed = true;
423 }
424
425 if should_disconnect {
426 this.flags.remove(Flags::LINGER);
427 this.flags.insert(Flags::READ_DISCONNECT | Flags::SHUTDOWN);
428 return Ok(Poll::Ready(()));
429 }
430
431 if !progressed {
432 return Ok(Poll::Pending);
433 }
434 }
435 }
436
437 fn send_response_inner(
438 self: Pin<&mut Self>,
439 res: Response<()>,
440 body: &impl MessageBody,
441 ) -> Result<BodySize, DispatchError> {
442 let this = self.project();
443
444 let size = body.size();
445
446 this.codec
447 .encode(Message::Item((res, size)), this.write_buf)
448 .map_err(|err| {
449 if let Some(mut payload) = this.payload.take() {
450 payload.set_error(PayloadError::Incomplete(None));
451 }
452
453 DispatchError::Io(err)
454 })?;
455
456 Ok(size)
457 }
458
459 fn send_response(
460 mut self: Pin<&mut Self>,
461 mut res: Response<()>,
462 body: B,
463 ) -> Result<(), DispatchError> {
464 let is_upgrade = res.upgrade();
465 let (draining, close_for_unread_payload) = {
466 let this = self.as_mut().project();
467 (
468 this.flags.contains(Flags::DRAINING),
469 !is_upgrade
470 && should_close_for_unread_payload(
471 this.payload.as_ref(),
472 *this.payload_drainable,
473 ),
474 )
475 };
476 let close_after_response = (!is_upgrade && draining) || close_for_unread_payload;
477
478 if close_after_response {
479 res.head_mut().set_connection_type(ConnectionType::Close);
480 }
481
482 let size = self.as_mut().send_response_inner(res, &body)?;
483 match size {
484 BodySize::None | BodySize::Sized(0) => {
485 let mut this = self.as_mut().project();
486
487 if close_for_unread_payload {
488 if this.config.client_disconnect_deadline().is_some() {
489 Self::enter_linger(this.flags);
490 } else {
491 this.flags.insert(Flags::SHUTDOWN | Flags::FINISHED);
492 }
493 } else {
494 this.flags.insert(Flags::FINISHED);
495 }
496
497 this.state.set(State::None);
498 }
499 _ => self
500 .as_mut()
501 .project()
502 .state
503 .set(State::SendPayload { body }),
504 }
505
506 Ok(())
507 }
508
509 fn send_error_response(
510 mut self: Pin<&mut Self>,
511 mut res: Response<()>,
512 body: BoxBody,
513 ) -> Result<(), DispatchError> {
514 let is_upgrade = res.upgrade();
515 let (draining, close_for_unread_payload) = {
516 let this = self.as_mut().project();
517 (
518 this.flags.contains(Flags::DRAINING),
519 !is_upgrade
520 && should_close_for_unread_payload(
521 this.payload.as_ref(),
522 *this.payload_drainable,
523 ),
524 )
525 };
526 let close_after_response = (!is_upgrade && draining) || close_for_unread_payload;
527
528 if close_after_response {
529 res.head_mut().set_connection_type(ConnectionType::Close);
530 }
531
532 let size = self.as_mut().send_response_inner(res, &body)?;
533 match size {
534 BodySize::None | BodySize::Sized(0) => {
535 let mut this = self.as_mut().project();
536
537 if close_for_unread_payload {
538 if this.config.client_disconnect_deadline().is_some() {
539 Self::enter_linger(this.flags);
540 } else {
541 this.flags.insert(Flags::SHUTDOWN | Flags::FINISHED);
542 }
543 } else {
544 this.flags.insert(Flags::FINISHED);
545 }
546
547 this.state.set(State::None);
548 }
549 _ => self
550 .as_mut()
551 .project()
552 .state
553 .set(State::SendErrorPayload { body }),
554 }
555
556 Ok(())
557 }
558
559 fn send_continue(self: Pin<&mut Self>) {
560 self.project()
561 .write_buf
562 .extend_from_slice(b"HTTP/1.1 100 Continue\r\n\r\n");
563 }
564
565 fn poll_response(
566 mut self: Pin<&mut Self>,
567 cx: &mut Context<'_>,
568 ) -> Result<PollResponse, DispatchError> {
569 'res: loop {
570 let mut this = self.as_mut().project();
571 match this.state.as_mut().project() {
572 StateProj::None if this.flags.contains(Flags::DRAINING) => {
573 this.messages.clear();
574 this.flags.remove(Flags::KEEP_ALIVE);
575
576 if !this.flags.contains(Flags::LINGER) {
577 this.flags.insert(Flags::SHUTDOWN);
578 }
579
580 return Ok(PollResponse::DoNothing);
581 }
582
583 StateProj::None => match this.messages.pop_front() {
585 Some(DispatcherMessage::Item(req)) => {
587 if req.head().expect() {
589 let fut = this.flow.expect.call(req);
591 this.state.set(State::ExpectCall { fut });
592 } else {
593 let fut = this.flow.service.call(req);
595 this.state.set(State::ServiceCall { fut });
596 };
597 }
598
599 Some(DispatcherMessage::Error(res)) => {
601 self.as_mut().send_error_response(res, BoxBody::new(()))?;
605 }
606
607 Some(DispatcherMessage::Upgrade(req)) => return Ok(PollResponse::Upgrade(req)),
609
610 None => {
612 this.flags.set(
614 Flags::KEEP_ALIVE,
615 this.payload.is_none() && this.codec.keep_alive(),
616 );
617
618 return Ok(PollResponse::DoNothing);
619 }
620 },
621
622 StateProj::ServiceCall { fut } => {
623 match fut.poll(cx) {
624 Poll::Ready(Ok(res)) => {
626 let (res, body) = res.into().replace_body(());
627 self.as_mut().send_response(res, body)?;
628 }
629
630 Poll::Ready(Err(err)) => {
632 let res: Response<BoxBody> = err.into();
633 let (res, body) = res.replace_body(());
634 self.as_mut().send_error_response(res, body)?;
635 }
636
637 Poll::Pending => {
640 if !self.as_mut().poll_request(cx)? {
643 return Ok(PollResponse::DoNothing);
644 }
645 }
647 }
648 }
649
650 StateProj::SendPayload { mut body } => {
651 while this.write_buf.len() < *this.h1_write_buffer_size {
654 match body.as_mut().poll_next(cx) {
655 Poll::Ready(Some(Ok(item))) => {
656 this.codec
657 .encode(Message::Chunk(Some(item)), this.write_buf)?;
658 }
659
660 Poll::Ready(None) => {
661 this.codec.encode(Message::Chunk(None), this.write_buf)?;
662
663 let close_for_unread_payload = should_close_for_unread_payload(
668 this.payload.as_ref(),
669 *this.payload_drainable,
670 );
671 let not_pipelined = this.messages.is_empty();
672
673 this.state.set(State::None);
676
677 if not_pipelined && close_for_unread_payload {
678 if this.config.client_disconnect_deadline().is_some() {
679 Self::enter_linger(this.flags);
680 } else {
681 this.flags.insert(Flags::SHUTDOWN | Flags::FINISHED);
682 }
683 } else {
684 this.flags.insert(Flags::FINISHED);
685 }
686
687 continue 'res;
688 }
689
690 Poll::Ready(Some(Err(err))) => {
691 let err = err.into();
692 tracing::error!("Response payload stream error: {err:?}");
693 this.flags.insert(Flags::FINISHED);
694 return Err(DispatchError::Body(err));
695 }
696
697 Poll::Pending => return Ok(PollResponse::DoNothing),
698 }
699 }
700
701 return Ok(PollResponse::DrainWriteBuf);
704 }
705
706 StateProj::SendErrorPayload { mut body } => {
707 while this.write_buf.len() < *this.h1_write_buffer_size {
712 match body.as_mut().poll_next(cx) {
713 Poll::Ready(Some(Ok(item))) => {
714 this.codec
715 .encode(Message::Chunk(Some(item)), this.write_buf)?;
716 }
717
718 Poll::Ready(None) => {
719 this.codec.encode(Message::Chunk(None), this.write_buf)?;
720
721 let close_for_unread_payload = should_close_for_unread_payload(
726 this.payload.as_ref(),
727 *this.payload_drainable,
728 );
729 let not_pipelined = this.messages.is_empty();
730
731 this.state.set(State::None);
734
735 if not_pipelined && close_for_unread_payload {
736 if this.config.client_disconnect_deadline().is_some() {
737 Self::enter_linger(this.flags);
738 } else {
739 this.flags.insert(Flags::SHUTDOWN | Flags::FINISHED);
740 }
741 } else {
742 this.flags.insert(Flags::FINISHED);
743 }
744
745 continue 'res;
746 }
747
748 Poll::Ready(Some(Err(err))) => {
749 tracing::error!("Response payload stream error: {err:?}");
750 this.flags.insert(Flags::FINISHED);
751 return Err(DispatchError::Body(
752 Error::new_body().with_cause(err).into(),
753 ));
754 }
755
756 Poll::Pending => return Ok(PollResponse::DoNothing),
757 }
758 }
759
760 return Ok(PollResponse::DrainWriteBuf);
763 }
764
765 StateProj::ExpectCall { fut } => {
766 trace!(" calling expect service");
767
768 match fut.poll(cx) {
769 Poll::Ready(Ok(req)) => {
772 this.write_buf
773 .extend_from_slice(b"HTTP/1.1 100 Continue\r\n\r\n");
774 let fut = this.flow.service.call(req);
775 this.state.set(State::ServiceCall { fut });
776 }
777
778 Poll::Ready(Err(err)) => {
780 let res: Response<BoxBody> = err.into();
781 let (res, body) = res.replace_body(());
782 self.as_mut().send_error_response(res, body)?;
783 }
784
785 Poll::Pending => return Ok(PollResponse::DoNothing),
787 }
788 }
789 }
790 }
791 }
792
793 fn handle_request(
794 mut self: Pin<&mut Self>,
795 req: Request,
796 cx: &mut Context<'_>,
797 ) -> Result<(), DispatchError> {
798 {
800 let mut this = self.as_mut().project();
801
802 if req.head().expect() {
804 let fut = this.flow.expect.call(req);
806 this.state.set(State::ExpectCall { fut });
807 } else {
808 let fut = this.flow.service.call(req);
810 this.state.set(State::ServiceCall { fut });
811 };
812 };
813
814 loop {
816 match self.as_mut().project().state.project() {
817 StateProj::ExpectCall { fut } => {
818 match fut.poll(cx) {
819 Poll::Ready(Ok(req)) => {
821 self.as_mut().send_continue();
822
823 let mut this = self.as_mut().project();
824 let fut = this.flow.service.call(req);
825 this.state.set(State::ServiceCall { fut });
826
827 continue;
828 }
829
830 Poll::Ready(Err(err)) => {
834 let res: Response<BoxBody> = err.into();
835 let (res, body) = res.replace_body(());
836 return self.send_error_response(res, body);
837 }
838
839 Poll::Pending => return Ok(()),
842 }
843 }
844
845 StateProj::ServiceCall { fut } => {
846 return match fut.poll(cx) {
848 Poll::Ready(Ok(res)) => {
852 let (res, body) = res.into().replace_body(());
853 self.as_mut().send_response(res, body)
854 }
855
856 Poll::Pending => Ok(()),
858
859 Poll::Ready(Err(err)) => {
861 let res: Response<BoxBody> = err.into();
862 let (res, body) = res.replace_body(());
863 self.as_mut().send_error_response(res, body)
864 }
865 };
866 }
867
868 _ => {
869 unreachable!("State must be set to ServiceCall or ExceptCall in handle_request")
870 }
871 }
872 }
873 }
874
875 fn poll_request(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Result<bool, DispatchError> {
879 if self.flags.contains(Flags::DRAINING) && self.state.is_none() {
880 return Ok(false);
881 }
882
883 let pipeline_queue_full = self.messages.len() >= MAX_PIPELINED_MESSAGES;
884 let can_not_read = !self.can_read(cx);
885
886 if pipeline_queue_full || can_not_read {
888 return Ok(false);
889 }
890
891 let mut this = self.as_mut().project();
892
893 let mut updated = false;
894
895 loop {
897 match this.codec.decode(this.read_buf) {
898 Ok(Some(msg)) => {
899 updated = true;
900
901 match msg {
902 Message::Item(mut req) => {
903 this.head_timer.clear(line!());
905
906 req.head_mut().peer_addr = *this.peer_addr;
907
908 req.conn_data.clone_from(this.conn_data);
909
910 match this.codec.message_type() {
911 MessageType::None => *this.payload_drainable = false,
913
914 MessageType::Stream if this.flow.upgrade.is_some() => {
918 *this.payload_drainable = false;
919 this.messages.push_back(DispatcherMessage::Upgrade(req));
920 break;
921 }
922
923 MessageType::Payload | MessageType::Stream => {
925 let (sender, payload) = Payload::create(false);
931 *req.payload() = crate::Payload::H1 { payload };
932 *this.payload = Some(sender);
933 *this.payload_drainable = req.chunked().unwrap_or(false);
934 }
935 }
936
937 if this.state.is_none() {
939 self.as_mut().handle_request(req, cx)?;
940 this = self.as_mut().project();
941 } else {
942 this.messages.push_back(DispatcherMessage::Item(req));
943 }
944 }
945
946 Message::Chunk(Some(chunk)) => {
947 if let Some(ref mut payload) = this.payload {
948 payload.feed_data(chunk);
949 } else {
950 error!("Internal server error: unexpected payload chunk");
951 this.flags.insert(Flags::READ_DISCONNECT);
952 this.messages.push_back(DispatcherMessage::Error(
953 Response::internal_server_error().drop_body(),
954 ));
955 *this.error = Some(DispatchError::InternalError);
956 break;
957 }
958 }
959
960 Message::Chunk(None) => {
961 if let Some(mut payload) = this.payload.take() {
962 payload.feed_eof();
963 *this.payload_drainable = false;
964 } else {
965 error!("Internal server error: unexpected eof");
966 this.flags.insert(Flags::READ_DISCONNECT);
967 this.messages.push_back(DispatcherMessage::Error(
968 Response::internal_server_error().drop_body(),
969 ));
970 *this.error = Some(DispatchError::InternalError);
971 break;
972 }
973 }
974 }
975 }
976
977 Ok(None) => break,
980
981 Err(ParseError::Io(err)) => {
982 trace!("I/O error: {}", &err);
983 self.as_mut().client_disconnected();
984 this = self.as_mut().project();
985 *this.error = Some(DispatchError::Io(err));
986 break;
987 }
988
989 Err(ParseError::TooLarge) => {
990 trace!("request head was too big; returning 431 response");
991
992 if let Some(mut payload) = this.payload.take() {
993 payload.set_error(PayloadError::Overflow);
994 }
995
996 this.messages
998 .push_back(DispatcherMessage::Error(Response::with_body(
999 StatusCode::REQUEST_HEADER_FIELDS_TOO_LARGE,
1000 (),
1001 )));
1002
1003 this.flags.insert(Flags::READ_DISCONNECT);
1004 *this.error = Some(ParseError::TooLarge.into());
1005
1006 break;
1007 }
1008
1009 Err(err) => {
1010 trace!("parse error {}", &err);
1011
1012 if let Some(mut payload) = this.payload.take() {
1013 payload.set_error(PayloadError::EncodingCorrupted);
1014 }
1015
1016 this.messages.push_back(DispatcherMessage::Error(
1018 Response::bad_request().drop_body(),
1019 ));
1020
1021 this.flags.insert(Flags::READ_DISCONNECT);
1022 *this.error = Some(err.into());
1023 break;
1024 }
1025 }
1026 }
1027
1028 Ok(updated)
1029 }
1030
1031 fn poll_head_timer(
1032 mut self: Pin<&mut Self>,
1033 cx: &mut Context<'_>,
1034 ) -> Result<(), DispatchError> {
1035 let this = self.as_mut().project();
1036
1037 if let TimerState::Active { timer } = this.head_timer {
1038 if timer.as_mut().poll(cx).is_ready() {
1039 trace!("timed out on slow request; replying with 408 and closing connection");
1042
1043 let _ = self.as_mut().send_error_response(
1044 Response::with_body(StatusCode::REQUEST_TIMEOUT, ()),
1045 BoxBody::new(()),
1046 );
1047
1048 self.project().flags.insert(Flags::SHUTDOWN);
1049 }
1050 };
1051
1052 Ok(())
1053 }
1054
1055 fn poll_ka_timer(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Result<(), DispatchError> {
1056 let this = self.as_mut().project();
1057 if let TimerState::Active { timer } = this.ka_timer {
1058 debug_assert!(
1059 this.flags.contains(Flags::KEEP_ALIVE),
1060 "keep-alive flag should be set when timer is active",
1061 );
1062 debug_assert!(
1063 this.state.is_none(),
1064 "dispatcher should not be in keep-alive phase if state is not none: {:?}",
1065 this.state,
1066 );
1067
1068 if timer.as_mut().poll(cx).is_ready() {
1080 trace!("timer timed out; closing connection");
1082 this.flags.insert(Flags::SHUTDOWN);
1083
1084 if let Some(deadline) = this.config.client_disconnect_deadline() {
1085 this.shutdown_timer
1087 .set_and_init(cx, sleep_until(deadline.into()), line!());
1088 } else {
1089 this.flags.insert(Flags::WRITE_DISCONNECT);
1091 }
1092 }
1093 }
1094
1095 Ok(())
1096 }
1097
1098 fn poll_shutdown_timer(
1099 mut self: Pin<&mut Self>,
1100 cx: &mut Context<'_>,
1101 ) -> Result<(), DispatchError> {
1102 let this = self.as_mut().project();
1103 if let TimerState::Active { timer } = this.shutdown_timer {
1104 debug_assert!(
1105 this.flags.intersects(Flags::LINGER | Flags::SHUTDOWN),
1106 "shutdown or linger flag should be set when timer is active",
1107 );
1108
1109 if timer.as_mut().poll(cx).is_ready() {
1110 if this.flags.contains(Flags::LINGER) {
1111 trace!("timed-out during linger; shutting down connection");
1112 this.flags.remove(Flags::LINGER);
1113 this.flags.insert(Flags::SHUTDOWN);
1114 this.shutdown_timer.clear(line!());
1115 } else {
1116 trace!("timed-out during shutdown");
1117 return Err(DispatchError::DisconnectTimeout);
1118 }
1119 }
1120 }
1121
1122 Ok(())
1123 }
1124
1125 fn poll_graceful_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) {
1126 let this = self.as_mut().project();
1127
1128 let notified = this
1129 .graceful_shutdown
1130 .as_mut()
1131 .is_some_and(|signal| signal.as_mut().poll(cx).is_ready());
1132
1133 if notified {
1134 *this.graceful_shutdown = None;
1135 this.flags.remove(Flags::KEEP_ALIVE);
1136 this.flags.insert(Flags::DRAINING);
1137
1138 if this.ka_timer.is_enabled() {
1139 this.ka_timer.clear(line!());
1140 }
1141 }
1142 }
1143
1144 fn poll_timers(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Result<(), DispatchError> {
1146 self.as_mut().poll_head_timer(cx)?;
1147 self.as_mut().poll_ka_timer(cx)?;
1148 self.as_mut().poll_shutdown_timer(cx)?;
1149
1150 Ok(())
1151 }
1152
1153 #[inline(always)] fn read_available(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Result<bool, DispatchError> {
1160 let this = self.project();
1161
1162 if this.flags.contains(Flags::READ_DISCONNECT) {
1163 return Ok(false);
1164 };
1165
1166 let mut io = Pin::new(this.io.as_mut().unwrap());
1167
1168 let mut read_some = false;
1169
1170 loop {
1171 if this.read_buf.len() >= MAX_BUFFER_SIZE {
1173 match this.payload.as_ref().map(|p| p.need_read(cx)) {
1192 Some(PayloadStatus::Pause) => {}
1194
1195 Some(PayloadStatus::Dropped) | Some(PayloadStatus::Read) | None => {
1197 cx.waker().wake_by_ref()
1198 }
1199 }
1200
1201 return Ok(false);
1202 }
1203
1204 let remaining = this.read_buf.capacity() - this.read_buf.len();
1206 if remaining < LW_BUFFER_SIZE {
1207 this.read_buf.reserve(HW_BUFFER_SIZE - remaining);
1208 }
1209
1210 match tokio_util::io::poll_read_buf(io.as_mut(), cx, this.read_buf) {
1211 Poll::Ready(Ok(n)) => {
1212 if !this.payload.as_ref().is_some_and(|pl| pl.is_dropped()) {
1215 this.flags.remove(Flags::FINISHED);
1216 }
1217
1218 if n == 0 {
1219 return Ok(true);
1220 }
1221
1222 read_some = true;
1223 }
1224
1225 Poll::Pending => {
1226 return Ok(false);
1227 }
1228
1229 Poll::Ready(Err(err)) => {
1230 return match err.kind() {
1231 io::ErrorKind::WouldBlock => Ok(false),
1233
1234 io::ErrorKind::ConnectionReset if read_some => Ok(true),
1236
1237 _ => Err(DispatchError::Io(err)),
1238 };
1239 }
1240 }
1241 }
1242 }
1243
1244 fn upgrade(self: Pin<&mut Self>, req: Request) -> U::Future {
1246 let this = self.project();
1247 let mut parts = FramedParts::with_read_buf(
1248 this.io.take().unwrap(),
1249 mem::take(this.codec),
1250 mem::take(this.read_buf),
1251 );
1252 parts.write_buf = mem::take(this.write_buf);
1253 let framed = Framed::from_parts(parts);
1254 this.flow.upgrade.as_ref().unwrap().call((req, framed))
1255 }
1256}
1257
1258impl<T, S, B, X, U> Future for Dispatcher<T, S, B, X, U>
1259where
1260 T: AsyncRead + AsyncWrite + Unpin,
1261
1262 S: Service<Request>,
1263 S::Error: Into<Response<BoxBody>>,
1264 S::Response: Into<Response<B>>,
1265
1266 B: MessageBody,
1267
1268 X: Service<Request, Response = Request>,
1269 X::Error: Into<Response<BoxBody>>,
1270
1271 U: Service<(Request, Framed<T, Codec>), Response = ()>,
1272 U::Error: fmt::Display,
1273{
1274 type Output = Result<(), DispatchError>;
1275
1276 #[inline]
1277 fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
1278 let this = self.as_mut().project();
1279
1280 #[cfg(test)]
1281 {
1282 *this.poll_count += 1;
1283 }
1284
1285 match this.inner.project() {
1286 DispatcherStateProj::Upgrade { fut: upgrade } => upgrade.poll(cx).map_err(|err| {
1287 error!("Upgrade handler error: {}", err);
1288 DispatchError::Upgrade
1289 }),
1290
1291 DispatcherStateProj::Normal { mut inner } => {
1292 trace!("start flags: {:?}", &inner.flags);
1293
1294 trace_timer_states(
1295 "start",
1296 &inner.head_timer,
1297 &inner.ka_timer,
1298 &inner.shutdown_timer,
1299 );
1300
1301 inner.as_mut().poll_graceful_shutdown(cx);
1302 inner.as_mut().poll_timers(cx)?;
1303
1304 let poll = if inner.flags.contains(Flags::LINGER) {
1305 match inner.as_mut().poll_linger(cx)? {
1306 Poll::Ready(()) => {
1307 cx.waker().wake_by_ref();
1308 Poll::Pending
1309 }
1310 Poll::Pending => Poll::Pending,
1311 }
1312 } else if inner.flags.contains(Flags::SHUTDOWN) {
1313 if inner.flags.contains(Flags::WRITE_DISCONNECT) {
1314 Poll::Ready(Ok(()))
1315 } else {
1316 ready!(inner.as_mut().poll_flush(cx))?;
1318 Pin::new(inner.as_mut().project().io.as_mut().unwrap())
1319 .poll_shutdown(cx)
1320 .map_err(DispatchError::from)
1321 }
1322 } else {
1323 let should_disconnect = inner.as_mut().read_available(cx)?;
1325
1326 if !inner.read_buf.is_empty() && inner.flags.contains(Flags::KEEP_ALIVE) {
1328 let inner = inner.as_mut().project();
1329 inner.flags.remove(Flags::KEEP_ALIVE);
1330 inner.ka_timer.clear(line!());
1331 }
1332
1333 if !inner.flags.contains(Flags::STARTED) {
1334 inner.as_mut().project().flags.insert(Flags::STARTED);
1335
1336 if let Some(deadline) = inner.config.client_request_deadline() {
1337 inner.as_mut().project().head_timer.set_and_init(
1338 cx,
1339 sleep_until(deadline.into()),
1340 line!(),
1341 );
1342 }
1343 }
1344
1345 inner.as_mut().poll_request(cx)?;
1346
1347 if should_disconnect {
1348 let inner = inner.as_mut().project();
1350 inner.flags.insert(Flags::READ_DISCONNECT);
1351 if let Some(mut payload) = inner.payload.take() {
1352 payload.set_error(PayloadError::Incomplete(None));
1353 payload.feed_eof();
1354 }
1355 };
1356
1357 loop {
1358 let drain = match inner.as_mut().poll_response(cx)? {
1361 PollResponse::DrainWriteBuf => true,
1362
1363 PollResponse::DoNothing => {
1364 if inner.flags.contains(Flags::KEEP_ALIVE | Flags::FINISHED) {
1367 if let Some(timer) = inner.config.keep_alive_deadline() {
1368 inner.as_mut().project().ka_timer.set_and_init(
1369 cx,
1370 sleep_until(timer.into()),
1371 line!(),
1372 );
1373 }
1374 }
1375
1376 false
1377 }
1378
1379 PollResponse::Upgrade(req) => {
1381 let upgrade = inner.upgrade(req);
1382 self.as_mut()
1383 .project()
1384 .inner
1385 .set(DispatcherState::Upgrade { fut: upgrade });
1386 return self.poll(cx);
1387 }
1388 };
1389
1390 let flush_was_ready = inner.as_mut().poll_flush(cx)?.is_ready();
1397
1398 if !flush_was_ready || !drain {
1403 break;
1404 }
1405 }
1406
1407 if inner.flags.contains(Flags::WRITE_DISCONNECT) {
1409 trace!("client is gone; disconnecting");
1410 return Poll::Ready(Ok(()));
1411 }
1412
1413 let inner_p = inner.as_mut().project();
1414 let state_is_none = inner_p.state.is_none();
1415
1416 if inner_p.flags.contains(Flags::READ_DISCONNECT)
1424 && (!inner_p.config.h1_allow_half_closed() || state_is_none)
1425 {
1426 trace!("read half closed; start shutdown");
1427 inner_p.flags.insert(Flags::SHUTDOWN);
1428 }
1429
1430 if state_is_none && inner_p.write_buf.is_empty() {
1432 if let Some(err) = inner_p.error.take() {
1433 error!("stream error: {}", &err);
1434 return Poll::Ready(Err(err));
1435 }
1436
1437 if inner_p.flags.contains(Flags::FINISHED)
1439 && !inner_p.flags.contains(Flags::KEEP_ALIVE)
1440 && inner_p.payload.is_none()
1441 {
1442 inner_p.flags.remove(Flags::FINISHED);
1443 inner_p.flags.insert(Flags::SHUTDOWN);
1444 return self.poll(cx);
1445 }
1446
1447 if inner_p.flags.contains(Flags::SHUTDOWN) {
1449 return self.poll(cx);
1450 }
1451 }
1452
1453 trace_timer_states(
1454 "end",
1455 inner_p.head_timer,
1456 inner_p.ka_timer,
1457 inner_p.shutdown_timer,
1458 );
1459
1460 if inner_p.flags.intersects(Flags::LINGER | Flags::SHUTDOWN) {
1461 cx.waker().wake_by_ref();
1462 }
1463 Poll::Pending
1464 };
1465
1466 trace!("end flags: {:?}", &inner.flags);
1467
1468 poll
1469 }
1470 }
1471 }
1472}
1473
1474fn should_close_for_unread_payload(
1475 payload: Option<&PayloadSender>,
1476 payload_drainable: bool,
1477) -> bool {
1478 let payload_unfinished = payload.is_some();
1479 let drain_payload = payload.is_some_and(|pl| pl.is_dropped()) && payload_drainable;
1480
1481 payload_unfinished && !drain_payload
1482}
1483
1484#[allow(dead_code)]
1485fn trace_timer_states(
1486 label: &str,
1487 head_timer: &TimerState,
1488 ka_timer: &TimerState,
1489 shutdown_timer: &TimerState,
1490) {
1491 trace!("{} timers:", label);
1492
1493 if head_timer.is_enabled() {
1494 trace!(" head {}", &head_timer);
1495 }
1496
1497 if ka_timer.is_enabled() {
1498 trace!(" keep-alive {}", &ka_timer);
1499 }
1500
1501 if shutdown_timer.is_enabled() {
1502 trace!(" shutdown {}", &shutdown_timer);
1503 }
1504}