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