1 // Copyright 2023 Google LLC
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License");
4 // you may not use this file except in compliance with the License.
5 // You may obtain a copy of the License at
6 //
7 // https://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS,
11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 // See the License for the specific language governing permissions and
13 // limitations under the License.
14
15 //! Server Response Writer module for micro HTTP server.
16 //!
17 //! This module implements a basic response writer that can pass
18 //! chunked http responses between a uri handler and the network.
19 //!
20 //! The main use is for streaming large files from the capture_handler()
21 //! to the Http client.
22 //!
23 //! This library is intended solely for serving netsim clients.
24
25 use std::io::Write;
26 use std::str::FromStr;
27
28 use http::{HeaderName, HeaderValue, Response};
29 use log::error;
30
31 pub type ResponseWritable<'a> = &'a mut dyn ServerResponseWritable;
32
33 pub type StrHeaders = Vec<(String, String)>;
34
35 // ServerResponseWritable trait is used by both the Http and gRPC
36 // servers.
37 pub trait ServerResponseWritable {
put_ok_with_length(&mut self, mime_type: &str, length: usize, headers: StrHeaders)38 fn put_ok_with_length(&mut self, mime_type: &str, length: usize, headers: StrHeaders);
put_chunk(&mut self, chunk: &[u8])39 fn put_chunk(&mut self, chunk: &[u8]);
put_ok(&mut self, mime_type: &str, body: &str, headers: StrHeaders)40 fn put_ok(&mut self, mime_type: &str, body: &str, headers: StrHeaders);
put_error(&mut self, error_code: u16, error_message: &str)41 fn put_error(&mut self, error_code: u16, error_message: &str);
put_ok_with_vec(&mut self, mime_type: &str, body: Vec<u8>, headers: StrHeaders)42 fn put_ok_with_vec(&mut self, mime_type: &str, body: Vec<u8>, headers: StrHeaders);
put_ok_switch_protocol(&mut self, connection: &str, headers: StrHeaders)43 fn put_ok_switch_protocol(&mut self, connection: &str, headers: StrHeaders);
44 }
45
46 // A response writer that can contain a TCP stream or other writable.
47 pub struct ServerResponseWriter<'a> {
48 writer: &'a mut dyn Write,
49 response: Option<Response<Vec<u8>>>,
50 }
51
52 impl<'a> ServerResponseWriter<'a> {
new<W: Write>(writer: &mut W) -> ServerResponseWriter53 pub fn new<W: Write>(writer: &mut W) -> ServerResponseWriter {
54 ServerResponseWriter { writer, response: None }
55 }
put_response(&mut self, response: Response<Vec<u8>>)56 pub fn put_response(&mut self, response: Response<Vec<u8>>) {
57 let reason = match response.status().as_u16() {
58 101 => "Switching Protocols",
59 200 => "OK",
60 404 => "Not Found",
61 _ => "Unknown Reason",
62 };
63 let mut buffer =
64 format!("HTTP/1.1 {} {}\r\n", response.status().as_str(), reason).into_bytes();
65 for (name, value) in response.headers() {
66 buffer
67 .extend_from_slice(format!("{}: {}\r\n", name, value.to_str().unwrap()).as_bytes());
68 }
69 buffer.extend_from_slice(b"\r\n");
70 buffer.extend_from_slice(response.body());
71 if let Err(e) = self.writer.write_all(&buffer) {
72 error!("handle_connection error {e}");
73 };
74 self.response = Some(response);
75 }
get_response(self) -> Option<Response<Vec<u8>>>76 pub fn get_response(self) -> Option<Response<Vec<u8>>> {
77 self.response
78 }
79 }
80
81 // Implement the ServerResponseWritable trait for the
82 // ServerResponseWriter struct. These methods are called
83 // by the handler methods.
84 impl ServerResponseWritable for ServerResponseWriter<'_> {
put_error(&mut self, error_code: u16, error_message: &str)85 fn put_error(&mut self, error_code: u16, error_message: &str) {
86 let body = error_message.as_bytes().to_vec();
87 self.put_response(
88 Response::builder()
89 .status(error_code)
90 .header("Content-Type", HeaderValue::from_static("text/plain"))
91 .header(
92 "Content-Length",
93 HeaderValue::from_str(body.len().to_string().as_str()).unwrap(),
94 )
95 .body(body)
96 .unwrap(),
97 );
98 }
put_chunk(&mut self, chunk: &[u8])99 fn put_chunk(&mut self, chunk: &[u8]) {
100 if let Err(e) = self.writer.write_all(chunk) {
101 error!("handle_connection error {e}");
102 };
103 self.writer.flush().unwrap();
104 }
put_ok_with_length(&mut self, mime_type: &str, length: usize, headers: StrHeaders)105 fn put_ok_with_length(&mut self, mime_type: &str, length: usize, headers: StrHeaders) {
106 let mut response = Response::builder()
107 .status(200)
108 .header("Content-Type", HeaderValue::from_str(mime_type).unwrap())
109 .header("Content-Length", HeaderValue::from_str(length.to_string().as_str()).unwrap())
110 .body(Vec::new())
111 .unwrap();
112 add_headers(&mut response, headers);
113 self.put_response(response);
114 }
put_ok(&mut self, mime_type: &str, body: &str, headers: StrHeaders)115 fn put_ok(&mut self, mime_type: &str, body: &str, headers: StrHeaders) {
116 let mut response = new_ok(mime_type, body.into());
117 add_headers(&mut response, headers);
118 self.put_response(response);
119 }
put_ok_with_vec(&mut self, mime_type: &str, body: Vec<u8>, headers: StrHeaders)120 fn put_ok_with_vec(&mut self, mime_type: &str, body: Vec<u8>, headers: StrHeaders) {
121 let mut response = new_ok(mime_type, body);
122 add_headers(&mut response, headers);
123 self.put_response(response);
124 }
put_ok_switch_protocol(&mut self, connection: &str, headers: StrHeaders)125 fn put_ok_switch_protocol(&mut self, connection: &str, headers: StrHeaders) {
126 let mut response = Response::builder()
127 .status(101)
128 .header("Upgrade", HeaderValue::from_str(connection).unwrap())
129 .header("Connection", HeaderValue::from_static("Upgrade"))
130 .body(Vec::new())
131 .unwrap();
132 add_headers(&mut response, headers);
133 self.put_response(response);
134 }
135 }
136
new_ok(content_type: &str, body: Vec<u8>) -> Response<Vec<u8>>137 fn new_ok(content_type: &str, body: Vec<u8>) -> Response<Vec<u8>> {
138 Response::builder()
139 .status(200)
140 .header("Content-Type", HeaderValue::from_str(content_type).unwrap())
141 .header("Content-Length", HeaderValue::from_str(body.len().to_string().as_str()).unwrap())
142 .body(body)
143 .unwrap()
144 }
145
add_headers(response: &mut Response<Vec<u8>>, headers: StrHeaders)146 fn add_headers(response: &mut Response<Vec<u8>>, headers: StrHeaders) {
147 for (key, value) in headers {
148 response.headers_mut().insert(
149 HeaderName::from_str(key.as_str()).unwrap(),
150 HeaderValue::from_str(value.as_str()).unwrap(),
151 );
152 }
153 }
154
155 #[cfg(test)]
156 mod tests {
157 use super::*;
158 use std::io::Cursor;
159
160 const SAMPLE_CHUNK: &[u8] = &[0, 1, 2, 3, 4];
161
162 #[test]
test_put_error()163 fn test_put_error() {
164 let mut stream = Cursor::new(Vec::new());
165 let mut writer = ServerResponseWriter::new(&mut stream);
166 writer.put_error(404, "Hello World");
167 let written_bytes = stream.get_ref();
168 let expected_bytes =
169 b"HTTP/1.1 404 Not Found\r\ncontent-type: text/plain\r\ncontent-length: 11\r\n\r\nHello World";
170 assert_eq!(written_bytes, expected_bytes);
171 }
172
173 #[test]
test_put_ok()174 fn test_put_ok() {
175 let mut stream = Cursor::new(Vec::new());
176 let mut writer = ServerResponseWriter::new(&mut stream);
177 writer.put_ok("text/plain", "Hello World", vec![]);
178 let written_bytes = stream.get_ref();
179 let expected_bytes =
180 b"HTTP/1.1 200 OK\r\ncontent-type: text/plain\r\ncontent-length: 11\r\n\r\nHello World";
181 assert_eq!(written_bytes, expected_bytes);
182 }
183
184 #[test]
test_put_ok_with_length()185 fn test_put_ok_with_length() {
186 let mut stream = Cursor::new(Vec::new());
187 let mut writer = ServerResponseWriter::new(&mut stream);
188 writer.put_ok_with_length("text/plain", 100, vec![]);
189 let written_bytes = stream.get_ref();
190 let expected_bytes =
191 b"HTTP/1.1 200 OK\r\ncontent-type: text/plain\r\ncontent-length: 100\r\n\r\n";
192 assert_eq!(written_bytes, expected_bytes);
193 }
194
195 #[test]
test_put_ok_with_vec()196 fn test_put_ok_with_vec() {
197 let mut stream = Cursor::new(Vec::new());
198 let mut writer = ServerResponseWriter::new(&mut stream);
199 writer.put_ok_with_vec("text/plain", SAMPLE_CHUNK.to_vec(), vec![]);
200 let written_bytes = stream.get_ref();
201 let expected_bytes = &[
202 b"HTTP/1.1 200 OK\r\ncontent-type: text/plain\r\ncontent-length: 5\r\n\r\n".to_vec(),
203 SAMPLE_CHUNK.to_vec(),
204 ]
205 .concat();
206 assert_eq!(written_bytes, expected_bytes);
207 }
208
209 #[test]
test_put_ok_switch_protocol()210 fn test_put_ok_switch_protocol() {
211 let mut stream = Cursor::new(Vec::new());
212 let mut writer = ServerResponseWriter::new(&mut stream);
213 writer.put_ok_switch_protocol("Websocket", vec![]);
214 let written_bytes = stream.get_ref();
215 let expected_bytes = b"HTTP/1.1 101 Switching Protocols\r\nupgrade: Websocket\r\nconnection: Upgrade\r\n\r\n";
216 assert_eq!(written_bytes, expected_bytes);
217 }
218
219 #[test]
test_put_chunk()220 fn test_put_chunk() {
221 let mut stream = Cursor::new(Vec::new());
222 let mut writer = ServerResponseWriter::new(&mut stream);
223 writer.put_chunk(SAMPLE_CHUNK);
224 let written_bytes = stream.get_ref();
225 let expected_bytes = SAMPLE_CHUNK;
226 assert_eq!(written_bytes, expected_bytes);
227 }
228 }
229