Line data Source code
1 : //! The little HTTP/1.1 that the agent server and the command line tool speak.
2 : //! Every exchange is one request and one response on its own connection.
3 :
4 : use tokio::io::{AsyncRead, AsyncReadExt};
5 :
6 : /// The size limit of a request head
7 : const MAX_HEAD_BYTES: usize = 64 * 1024;
8 :
9 : /// The head of a request
10 : #[derive(Debug)]
11 : pub struct Request {
12 : pub method: String,
13 : /// The request target, e.g. `/v1/call?x=1`
14 : pub target: String,
15 : /// The headers with lowercase names
16 : pub headers: Vec<(String, String)>,
17 : /// The bytes after the head that were read with it
18 : pub rest: Vec<u8>,
19 : }
20 :
21 : impl Request {
22 218 : pub fn header(&self, name: &str) -> Option<&str> {
23 712 : self.headers.iter().find(|(n, _)| n == name).map(|(_, v)| v.as_str())
24 218 : }
25 :
26 : /// The target without its query
27 48 : pub fn path(&self) -> &str {
28 48 : self.target.split('?').next().unwrap_or_default()
29 48 : }
30 : }
31 :
32 : /// Read the head of a request, None if the peer closed the connection before sending one
33 41 : pub async fn read_request<S: AsyncRead + Unpin>(stream: &mut S) -> std::io::Result<Option<Request>> {
34 41 : let invalid = |message: &str| std::io::Error::new(std::io::ErrorKind::InvalidData, message.to_string());
35 41 : let mut buffer = Vec::new();
36 41 : let mut chunk = [0; 4096];
37 : loop {
38 41 : let n = stream.read(&mut chunk).await?;
39 41 : if n == 0 {
40 1 : return Ok(None);
41 40 : }
42 40 : buffer.extend_from_slice(&chunk[..n]);
43 40 : let mut headers = [httparse::EMPTY_HEADER; 64];
44 40 : let mut request = httparse::Request::new(&mut headers);
45 40 : match request.parse(&buffer).map_err(|e| invalid(&e.to_string()))? {
46 40 : httparse::Status::Complete(length) => {
47 : return Ok(Some(Request {
48 40 : method: request.method.unwrap_or_default().to_string(),
49 40 : target: request.path.unwrap_or_default().to_string(),
50 40 : headers: request
51 40 : .headers
52 40 : .iter()
53 206 : .map(|h| (h.name.to_ascii_lowercase(), String::from_utf8_lossy(h.value).into_owned()))
54 40 : .collect(),
55 40 : rest: buffer[length..].to_vec(),
56 : }));
57 : }
58 0 : httparse::Status::Partial if buffer.len() > MAX_HEAD_BYTES => {
59 0 : return Err(invalid("the request head is too big"))
60 : }
61 0 : httparse::Status::Partial => {}
62 : }
63 : }
64 41 : }
65 :
66 : /// Read the body of a request with a Content-Length, up to a limit
67 10 : pub async fn read_body<S: AsyncRead + Unpin>(
68 10 : stream: &mut S,
69 10 : request: &Request,
70 10 : limit: usize,
71 10 : ) -> Result<Vec<u8>, String> {
72 10 : let length: usize = match request.header("content-length") {
73 10 : Some(length) => length.trim().parse().map_err(|_| "the content length is invalid".to_string())?,
74 0 : None => 0,
75 : };
76 10 : if length > limit {
77 1 : return Err("the request is too big".into());
78 9 : }
79 9 : let mut body = request.rest.clone();
80 9 : body.truncate(length);
81 9 : let mut remaining = (&mut *stream).take((length - body.len()) as u64);
82 9 : remaining.read_to_end(&mut body).await.map_err(|e| e.to_string())?;
83 9 : if body.len() < length {
84 0 : return Err("the request ended early".into());
85 9 : }
86 9 : Ok(body)
87 10 : }
88 :
89 : /// A response that closes the connection
90 15 : pub fn response(status: u16, reason: &str, content_type: Option<&str>, body: &[u8]) -> Vec<u8> {
91 15 : let content_type =
92 15 : content_type.map(|t| format!("Content-Type: {t}\r\nCache-Control: no-store\r\n")).unwrap_or_default();
93 15 : let mut out = format!(
94 : "HTTP/1.1 {status} {reason}\r\n{content_type}Content-Length: {}\r\nConnection: close\r\n\r\n",
95 15 : body.len()
96 : )
97 15 : .into_bytes();
98 15 : out.extend_from_slice(body);
99 15 : out
100 15 : }
101 :
102 : /// The status and the body of a complete response
103 41 : pub fn parse_response(bytes: &[u8]) -> Option<(u16, &[u8])> {
104 41 : let mut headers = [httparse::EMPTY_HEADER; 64];
105 41 : let mut response = httparse::Response::new(&mut headers);
106 41 : match response.parse(bytes).ok()? {
107 40 : httparse::Status::Complete(length) => Some((response.code?, &bytes[length..])),
108 1 : httparse::Status::Partial => None,
109 : }
110 41 : }
111 :
112 : #[cfg(test)]
113 : mod tests {
114 : use super::*;
115 :
116 : #[tokio::test]
117 1 : async fn reads_requests_with_bodies() {
118 1 : let mut input: &[u8] = b"POST /v1/call?x HTTP/1.1\r\nHost: 127.0.0.1:1\r\nContent-Length: 5\r\n\r\nhello";
119 1 : let request = read_request(&mut input).await.unwrap().unwrap();
120 1 : assert_eq!((request.method.as_str(), request.path()), ("POST", "/v1/call"));
121 1 : assert_eq!(request.header("host"), Some("127.0.0.1:1"));
122 1 : assert_eq!(read_body(&mut input, &request, 100).await.unwrap(), b"hello");
123 1 : assert_eq!(read_body(&mut input, &request, 4).await.unwrap_err(), "the request is too big");
124 1 : assert!(read_request(&mut &b""[..]).await.unwrap().is_none());
125 1 : }
126 :
127 : #[test]
128 1 : fn writes_and_parses_responses() {
129 1 : let bytes = response(404, "Not Found", Some("application/json"), b"{}");
130 1 : assert_eq!(parse_response(&bytes), Some((404, &b"{}"[..])));
131 1 : assert_eq!(parse_response(b"HTTP/1.1 200 OK\r\n"), None);
132 1 : }
133 : }
|
