use std::borrow::Cow;
use std::error;
use std::fmt;
use std::io;
use std::io::BufRead;
use std::io::Error as IoError;
use std::io::Read;
use std::io::Write;
use std::net::TcpStream;
use std::net::ToSocketAddrs;
use std::time::Duration;
use Request;
use Response;
use ResponseBody;
#[derive(Debug)]
pub enum ProxyError {
BodyAlreadyExtracted,
IoError(IoError),
HttpParseError,
}
impl From<IoError> for ProxyError {
fn from(err: IoError) -> ProxyError {
ProxyError::IoError(err)
}
}
impl error::Error for ProxyError {
#[inline]
fn source(&self) -> Option<&(dyn error::Error + 'static)> {
match *self {
ProxyError::IoError(ref e) => Some(e),
_ => None,
}
}
}
impl fmt::Display for ProxyError {
#[inline]
fn fmt(&self, fmt: &mut fmt::Formatter) -> Result<(), fmt::Error> {
let description = match *self {
ProxyError::BodyAlreadyExtracted => "the body of the request was already extracted",
ProxyError::IoError(_) => {
"could not read the body from the request, or could not connect to the remote \
server, or the connection to the remote server closed unexpectedly"
}
ProxyError::HttpParseError => "the destination server didn't produce compliant HTTP",
};
write!(fmt, "{}", description)
}
}
#[derive(Debug, Clone)]
pub struct ProxyConfig<A> {
pub addr: A,
pub replace_host: Option<Cow<'static, str>>,
}
pub fn proxy<A>(request: &Request, config: ProxyConfig<A>) -> Result<Response, ProxyError>
where
A: ToSocketAddrs,
{
let mut socket = TcpStream::connect(config.addr)?;
socket.set_read_timeout(Some(Duration::from_secs(60)))?;
socket.set_write_timeout(Some(Duration::from_secs(60)))?;
let mut data = match request.data() {
Some(d) => d,
None => return Err(ProxyError::BodyAlreadyExtracted),
};
socket
.write_all(format!("{} {} HTTP/1.1\r\n", request.method(), request.raw_url()).as_bytes())?;
for (header, value) in request.headers() {
let value = if header == "Host" {
if let Some(ref replace) = config.replace_host {
&**replace
} else {
value
}
} else {
value
};
if header == "Connection" {
continue;
}
socket.write_all(format!("{}: {}\r\n", header, value).as_bytes())?;
}
socket.write_all(b"Connection: close\r\n\r\n")?;
io::copy(&mut data, &mut socket)?;
let mut socket = io::BufReader::new(socket);
let mut headers = Vec::new();
let status_code;
{
let mut lines = socket.by_ref().lines();
{
let line = match lines.next() {
Some(l) => l,
None => return Err(ProxyError::HttpParseError),
}?;
let mut splits = line.splitn(3, ' ');
let _ = splits.next();
let status_str = match splits.next() {
Some(l) => l,
None => return Err(ProxyError::HttpParseError),
};
status_code = match status_str.parse() {
Ok(s) => s,
Err(_) => return Err(ProxyError::HttpParseError),
};
}
for header in lines {
let header = header?;
if header.is_empty() {
break;
}
let mut splits = header.splitn(2, ':');
let header = match splits.next() {
Some(v) => v,
None => return Err(ProxyError::HttpParseError),
};
let val = match splits.next() {
Some(v) => v,
None => return Err(ProxyError::HttpParseError),
};
let val = &val[1..];
headers.push((header.to_owned().into(), val.to_owned().into()));
}
}
Ok(Response {
status_code,
headers,
data: ResponseBody::from_reader(socket),
upgrade: None,
})
}
#[derive(Debug)]
pub enum FullProxyError {
BodyAlreadyExtracted,
}
impl error::Error for FullProxyError {}
impl fmt::Display for FullProxyError {
#[inline]
fn fmt(&self, fmt: &mut fmt::Formatter) -> Result<(), fmt::Error> {
let description = match *self {
FullProxyError::BodyAlreadyExtracted => "the body of the request was already extracted",
};
write!(fmt, "{}", description)
}
}
pub fn full_proxy<A>(request: &Request, config: ProxyConfig<A>) -> Result<Response, FullProxyError>
where
A: ToSocketAddrs,
{
match proxy(request, config) {
Ok(r) => Ok(r),
Err(ProxyError::IoError(_)) => Ok(Response::text("Gateway Time-out").with_status_code(504)),
Err(ProxyError::HttpParseError) => Ok(Response::text("Bad Gateway").with_status_code(502)),
Err(ProxyError::BodyAlreadyExtracted) => Err(FullProxyError::BodyAlreadyExtracted),
}
}