set unit on response

This commit is contained in:
Martin Algesten
2018-06-30 14:37:48 +02:00
parent 3a82ad05ab
commit 317f75f8bf
3 changed files with 108 additions and 107 deletions

View File

@@ -93,8 +93,8 @@ impl Request {
self.to_url()
.and_then(|url| {
let reader = payload.into_read();
let mut unit = Unit::new(&self, &url, &reader);
unit.connect(url, &self.method, self.redirects, reader)
let unit = Unit::new(&self, &url, &reader);
connect(unit, url, &self.method, self.redirects, reader)
})
.unwrap_or_else(|e| e.into())
}

View File

@@ -1,3 +1,4 @@
use agent::Unit;
use ascii::AsciiString;
use chunked_transfer;
use header::Header;
@@ -44,7 +45,7 @@ pub struct Response {
index: (usize, usize), // index into status_line where we split: HTTP/1.1 200 OK
status: u16,
headers: Vec<Header>,
is_head: bool,
unit: Option<Unit>,
stream: Option<Stream>,
}
@@ -260,7 +261,8 @@ impl Response {
let reader = self.stream.expect("No reader in response?!");
// head requests never have a body
if self.is_head {
let is_head = self.unit.map(|u| u.is_head).unwrap_or(false);
if is_head {
return Box::new(LimitedRead::new(reader, 0)) as Box<Read>;
}
@@ -389,7 +391,7 @@ impl Response {
index,
status,
headers,
is_head: false,
unit: None,
stream: None,
})
}
@@ -432,7 +434,7 @@ impl FromStr for Response {
let bytes = s.as_bytes().to_owned();
let mut cursor = Cursor::new(bytes);
let mut resp = Self::do_from_read(&mut cursor)?;
set_stream(&mut resp, Stream::Cursor(cursor), false);
set_stream(&mut resp, None, Stream::Cursor(cursor));
Ok(resp)
}
}
@@ -448,8 +450,8 @@ impl Into<Response> for Error {
}
}
pub fn set_stream(resp: &mut Response, stream: Stream, is_head: bool) {
resp.is_head = is_head;
pub fn set_stream(resp: &mut Response, unit: Option<Unit>, stream: Stream) {
resp.unit = unit;
resp.stream = Some(stream);
}

View File

@@ -1,7 +1,7 @@
use url::Url;
use stream::{connect_http, connect_https, connect_test};
use body::{send_body, Payload, SizedReader};
use std::io::Write;
use body::{Payload, SizedReader, send_body};
use stream::{connect_http, connect_https, connect_test};
use url::Url;
//
pub struct Unit {
@@ -82,8 +82,22 @@ impl Unit {
}
}
fn connect(
&mut self,
#[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(
unit: Unit,
url: Url,
method: &str,
redirects: u32,
@@ -93,9 +107,9 @@ impl Unit {
// open socket
let mut stream = match url.scheme() {
"http" => connect_http(self),
"https" => connect_https(self),
"test" => connect_test(self),
"http" => connect_http(&unit),
"https" => connect_https(&unit),
"test" => connect_test(&unit),
_ => Err(Error::UnknownScheme(url.scheme().to_string())),
}?;
@@ -106,12 +120,12 @@ impl Unit {
"{} {}{} HTTP/1.1\r\n",
method,
url.path(),
self.query_string
&unit.query_string
)?;
if !has_header(&self.headers, "host") {
if !has_header(&unit.headers, "host") {
write!(prelude, "Host: {}\r\n", url.host().unwrap())?;
}
for header in &self.headers {
for header in &unit.headers {
write!(prelude, "{}: {}\r\n", header.name(), header.value())?;
}
write!(prelude, "\r\n")?;
@@ -123,13 +137,13 @@ impl Unit {
// squirrel away cookies
{
let mut state = self.agent.lock().unwrap();
let state = &mut unit.agent.lock().unwrap();
if let Some(add_jar) = state.as_mut().map(|state| &mut state.jar) {
for raw_cookie in resp.all("set-cookie").iter() {
let to_parse = if raw_cookie.to_lowercase().contains("domain=") {
raw_cookie.to_string()
} else {
format!("{}; Domain={}", raw_cookie, self.hostname)
format!("{}; Domain={}", raw_cookie, &unit.hostname)
};
match Cookie::parse_encoded(&to_parse[..]) {
Err(_) => (), // ignore unparseable cookies
@@ -159,40 +173,25 @@ impl Unit {
// perform the redirect differently depending on 3xx code.
return match resp.status() {
301 | 302 | 303 => {
send_body(body, self.is_chunked, &mut stream)?;
send_body(body, unit.is_chunked, &mut stream)?;
let empty = Payload::Empty.into_read();
self.connect(new_url, "GET", redirects - 1, empty)
connect(unit, new_url, "GET", redirects - 1, empty)
}
307 | 308 | _ => self.connect(new_url, method, redirects - 1, body),
307 | 308 | _ => connect(unit, new_url, method, redirects - 1, body),
};
}
}
// send the body (which can be empty now depending on redirects)
send_body(body, self.is_chunked, &mut stream)?;
send_body(body, unit.is_chunked, &mut stream)?;
// since it is not a redirect, give away the incoming stream to the response object
response::set_stream(&mut resp, stream, self.is_head);
response::set_stream(&mut resp, Some(unit), stream);
// release the response
Ok(resp)
}
#[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)
}
}
// TODO check so cookies can't be set for tld:s
fn match_cookies<'a>(jar: &'a CookieJar, domain: &str, path: &str, is_secure: bool) -> Vec<Header> {
jar.iter()