Conn 的读写状态机

pub(crate) struct Conn<I, B, T> {
    io: Buffered<I, EncodedBuf<B>>,
    state: State,
    _marker: PhantomData<fn(T)>,
}

#[derive(Debug)]
enum Reading {
    Init,
    Continue(Decoder),
    Body(Decoder),
    KeepAlive,
    Closed,
}

enum Writing {
    Init,
    Body(Encoder),
    KeepAlive,
    Closed,
}

#[derive(Clone, Copy, Debug, Default)]
enum KA {
    Idle,
    #[default]
    Busy,
    Disabled,
}

        self.state.busy();
        self.state.keep_alive &= msg.keep_alive;
        self.state.version = msg.head.version;

        let mut wants = if msg.wants_upgrade {
            Wants::UPGRADE
        } else {
            Wants::EMPTY
        };

        if msg.decode == DecodedLength::ZERO {
            if msg.expect_continue {
                debug!("ignoring expect-continue since body is empty");
            }
            self.state.reading = Reading::KeepAlive;
            if !T::should_read_first() {
                self.try_keep_alive(cx);
            }
        } else if msg.expect_continue && msg.head.version.gt(&Version::HTTP_10) {
            let h1_max_header_size = self.state.h1_max_header_size;
            self.state.reading = Reading::Continue(Decoder::new(
                msg.decode,
                self.state.h1_max_headers,
                h1_max_header_size,
            ));
            wants = wants.add(Wants::EXPECT);
        } else {
            let h1_max_header_size = self.state.h1_max_header_size;
            self.state.reading = Reading::Body(Decoder::new(
                msg.decode,
                self.state.h1_max_headers,
                h1_max_header_size,
            ));
        }

    pub(crate) fn write_head(&mut self, head: MessageHead<T::Outgoing>, body: Option<BodyLength>) {
        if let Some(encoder) = self.encode_head(head, body) {
            self.state.writing = if !encoder.is_eof() {
                Writing::Body(encoder)
            } else if encoder.is_last() {
                Writing::Closed
            } else {
                Writing::KeepAlive
            };
        }
    }

    fn enforce_version(&mut self, head: &mut MessageHead<T::Outgoing>) {
        match self.state.version {
            Version::HTTP_10 => {
                // Fixes response or connection when keep-alive header is not present
                self.fix_keep_alive(head);
                // If the remote only knows HTTP/1.0, we should force ourselves
                // to do only speak HTTP/1.0 as well.
                head.version = Version::HTTP_10;
            }
            Version::HTTP_11 => {
                if let KA::Disabled = self.state.keep_alive.status() {
                    match head.headers.entry(CONNECTION) {
                        Entry::Occupied(mut entry) => {
                            if !entry.iter().any(headers::connection_close) {
                                entry.append(HeaderValue::from_static("close"));
                            }
                        }
                        Entry::Vacant(entry) => {
                            entry.insert(HeaderValue::from_static("close"));
                        }
                    }
                }
            }
            _ => (),
        }
        // ...
    }

    fn try_keep_alive<T: Http1Transaction>(&mut self) {
        match (&self.reading, &self.writing) {
            (&Reading::KeepAlive, &Writing::KeepAlive) => {
                if let KA::Busy = self.keep_alive.status() {
                    self.idle::<T>();
                } else {
                    trace!(
                        "try_keep_alive({}): could keep-alive, but status = {:?}",
                        T::LOG,
                        self.keep_alive
                    );
                    self.close();
                }
            }
            (&Reading::Closed, &Writing::KeepAlive) | (&Reading::KeepAlive, &Writing::Closed) => {
                self.close();
            }
            _ => (),
        }
    }

hyper = { version = "=1.12.0", features = ["server", "http1"] }
hyper-util = { version = "=0.1.21", features = ["tokio"] }
http-body-util = "0.1"
tokio = { version = "=1.54.1", features = ["full"] }

use std::{convert::Infallible, time::Duration};
use http_body_util::Full;
use hyper::{body::{Bytes, Incoming}, server::conn::http1, service::service_fn, Request, Response};
use hyper_util::rt::TokioIo;
use tokio::{io::{AsyncReadExt, AsyncWriteExt}, net::{TcpListener, TcpStream}, time::timeout};

async fn probe(addr: std::net::SocketAddr, name: &str, req: &str) {
  let mut s = TcpStream::connect(addr).await.unwrap();
  s.write_all(req.as_bytes()).await.unwrap();
  let mut out = Vec::new();
  let mut buf = [0u8; 1024];
  let end = loop {
    match timeout(Duration::from_millis(300), s.read(&mut buf)).await {
      Ok(Ok(0)) => break "server closed",
      Ok(Ok(n)) => out.extend_from_slice(&buf[..n]),
      Ok(Err(e)) => panic!("{e}"),
      Err(_) => break "still open after 300ms",
    }
  };
  println!("== {name}\n{}\n-- {end}", String::from_utf8_lossy(&out).replace("\r\n", "\\r\\n\n"));
}

#[tokio::main]
async fn main() {
  let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
  let addr = listener.local_addr().unwrap();
  tokio::spawn(async move {
    loop {
      let (stream, _) = listener.accept().await.unwrap();
      let svc = service_fn(|req: Request<Incoming>| async move {
        let mut rsp = Response::new(Full::new(Bytes::from("ok")));
        if req.uri().path() == "/bye" {
          rsp.headers_mut().insert("connection", "close".parse().unwrap());
        }
        Ok::<_, Infallible>(rsp)
      });
      tokio::spawn(async move {
        let r = http1::Builder::new()
          .auto_date_header(false) // keep the output stable
          .serve_connection(TokioIo::new(stream), svc)
          .await;
        if let Err(e) = r {
          println!("server conn error: {e}");
        }
      });
    }
  });

  probe(addr, "1.1, no Connection header", "GET / HTTP/1.1\r\nHost: x\r\n\r\n").await;
  probe(addr, "1.1, Connection: close", "GET / HTTP/1.1\r\nHost: x\r\nConnection: close\r\n\r\n").await;
  probe(addr, "1.1, handler sets Connection: close", "GET /bye HTTP/1.1\r\nHost: x\r\n\r\n").await;
  probe(addr, "1.0, no Connection header", "GET / HTTP/1.0\r\n\r\n").await;
  probe(addr, "1.0, Connection: keep-alive", "GET / HTTP/1.0\r\nConnection: keep-alive\r\n\r\n").await;
  probe(
    addr,
    "1.1, close then keep-alive",
    "GET / HTTP/1.1\r\nHost: x\r\nConnection: close\r\nConnection: keep-alive\r\n\r\n",
  )
  .await;
}

== 1.1, no Connection header
HTTP/1.1 200 OK\r\n
content-length: 2\r\n
\r\n
ok
-- still open after 300ms
== 1.1, Connection: close
HTTP/1.1 200 OK\r\n
connection: close\r\n
content-length: 2\r\n
\r\n
ok
-- server closed
== 1.1, handler sets Connection: close
HTTP/1.1 200 OK\r\n
connection: close\r\n
content-length: 2\r\n
\r\n
ok
-- server closed
== 1.0, no Connection header
HTTP/1.0 200 OK\r\n
content-length: 2\r\n
\r\n
ok
-- server closed
== 1.0, Connection: keep-alive
HTTP/1.0 200 OK\r\n
connection: keep-alive\r\n
content-length: 2\r\n
\r\n
ok
-- still open after 300ms
== 1.1, close then keep-alive
HTTP/1.1 200 OK\r\n
connection: close\r\n
content-length: 2\r\n
\r\n
ok
-- server closed

            Reading::Continue(decoder) => {
                // Write the 100 Continue if not already responded...
                if let Writing::Init = self.state.writing {
                    trace!("automatically sending 100 Continue");
                    let cont = b"HTTP/1.1 100 Continue\r\n\r\n";
                    self.io.headers_buf().extend_from_slice(cont);
                }

                // And now recurse once in the Reading::Body state...
                self.state.reading = Reading::Body(decoder.clone());
                return self.poll_read_body(cx);
            }

hyper = { version = "=1.12.0", features = ["server", "http1"] }
hyper-util = { version = "=0.1.21", features = ["tokio"] }
http-body-util = "0.1"
tokio = { version = "=1.54.1", features = ["full"] }

use std::{convert::Infallible, time::{Duration, Instant}};
use http_body_util::{BodyExt, Full};
use hyper::{body::{Bytes, Incoming}, server::conn::http1, service::service_fn, Request, Response};
use hyper_util::rt::TokioIo;
use tokio::{io::{AsyncReadExt, AsyncWriteExt}, net::{TcpListener, TcpStream}, time::timeout};

#[tokio::main]
async fn main() {
  let start = Instant::now();
  let ms = move || start.elapsed().as_millis();
  let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
  let addr = listener.local_addr().unwrap();
  tokio::spawn(async move {
    for i in 1.. {
      let (stream, _) = listener.accept().await.unwrap();
      let svc = service_fn(move |req: Request<Incoming>| async move {
        if req.uri().path() == "/read" {
          println!("{:>4}ms server: handler called, reads body in 200ms", ms());
          tokio::time::sleep(Duration::from_millis(200)).await;
          let body = req.into_body().collect().await.unwrap().to_bytes();
          println!("{:>4}ms server: body = {body:?}", ms());
        } else {
          println!("{:>4}ms server: handler called, rejects without reading body", ms());
        }
        Ok::<_, Infallible>(Response::new(Full::new(Bytes::from("done"))))
      });
      tokio::spawn(async move {
        let r = http1::Builder::new()
          .auto_date_header(false)
          .serve_connection(TokioIo::new(stream), svc)
          .await;
        println!("{:>4}ms server: connection {i} ended {r:?}", ms());
      });
    }
  });

  for path in ["/read", "/reject"] {
    println!("== {path}");
    let mut s = TcpStream::connect(addr).await.unwrap();
    let head = format!("POST {path} HTTP/1.1\r\nHost: x\r\nContent-Length: 5\r\nExpect: 100-continue\r\n\r\n");
    s.write_all(head.as_bytes()).await.unwrap();
    let mut buf = [0u8; 1024];
    let mut sent_body = false;
    loop {
      match timeout(Duration::from_millis(500), s.read(&mut buf)).await {
        Ok(Ok(0)) => {
          println!("{:>4}ms client: EOF", ms());
          break;
        }
        Ok(Ok(n)) => {
          let text = String::from_utf8_lossy(&buf[..n]).replace("\r\n", "\\r\\n");
          println!("{:>4}ms client: got {text}", ms());
          if text.starts_with("HTTP/1.1 100") && !sent_body {
            s.write_all(b"hello").await.unwrap();
            sent_body = true;
          }
        }
        Ok(Err(e)) => panic!("{e}"),
        Err(_) => {
          println!("{:>4}ms client: nothing more for 500ms", ms());
          break;
        }
      }
    }
  }
}

== /read
   0ms server: handler called, reads body in 200ms
 203ms client: got HTTP/1.1 100 Continue\r\n\r\n
 203ms server: body = b"hello"
 204ms client: got HTTP/1.1 200 OK\r\ncontent-length: 4\r\n\r\ndone
 705ms client: nothing more for 500ms
== /reject
 706ms server: connection 1 ended Ok(())
 706ms server: handler called, rejects without reading body
 706ms server: connection 2 ended Ok(())
 706ms client: got HTTP/1.1 200 OK\r\ncontent-length: 4\r\n\r\ndone
 706ms client: EOF

    /// Set whether HTTP/1 connections should support half-closures.
    ///
    /// Clients can chose to shutdown their write-side while waiting
    /// for the server to respond. Setting this to `true` will
    /// prevent closing the connection immediately if `read`
    /// detects an EOF in the middle of a request.
    ///
    /// Default is `false`.

    pub(crate) fn poll_read_keep_alive(&mut self, cx: &mut Context<'_>) -> Poll<crate::Result<()>> {
        debug_assert!(!self.can_read_head() && !self.can_read_body());

        if self.is_read_closed() {
            Poll::Pending
        } else if self.is_mid_message() {
            self.mid_message_detect_eof(cx)
        } else {
            self.require_empty_read(cx)
        }
    }

    fn mid_message_detect_eof(&mut self, cx: &mut Context<'_>) -> Poll<crate::Result<()>> {
        debug_assert!(!self.can_read_head() && !self.can_read_body() && !self.is_read_closed());
        debug_assert!(self.is_mid_message());

        if self.state.allow_half_close || !self.io.read_buf().is_empty() {
            return Poll::Pending;
        }

        let num_read = ready!(self.force_io_read(cx)).map_err(crate::Error::new_io)?;

        if num_read == 0 {
            trace!("found unexpected EOF on busy connection: {:?}", self.state);
            self.state.close_read();
            Poll::Ready(Err(crate::Error::new_incomplete()))
        } else {
            Poll::Ready(Ok(()))
        }
    }

half_close(false): connection Err(hyper::Error(IncompleteMessage)), client got 2 responses
half_close(true): connection Ok(()), client got 3 responses