use base64;
use body::{send_body, Payload, SizedReader};
use std::io::{Result as IoResult, Write};
use stream::{connect_http, connect_https, connect_test, Stream};
use url::Url;
use pool::DEFAULT_HOST;
#[derive(Debug)]
pub struct Unit {
pub agent: Arc<Mutex<Option<AgentState>>>,
pub url: Url,
pub is_chunked: bool,
pub is_head: bool,
pub hostname: String,
pub query_string: String,
pub headers: Vec<Header>,
pub timeout_connect: u64,
pub timeout_read: u64,
pub timeout_write: u64,
}
impl Unit {
fn new(req: &Request, url: &Url, body: &SizedReader) -> Self {
let is_chunked = req.header("transfer-encoding")
.map(|enc| enc.len() > 0)
.unwrap_or(false);
let is_secure = url.scheme().eq_ignore_ascii_case("https");
let is_head = req.method.eq_ignore_ascii_case("head");
let hostname = url.host_str().unwrap_or(DEFAULT_HOST).to_string();
let query_string = combine_query(&url, &req.query);
let cookie_headers: Vec<_> = {
let mut state = req.agent.lock().unwrap();
match state.as_ref().map(|state| &state.jar) {
None => vec![],
Some(jar) => match_cookies(jar, &hostname, url.path(), is_secure),
}
};
let extra_headers = {
let mut extra = vec![];
if !is_chunked && !req.has("content-length") {
if let Some(size) = body.size {
extra.push(Header::new("Content-Length", &format!("{}", size)));
}
}
let username = url.username();
let password = url.password().unwrap_or("");
if (username != "" || password != "") && !req.has("authorization") {
let encoded = base64::encode(&format!("{}:{}", username, password));
extra.push(Header::new("Authorization", &format!("Basic {}", encoded)));
}
extra
};
let headers: Vec<_> = req
.headers
.iter()
.chain(cookie_headers.iter())
.chain(extra_headers.iter())
.cloned()
.collect();
Unit {
agent: Arc::clone(&req.agent),
url: url.clone(),
is_chunked,
is_head,
hostname,
query_string,
headers,
timeout_connect: req.timeout_connect,
timeout_read: req.timeout_read,
timeout_write: req.timeout_write,
}
}
#[cfg(test)]
pub fn header<'a>(&self, name: &'a str) -> Option<&str> {
get_header(&self.headers, name)
}
#[cfg(test)]
pub fn has<'a>(&self, name: &'a str) -> bool {
has_header(&self.headers, name)
}
#[cfg(test)]
pub fn all<'a>(&self, name: &'a str) -> Vec<&str> {
get_all_headers(&self.headers, name)
}
}
pub fn connect(
mut unit: Unit,
method: &str,
use_pooled: bool,
redirects: u32,
body: SizedReader,
) -> Result<Response, Error> {
let (mut stream, is_recycled) = connect_socket(&unit, use_pooled)?;
let send_result = send_prelude(&unit, method, &mut stream);
if send_result.is_err() {
if is_recycled {
return connect(unit, method, false, redirects, body);
} else {
return Err(send_result.unwrap_err().into());
}
}
send_body(body, unit.is_chunked, &mut stream)?;
let mut resp = Response::from_read(&mut stream);
save_cookies(&unit, &resp);
if resp.redirect() {
if redirects == 0 {
return Err(Error::TooManyRedirects);
}
let location = resp.header("location");
if let Some(location) = location {
let new_url = unit
.url
.join(location)
.map_err(|_| Error::BadUrl(format!("Bad redirection: {}", location)))?;
unit.url = new_url;
match resp.status() {
301 | 302 | 303 => {
let empty = Payload::Empty.into_read();
return connect(unit, "GET", use_pooled, redirects - 1, empty);
}
, _ => (),
};
}
}
response::set_stream(&mut resp, Some(unit), stream);
Ok(resp)
}
fn match_cookies<'a>(jar: &'a CookieJar, domain: &str, path: &str, is_secure: bool) -> Vec<Header> {
jar.iter()
.filter(|c| {
let domain_ok = c
.domain()
.map(|cdom| domain.contains(cdom))
.unwrap_or(false);
let path_ok = c
.path()
.map(|cpath| path.find(cpath).map(|pos| pos == 0).unwrap_or(false))
.unwrap_or(true);
let secure_ok = !c.secure().unwrap_or(false) || is_secure;
domain_ok && path_ok && secure_ok
})
.map(|c| {
let name = c.name().to_string();
let value = c.value().to_string();
let nameval = Cookie::new(name, value).encoded().to_string();
Header::new("Cookie", &nameval)
})
.collect()
}
fn combine_query(url: &Url, query: &QString) -> String {
match (url.query(), query.len() > 0) {
(Some(urlq), true) => format!("?{}&{}", urlq, query),
(Some(urlq), false) => format!("?{}", urlq),
(None, true) => format!("?{}", query),
(None, false) => "".to_string(),
}
}
fn connect_socket(unit: &Unit, use_pooled: bool) -> Result<(Stream, bool), Error> {
if use_pooled {
let state = &mut unit.agent.lock().unwrap();
if let Some(agent) = state.as_mut() {
if let Some(stream) = agent.pool.try_get_connection(&unit.url) {
return Ok((stream, true));
}
}
}
let stream = match unit.url.scheme() {
"http" => connect_http(&unit),
"https" => connect_https(&unit),
"test" => connect_test(&unit),
_ => Err(Error::UnknownScheme(unit.url.scheme().to_string())),
};
Ok((stream?, false))
}
fn send_prelude(unit: &Unit, method: &str, stream: &mut Stream) -> IoResult<()> {
let mut prelude: Vec<u8> = vec![];
write!(
prelude,
"{} {}{} HTTP/1.1\r\n",
method,
unit.url.path(),
&unit.query_string
)?;
if !has_header(&unit.headers, "host") {
write!(prelude, "Host: {}\r\n", unit.url.host().unwrap())?;
}
for header in &unit.headers {
write!(prelude, "{}: {}\r\n", header.name(), header.value())?;
}
write!(prelude, "\r\n")?;
stream.write_all(&mut prelude[..])?;
Ok(())
}
fn save_cookies(unit: &Unit, resp: &Response) {
let cookies = resp.all("set-cookie");
if cookies.is_empty() {
return;
}
let state = &mut unit.agent.lock().unwrap();
if let Some(add_jar) = state.as_mut().map(|state| &mut state.jar) {
for raw_cookie in cookies.iter() {
let to_parse = if raw_cookie.to_lowercase().contains("domain=") {
raw_cookie.to_string()
} else {
format!("{}; Domain={}", raw_cookie, &unit.hostname)
};
match Cookie::parse_encoded(&to_parse[..]) {
Err(_) => (), Ok(mut cookie) => {
let cookie = cookie.into_owned();
add_jar.add(cookie)
}
}
}
}
}