use std::io::{self, Read, Write};
use crate::error::HttpError;
use crate::header::Header;
use crate::parse::transfer::Body;
use crate::status;
pub trait ResponseWriter: Send {
fn header(&mut self) -> &mut Header;
fn write(&mut self, buf: &[u8]) -> Result<usize, HttpError>;
fn write_header(&mut self, status_code: u16);
}
pub struct ConnResponseWriter<W: Write + Send> {
pub(crate) inner: W,
header: Header,
status: u16,
header_written: bool,
chunked: bool,
}
impl<W: Write + Send> ConnResponseWriter<W> {
pub fn new(inner: W) -> Self {
let mut header = Header::new();
header.set("Content-Type", "text/plain; charset=utf-8");
Self {
inner,
header,
status: 200,
header_written: false,
chunked: true, }
}
fn flush_headers(&mut self) -> Result<(), HttpError> {
if self.header_written {
return Ok(());
}
self.header_written = true;
let has_cl = self.header.get("Content-Length").is_some();
if !has_cl {
self.header.set("Transfer-Encoding", "chunked");
}
let status_text = status::status_text(self.status);
write!(self.inner, "HTTP/1.1 {} {}\r\n", self.status, status_text)?;
self.header.write_to(&mut self.inner)?;
self.inner.write_all(b"\r\n")?;
Ok(())
}
}
impl<W: Write + Send> ResponseWriter for ConnResponseWriter<W> {
fn header(&mut self) -> &mut Header {
&mut self.header
}
fn write(&mut self, buf: &[u8]) -> Result<usize, HttpError> {
self.flush_headers()?;
if self.chunked && self.header.get("Content-Length").is_none() {
write!(self.inner, "{:X}\r\n", buf.len())?;
self.inner.write_all(buf)?;
self.inner.write_all(b"\r\n")?;
Ok(buf.len())
} else {
self.inner.write_all(buf)?;
Ok(buf.len())
}
}
fn write_header(&mut self, status_code: u16) {
if self.header_written {
return;
}
self.status = status_code;
}
}
impl<W: Write + Send> ConnResponseWriter<W> {
pub fn finish(&mut self) -> Result<(), HttpError> {
self.flush_headers()?;
if self.chunked && self.header.get("Content-Length").is_none() {
self.inner.write_all(b"0\r\n\r\n")?;
}
self.inner.flush()?;
Ok(())
}
}
pub struct Response {
pub status: u16,
pub status_text: String,
pub proto: String,
pub proto_major: u8,
pub proto_minor: u8,
pub header: Header,
pub body: Option<Body>,
pub content_length: i64,
pub transfer_encoding: Vec<String>,
pub trailer: Header,
}
impl Response {
pub fn body_bytes(&mut self) -> Result<Vec<u8>, HttpError> {
let mut out = Vec::new();
if let Some(body) = self.body.take() {
BodyReader(body)
.read_to_end(&mut out)
.map_err(|_| HttpError::BodyRead)?;
}
Ok(out)
}
pub fn body_string(&mut self) -> Result<String, HttpError> {
let bytes = self.body_bytes()?;
String::from_utf8(bytes)
.map_err(|_| HttpError::BodyRead)
}
}
struct BodyReader(Body);
impl Read for BodyReader {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
self.0.read(buf)
}
}
#[cfg(test)]
mod tests {
use super::*;
fn make_writer() -> ConnResponseWriter<Vec<u8>> {
ConnResponseWriter::new(Vec::new())
}
#[test]
fn default_200() {
let mut w = make_writer();
w.write(b"hello").unwrap();
w.finish().unwrap();
let out = String::from_utf8(w.inner).unwrap();
assert!(out.starts_with("HTTP/1.1 200 OK\r\n"));
assert!(out.contains("hello"));
}
#[test]
fn custom_status() {
let mut w = make_writer();
w.write_header(404);
w.write(b"not found").unwrap();
w.finish().unwrap();
let out = String::from_utf8(w.inner).unwrap();
assert!(out.starts_with("HTTP/1.1 404 Not Found\r\n"));
}
#[test]
fn write_header_idempotent() {
let mut w = make_writer();
w.write_header(201);
w.write(b"body").unwrap(); w.write_header(500); w.finish().unwrap();
let out = String::from_utf8(w.inner).unwrap();
assert!(out.starts_with("HTTP/1.1 201"));
}
}