From 317f75f8bfc168a8d8dacb43e2394fb3774ee1c7 Mon Sep 17 00:00:00 2001 From: Martin Algesten Date: Sat, 30 Jun 2018 14:37:48 +0200 Subject: [PATCH] set unit on response --- src/request.rs | 4 +- src/response.rs | 14 ++-- src/unit.rs | 197 ++++++++++++++++++++++++------------------------ 3 files changed, 108 insertions(+), 107 deletions(-) diff --git a/src/request.rs b/src/request.rs index 8b29b41..5f69131 100644 --- a/src/request.rs +++ b/src/request.rs @@ -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()) } diff --git a/src/response.rs b/src/response.rs index 96642a9..a8c3b98 100644 --- a/src/response.rs +++ b/src/response.rs @@ -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
, - is_head: bool, + unit: Option, stream: Option, } @@ -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; } @@ -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 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, stream: Stream) { + resp.unit = unit; resp.stream = Some(stream); } diff --git a/src/unit.rs b/src/unit.rs index d70ccf3..9f95c9c 100644 --- a/src/unit.rs +++ b/src/unit.rs @@ -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,102 +82,6 @@ impl Unit { } } - fn connect( - &mut self, - url: Url, - method: &str, - redirects: u32, - body: SizedReader, - ) -> Result { - // - - // open socket - let mut stream = match url.scheme() { - "http" => connect_http(self), - "https" => connect_https(self), - "test" => connect_test(self), - _ => Err(Error::UnknownScheme(url.scheme().to_string())), - }?; - - // send the request start + headers - let mut prelude: Vec = vec![]; - write!( - prelude, - "{} {}{} HTTP/1.1\r\n", - method, - url.path(), - self.query_string - )?; - if !has_header(&self.headers, "host") { - write!(prelude, "Host: {}\r\n", url.host().unwrap())?; - } - for header in &self.headers { - write!(prelude, "{}: {}\r\n", header.name(), header.value())?; - } - write!(prelude, "\r\n")?; - - stream.write_all(&mut prelude[..])?; - - // start reading the response to process cookies and redirects. - let mut resp = Response::from_read(&mut stream); - - // squirrel away cookies - { - let mut state = self.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) - }; - match Cookie::parse_encoded(&to_parse[..]) { - Err(_) => (), // ignore unparseable cookies - Ok(mut cookie) => { - let cookie = cookie.into_owned(); - add_jar.add(cookie) - } - } - } - } - } - - // handle redirects - if resp.redirect() { - if redirects == 0 { - return Err(Error::TooManyRedirects); - } - - // the location header - let location = resp.header("location"); - if let Some(location) = location { - // join location header to current url in case it it relative - let new_url = url - .join(location) - .map_err(|_| Error::BadUrl(format!("Bad redirection: {}", location)))?; - - // perform the redirect differently depending on 3xx code. - return match resp.status() { - 301 | 302 | 303 => { - send_body(body, self.is_chunked, &mut stream)?; - let empty = Payload::Empty.into_read(); - self.connect(new_url, "GET", redirects - 1, empty) - } - 307 | 308 | _ => self.connect(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)?; - - // since it is not a redirect, give away the incoming stream to the response object - response::set_stream(&mut resp, stream, self.is_head); - - // release the response - Ok(resp) - } - #[cfg(test)] pub fn header<'a>(&self, name: &'a str) -> Option<&str> { get_header(&self.headers, name) @@ -190,7 +94,102 @@ impl Unit { 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, + body: SizedReader, +) -> Result { + // + + // open socket + let mut stream = match url.scheme() { + "http" => connect_http(&unit), + "https" => connect_https(&unit), + "test" => connect_test(&unit), + _ => Err(Error::UnknownScheme(url.scheme().to_string())), + }?; + + // send the request start + headers + let mut prelude: Vec = vec![]; + write!( + prelude, + "{} {}{} HTTP/1.1\r\n", + method, + url.path(), + &unit.query_string + )?; + if !has_header(&unit.headers, "host") { + write!(prelude, "Host: {}\r\n", 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[..])?; + + // start reading the response to process cookies and redirects. + let mut resp = Response::from_read(&mut stream); + + // squirrel away cookies + { + 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, &unit.hostname) + }; + match Cookie::parse_encoded(&to_parse[..]) { + Err(_) => (), // ignore unparseable cookies + Ok(mut cookie) => { + let cookie = cookie.into_owned(); + add_jar.add(cookie) + } + } + } + } + } + + // handle redirects + if resp.redirect() { + if redirects == 0 { + return Err(Error::TooManyRedirects); + } + + // the location header + let location = resp.header("location"); + if let Some(location) = location { + // join location header to current url in case it it relative + let new_url = url + .join(location) + .map_err(|_| Error::BadUrl(format!("Bad redirection: {}", location)))?; + + // perform the redirect differently depending on 3xx code. + return match resp.status() { + 301 | 302 | 303 => { + send_body(body, unit.is_chunked, &mut stream)?; + let empty = Payload::Empty.into_read(); + connect(unit, new_url, "GET", redirects - 1, empty) + } + 307 | 308 | _ => connect(unit, new_url, method, redirects - 1, body), + }; + } + } + + // send the body (which can be empty now depending on redirects) + 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, Some(unit), stream); + + // release the response + Ok(resp) } // TODO check so cookies can't be set for tld:s