1
0
mirror of https://github.com/actix/actix-extras.git synced 2024-11-25 00:12:59 +01:00
actix-extras/src/server/h1writer.rs

285 lines
9.2 KiB
Rust
Raw Normal View History

2018-02-26 23:33:56 +01:00
#![cfg_attr(feature = "cargo-clippy", allow(redundant_field_names))]
2018-02-20 02:21:04 +01:00
use std::{io, mem};
2018-03-18 19:05:44 +01:00
use std::rc::Rc;
2017-12-25 22:40:06 +01:00
use bytes::BufMut;
use futures::{Async, Poll};
use tokio_io::AsyncWrite;
2018-01-21 01:12:38 +01:00
use http::{Method, Version};
2018-03-18 19:05:44 +01:00
use http::header::{HeaderValue, CONNECTION, CONTENT_LENGTH, DATE};
2017-12-14 07:54:52 +01:00
use helpers;
2018-01-14 22:50:38 +01:00
use body::{Body, Binary};
use headers::ContentEncoding;
use httprequest::HttpInnerMessage;
use httpresponse::HttpResponse;
2018-01-12 03:41:33 +01:00
use super::{Writer, WriterState, MAX_WRITE_BUFFER_SIZE};
2018-01-15 02:00:28 +01:00
use super::shared::SharedBytes;
2018-02-25 09:43:00 +01:00
use super::encoding::ContentEncoder;
2018-03-18 19:05:44 +01:00
use super::settings::WorkerSettings;
const AVERAGE_HEADER_SIZE: usize = 30; // totally scientific
bitflags! {
struct Flags: u8 {
const STARTED = 0b0000_0001;
const UPGRADE = 0b0000_0010;
const KEEPALIVE = 0b0000_0100;
const DISCONNECTED = 0b0000_1000;
}
}
2018-03-18 19:05:44 +01:00
pub(crate) struct H1Writer<T: AsyncWrite, H: 'static> {
flags: Flags,
2017-12-14 01:44:35 +01:00
stream: T,
2018-02-25 09:43:00 +01:00
encoder: ContentEncoder,
2017-11-10 21:29:54 +01:00
written: u64,
headers_size: u32,
2017-12-15 04:34:31 +01:00
buffer: SharedBytes,
2018-03-09 19:00:15 +01:00
buffer_capacity: usize,
2018-03-18 19:05:44 +01:00
settings: Rc<WorkerSettings<H>>,
}
2018-03-18 19:05:44 +01:00
impl<T: AsyncWrite, H: 'static> H1Writer<T, H> {
2018-03-18 19:05:44 +01:00
pub fn new(stream: T, buf: SharedBytes, settings: Rc<WorkerSettings<H>>)
-> H1Writer<T, H>
{
H1Writer {
flags: Flags::empty(),
2018-02-25 09:43:00 +01:00
encoder: ContentEncoder::empty(buf.clone()),
2017-11-10 21:29:54 +01:00
written: 0,
headers_size: 0,
2017-12-15 04:34:31 +01:00
buffer: buf,
2018-03-09 19:00:15 +01:00
buffer_capacity: 0,
2018-02-26 23:33:56 +01:00
stream,
2018-03-18 19:05:44 +01:00
settings,
}
}
pub fn get_mut(&mut self) -> &mut T {
2017-12-14 01:44:35 +01:00
&mut self.stream
}
2017-12-14 01:44:35 +01:00
pub fn reset(&mut self) {
self.written = 0;
self.flags = Flags::empty();
}
pub fn disconnected(&mut self) {
2018-01-15 02:00:28 +01:00
self.buffer.take();
}
pub fn keepalive(&self) -> bool {
self.flags.contains(Flags::KEEPALIVE) && !self.flags.contains(Flags::UPGRADE)
}
2018-03-10 19:42:46 +01:00
fn write_data(&mut self, data: &[u8]) -> io::Result<usize> {
let mut written = 0;
while written < data.len() {
match self.stream.write(&data[written..]) {
2018-01-15 04:28:34 +01:00
Ok(0) => {
self.disconnected();
2018-03-10 19:42:46 +01:00
return Err(io::Error::new(io::ErrorKind::WriteZero, ""))
2017-12-14 01:44:35 +01:00
},
Ok(n) => written += n,
2017-12-14 01:44:35 +01:00
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
2018-03-10 19:42:46 +01:00
return Ok(written)
}
2017-12-14 01:44:35 +01:00
Err(err) => return Err(err),
}
}
2018-03-10 19:42:46 +01:00
Ok(written)
}
}
2018-03-18 19:05:44 +01:00
impl<T: AsyncWrite, H: 'static> Writer for H1Writer<T, H> {
2018-01-12 01:22:27 +01:00
#[inline]
2017-11-10 21:29:54 +01:00
fn written(&self) -> u64 {
2018-01-02 23:53:51 +01:00
self.written
2017-11-10 21:29:54 +01:00
}
fn start(&mut self,
req: &mut HttpInnerMessage,
msg: &mut HttpResponse,
encoding: ContentEncoding) -> io::Result<WriterState>
{
// prepare task
2018-02-25 09:43:00 +01:00
self.encoder = ContentEncoder::for_server(self.buffer.clone(), req, msg, encoding);
if msg.keep_alive().unwrap_or_else(|| req.keep_alive()) {
2018-02-10 09:05:20 +01:00
self.flags.insert(Flags::STARTED | Flags::KEEPALIVE);
} else {
self.flags.insert(Flags::STARTED);
}
// Connection upgrade
2017-12-09 13:33:40 +01:00
let version = msg.version().unwrap_or_else(|| req.version);
if msg.upgrade() {
2018-02-10 09:05:20 +01:00
self.flags.insert(Flags::UPGRADE);
2017-12-08 02:38:18 +01:00
msg.headers_mut().insert(CONNECTION, HeaderValue::from_static("upgrade"));
}
// keep-alive
else if self.flags.contains(Flags::KEEPALIVE) {
if version < Version::HTTP_11 {
2017-12-08 02:38:18 +01:00
msg.headers_mut().insert(CONNECTION, HeaderValue::from_static("keep-alive"));
}
} else if version >= Version::HTTP_11 {
2017-12-08 02:38:18 +01:00
msg.headers_mut().insert(CONNECTION, HeaderValue::from_static("close"));
}
2017-12-16 16:29:15 +01:00
let body = msg.replace_body(Body::Empty);
// render message
2017-11-09 01:44:23 +01:00
{
2018-01-15 02:00:28 +01:00
let mut buffer = self.buffer.get_mut();
2018-03-18 19:05:44 +01:00
let mut is_bin = if let Body::Binary(ref bytes) = body {
buffer.reserve(
256 + msg.headers().len() * AVERAGE_HEADER_SIZE + bytes.len());
true
2017-11-09 01:44:23 +01:00
} else {
2017-12-15 04:34:31 +01:00
buffer.reserve(256 + msg.headers().len() * AVERAGE_HEADER_SIZE);
2018-03-18 19:05:44 +01:00
false
};
2017-12-25 22:40:06 +01:00
// status line
helpers::write_status_line(version, msg.status().as_u16(), &mut buffer);
2018-02-20 02:21:04 +01:00
SharedBytes::extend_from_slice_(buffer, msg.reason().as_bytes());
2017-12-16 16:29:15 +01:00
match body {
2017-12-26 04:42:55 +01:00
Body::Empty =>
2018-01-21 01:12:38 +01:00
if req.method != Method::HEAD {
2018-03-18 19:05:44 +01:00
SharedBytes::put_slice(
buffer, b"\r\ncontent-length: 0\r\n");
2018-01-21 01:12:38 +01:00
} else {
2018-03-18 19:05:44 +01:00
SharedBytes::put_slice(buffer, b"\r\n");
2018-01-21 01:12:38 +01:00
},
2017-12-26 04:42:55 +01:00
Body::Binary(ref bytes) =>
helpers::write_content_length(bytes.len(), &mut buffer),
_ =>
2018-03-18 19:05:44 +01:00
SharedBytes::put_slice(buffer, b"\r\n"),
2017-12-16 16:29:15 +01:00
}
2017-12-25 22:40:06 +01:00
// write headers
2018-02-20 02:21:04 +01:00
let mut pos = 0;
2018-03-18 19:05:44 +01:00
let mut has_date = false;
2018-02-20 02:21:04 +01:00
let mut remaining = buffer.remaining_mut();
let mut buf: &mut [u8] = unsafe{ mem::transmute(buffer.bytes_mut()) };
2017-12-25 22:40:06 +01:00
for (key, value) in msg.headers() {
2018-03-18 19:05:44 +01:00
if is_bin && key == CONTENT_LENGTH {
is_bin = false;
continue
}
has_date = has_date || key == DATE;
2017-12-25 22:40:06 +01:00
let v = value.as_ref();
let k = key.as_str().as_bytes();
2018-02-20 02:21:04 +01:00
let len = k.len() + v.len() + 4;
if len > remaining {
unsafe{buffer.advance_mut(pos)};
pos = 0;
buffer.reserve(len);
remaining = buffer.remaining_mut();
buf = unsafe{ mem::transmute(buffer.bytes_mut()) };
}
buf[pos..pos+k.len()].copy_from_slice(k);
pos += k.len();
buf[pos..pos+2].copy_from_slice(b": ");
pos += 2;
buf[pos..pos+v.len()].copy_from_slice(v);
pos += v.len();
buf[pos..pos+2].copy_from_slice(b"\r\n");
pos += 2;
remaining -= len;
//buffer.put_slice(k);
//buffer.put_slice(b": ");
//buffer.put_slice(v);
//buffer.put_slice(b"\r\n");
2017-12-25 22:40:06 +01:00
}
2018-02-20 02:21:04 +01:00
unsafe{buffer.advance_mut(pos)};
2017-12-25 22:40:06 +01:00
2018-03-18 19:05:44 +01:00
// optimized date header
if !has_date {
self.settings.set_date(&mut buffer);
2017-12-25 22:40:06 +01:00
} else {
// msg eof
2018-02-20 02:21:04 +01:00
SharedBytes::extend_from_slice_(buffer, b"\r\n");
2017-11-09 01:44:23 +01:00
}
self.headers_size = buffer.len() as u32;
}
2017-12-16 16:29:15 +01:00
if let Body::Binary(bytes) = body {
2018-01-02 23:53:51 +01:00
self.written = bytes.len() as u64;
2018-01-14 23:40:39 +01:00
self.encoder.write(bytes)?;
2017-12-16 16:29:15 +01:00
} else {
2018-03-09 19:00:15 +01:00
// capacity, makes sense only for streaming or actor
self.buffer_capacity = msg.write_buffer_capacity();
2017-12-16 16:29:15 +01:00
msg.replace_body(body);
2017-11-09 01:44:23 +01:00
}
Ok(WriterState::Done)
}
2018-01-14 22:50:38 +01:00
fn write(&mut self, payload: Binary) -> io::Result<WriterState> {
2018-01-02 23:53:51 +01:00
self.written += payload.len() as u64;
if !self.flags.contains(Flags::DISCONNECTED) {
if self.flags.contains(Flags::STARTED) {
2018-02-10 09:05:20 +01:00
// shortcut for upgraded connection
if self.flags.contains(Flags::UPGRADE) {
if self.buffer.is_empty() {
2018-03-10 19:42:46 +01:00
let n = self.write_data(payload.as_ref())?;
if payload.len() < n {
self.buffer.extend_from_slice(&payload.as_ref()[n..]);
return Ok(WriterState::Done);
2018-02-10 09:05:20 +01:00
}
} else {
self.buffer.extend(payload);
}
} else {
// TODO: add warning, write after EOF
self.encoder.write(payload)?;
}
} else {
2017-11-09 01:44:23 +01:00
// might be response to EXCEPT
2018-01-15 02:00:28 +01:00
self.buffer.extend_from_slice(payload.as_ref())
}
}
if self.buffer.len() > self.buffer_capacity {
Ok(WriterState::Pause)
} else {
Ok(WriterState::Done)
}
}
2018-01-14 22:50:38 +01:00
fn write_eof(&mut self) -> io::Result<WriterState> {
2017-11-09 01:44:23 +01:00
self.encoder.write_eof()?;
if !self.encoder.is_eof() {
Err(io::Error::new(io::ErrorKind::Other,
"Last payload item, but eof is not reached"))
2018-01-15 02:00:28 +01:00
} else if self.buffer.len() > MAX_WRITE_BUFFER_SIZE {
Ok(WriterState::Pause)
} else {
Ok(WriterState::Done)
}
}
2018-01-12 01:22:27 +01:00
#[inline]
fn poll_completed(&mut self, shutdown: bool) -> Poll<(), io::Error> {
if !self.buffer.is_empty() {
let buf: &[u8] = unsafe{mem::transmute(self.buffer.as_ref())};
2018-03-10 19:42:46 +01:00
let written = self.write_data(buf)?;
let _ = self.buffer.split_to(written);
if self.buffer.len() > self.buffer_capacity {
return Ok(Async::NotReady)
}
}
if shutdown {
self.stream.shutdown()
} else {
Ok(Async::Ready(()))
}
}
}