Line data    Source code

       1              : use std::net::SocketAddr;
       2              : 
       3              : use tokio::io::{AsyncReadExt, AsyncWriteExt};
       4              : use tokio::net::{TcpListener, TcpStream};
       5              : 
       6              : /// A request the server received
       7              : #[derive(Debug)]
       8              : pub struct ReceivedRequest {
       9              :     /// The request line, e.g. `POST /path?a=1 HTTP/1.1`
      10              :     pub line: String,
      11              :     /// The headers with lowercase names
      12              :     pub headers: Vec<(String, String)>,
      13              :     /// The body
      14              :     pub body: Vec<u8>,
      15              : }
      16              : 
      17              : impl ReceivedRequest {
      18            2 :     pub fn header(&self, name: &str) -> Option<&str> {
      19            5 :         self.headers.iter().find(|(key, _)| key == name).map(|(_, value)| value.as_str())
      20            2 :     }
      21              : }
      22              : 
      23              : /// The connection of a received request, a test writes the response by hand to control its timing
      24              : pub struct Connection(pub TcpStream);
      25              : 
      26              : impl Connection {
      27           44 :     pub async fn write(&mut self, text: &str) {
      28           44 :         self.0.write_all(text.as_bytes()).await.unwrap();
      29           44 :         self.0.flush().await.unwrap();
      30           44 :     }
      31              : 
      32              :     /// Start a chunked response
      33            9 :     pub async fn write_head(&mut self, status: &str, headers: &[(&str, &str)]) {
      34            9 :         let headers: String = headers.iter().map(|(key, value)| format!("{key}: {value}\r\n")).collect();
      35            9 :         self.write(&format!("HTTP/1.1 {status}\r\ntransfer-encoding: chunked\r\n{headers}\r\n")).await;
      36            9 :     }
      37              : 
      38           29 :     pub async fn write_chunk(&mut self, chunk: &str) {
      39           29 :         self.write(&format!("{:x}\r\n{chunk}\r\n", chunk.len())).await;
      40           29 :     }
      41              : 
      42            5 :     pub async fn finish(&mut self) {
      43            5 :         self.write("0\r\n\r\n").await;
      44            5 :     }
      45              : 
      46              :     /// Wait until the client closes the connection
      47            3 :     pub async fn closed(&mut self) {
      48            3 :         let mut buffer = [0; 1024];
      49            3 :         while !matches!(self.0.read(&mut buffer).await, Ok(0) | Err(_)) {}
      50            3 :     }
      51              : }
      52              : 
      53              : /// An HTTP server that hands every request and its connection to the test
      54           15 : pub async fn spawn_http_server() -> (SocketAddr, tokio::sync::mpsc::Receiver<(ReceivedRequest, Connection)>) {
      55           15 :     let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
      56           15 :     let address = listener.local_addr().unwrap();
      57           15 :     let (sender, receiver) = tokio::sync::mpsc::channel(4);
      58           15 :     tokio::spawn(async move {
      59           30 :         while let Ok((mut stream, _)) = listener.accept().await {
      60           15 :             let sender = sender.clone();
      61           15 :             tokio::spawn(async move {
      62           15 :                 let mut buffer = Vec::new();
      63           15 :                 let head_end = loop {
      64          837 :                     if let Some(position) = buffer.windows(4).position(|window| window == b"\r\n\r\n") {
      65           15 :                         break position;
      66           15 :                     }
      67           15 :                     let mut chunk = [0; 4096];
      68           15 :                     match stream.read(&mut chunk).await {
      69            0 :                         Ok(0) | Err(_) => return,
      70           15 :                         Ok(read) => buffer.extend_from_slice(&chunk[..read]),
      71              :                     }
      72              :                 };
      73           15 :                 let head = String::from_utf8_lossy(&buffer[..head_end]).to_string();
      74           15 :                 let mut lines = head.split("\r\n");
      75           15 :                 let line = lines.next().unwrap().to_string();
      76           15 :                 let headers: Vec<(String, String)> = lines
      77           32 :                     .filter_map(|line| line.split_once(':'))
      78           32 :                     .map(|(key, value)| (key.trim().to_lowercase(), value.trim().to_string()))
      79           15 :                     .collect();
      80           15 :                 let length: usize = headers.iter()
      81           32 :                     .find(|(key, _)| key == "content-length")
      82           15 :                     .map_or(0, |(_, value)| value.parse().unwrap());
      83           15 :                 let mut body = buffer[head_end + 4..].to_vec();
      84           15 :                 while body.len() < length {
      85            0 :                     let mut chunk = [0; 4096];
      86            0 :                     match stream.read(&mut chunk).await {
      87            0 :                         Ok(0) | Err(_) => return,
      88            0 :                         Ok(read) => body.extend_from_slice(&chunk[..read]),
      89              :                     }
      90              :                 }
      91           15 :                 let _ = sender.send((ReceivedRequest { line, headers, body }, Connection(stream))).await;
      92           15 :             });
      93              :         }
      94            0 :     });
      95           15 :     (address, receiver)
      96           15 : }