This commit is contained in:
Martin Algesten
2018-07-01 11:02:17 +02:00
parent a14e75fcbc
commit 2d2daf58ba
6 changed files with 112 additions and 26 deletions

View File

@@ -40,13 +40,20 @@ include!("unit.rs");
/// ```
#[derive(Debug, Default, Clone)]
pub struct Agent {
/// Copied into each request of this agent.
headers: Vec<Header>,
/// Reused agent state for repeated requests from this agent.
state: Arc<Mutex<Option<AgentState>>>,
}
/// Container of the state
///
/// *Internal API*.
#[derive(Debug)]
pub struct AgentState {
/// Reused connections between requests.
pool: ConnectionPool,
/// Cookies saved between requests.
jar: CookieJar,
}

View File

@@ -4,6 +4,11 @@ use std::io::{Read, Result as IoResult};
use stream::Stream;
use url::Url;
pub const DEFAULT_HOST: &'static str = "localhost";
/// Holder of recycled connections.
///
/// *Internal API*
#[derive(Default, Debug)]
pub struct ConnectionPool {
// the actual pooled connection. however only one per hostname:port.
@@ -17,6 +22,7 @@ impl ConnectionPool {
}
}
/// How the unit::connect tries to get a pooled connection.
pub fn try_get_connection(&mut self, url: &Url) -> Option<Stream> {
self.recycle.remove(&PoolKey::new(url))
}
@@ -45,13 +51,18 @@ struct PoolKey {
impl PoolKey {
fn new(url: &Url) -> Self {
PoolKey {
hostname: url.host_str().unwrap_or("localhost").into(),
hostname: url.host_str().unwrap_or(DEFAULT_HOST).into(),
port: url.port_or_known_default().unwrap_or(0),
}
}
}
/// Read wrapper that returns the stream to the pool once the
/// read is exhausted (reached a 0).
///
/// *Internal API*
pub struct PoolReturnRead<R: Read + Sized> {
// unit that contains the agent where we want to return the reader.
unit: Option<Unit>,
// pointer to underlying stream
stream: *mut Stream,
@@ -69,6 +80,7 @@ impl<R: Read + Sized> PoolReturnRead<R> {
}
fn return_connection(&mut self) {
// guard we only do this once.
if let Some(unit) = self.unit.take() {
// this frees up the wrapper type around the Stream so
// we can safely bring the stream pointer back.

View File

@@ -445,11 +445,6 @@ impl Request {
self
}
/// Returns the method.
pub fn method(&self) -> &str {
&self.method
}
// pub fn retry(&self, times: u16) -> Request {
// unimplemented!()
// }

View File

@@ -251,8 +251,7 @@ impl Response {
//
let is_http10 = self.http_version().eq_ignore_ascii_case("HTTP/1.0");
let is_close = self
.header("connection")
let is_close = self.header("connection")
.map(|c| c.eq_ignore_ascii_case("close"))
.unwrap_or(false);
@@ -420,8 +419,10 @@ impl Response {
}
}
/// parse a line like: HTTP/1.1 200 OK\r\n
fn parse_status_line(line: &str) -> Result<((usize, usize), u16), Error> {
// HTTP/1.1 200 OK\r\n
//
let mut split = line.splitn(3, ' ');
let http_version = split.next().ok_or_else(|| Error::BadStatus)?;
@@ -448,6 +449,20 @@ fn parse_status_line(line: &str) -> Result<((usize, usize), u16), Error> {
impl FromStr for Response {
type Err = Error;
/// Parse a response from a string.
///
/// Example:
/// ```
/// let s = "HTTP/1.1 200 OK\r\n\
/// X-Forwarded-For: 1.2.3.4\r\n\
/// Content-Type: text/plain\r\n\
/// \r\n\
/// Hello World!!!";
/// let resp = s.parse::<ureq::Response>().unwrap();
/// assert!(resp.has("X-Forwarded-For"));
/// let body = resp.into_string().unwrap();
/// assert_eq!(body, "Hello World!!!");
/// ```
fn from_str(s: &str) -> Result<Self, Self::Err> {
let bytes = s.as_bytes().to_owned();
let mut cursor = Cursor::new(bytes);
@@ -468,12 +483,14 @@ impl Into<Response> for Error {
}
}
/// "Give away" Unit and Stream to the response.
///
/// *Internal API*
pub fn set_stream(resp: &mut Response, unit: Option<Unit>, stream: Stream) {
resp.unit = unit;
resp.stream = Some(stream);
}
// application/x-www-form-urlencoded, application/json, and multipart/form-data
fn read_next_line<R: Read>(reader: &mut R) -> IoResult<AsciiString> {
let mut buf = Vec::new();
@@ -500,6 +517,8 @@ fn read_next_line<R: Read>(reader: &mut R) -> IoResult<AsciiString> {
}
/// Read Wrapper around an (unsafe) pointer to a Stream.
///
/// *Internal API*
struct YoloRead {
stream: *mut Stream,
}
@@ -519,6 +538,9 @@ impl Read for YoloRead {
}
}
/// Limits a YoloRead to a content size (as set by a "Content-Length" header).
///
/// *Internal API*
struct LimitedRead {
reader: YoloRead,
limit: usize,
@@ -556,6 +578,11 @@ impl Read for LimitedRead {
}
}
/// Extract the charset from a "Content-Type" header.
///
/// "Content-Type: text/plain; charset=iso8859-1" -> "iso8859-1"
///
/// *Internal API*
pub fn charset_from_content_type(header: Option<&str>) -> &str {
header
.and_then(|header| {
@@ -574,14 +601,20 @@ mod tests {
#[test]
fn content_type_without_charset() {
let s = "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\n\r\nOK";
let s = "HTTP/1.1 200 OK\r\n\
Content-Type: application/json\r\n\
\r\n\
OK";
let resp = s.parse::<Response>().unwrap();
assert_eq!("application/json", resp.content_type());
}
#[test]
fn content_type_with_charset() {
let s = "HTTP/1.1 200 OK\r\nContent-Type: application/json; charset=iso-8859-4\r\n\r\nOK";
let s = "HTTP/1.1 200 OK\r\n\
Content-Type: application/json; charset=iso-8859-4\r\n\
\r\n\
OK";
let resp = s.parse::<Response>().unwrap();
assert_eq!("application/json", resp.content_type());
}
@@ -595,21 +628,35 @@ mod tests {
#[test]
fn charset() {
let s = "HTTP/1.1 200 OK\r\nContent-Type: application/json; charset=iso-8859-4\r\n\r\nOK";
let s = "HTTP/1.1 200 OK\r\n\
Content-Type: application/json; charset=iso-8859-4\r\n\
\r\n\
OK";
let resp = s.parse::<Response>().unwrap();
assert_eq!("iso-8859-4", resp.charset());
}
#[test]
fn charset_default() {
let s = "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\n\r\nOK";
let s = "HTTP/1.1 200 OK\r\n\
Content-Type: application/json\r\n\
\r\n\
OK";
let resp = s.parse::<Response>().unwrap();
assert_eq!("utf-8", resp.charset());
}
#[test]
fn chunked_transfer() {
let s = "HTTP/1.1 200 OK\r\nTransfer-Encoding: Chunked\r\n\r\n3\r\nhel\r\nb\r\nlo world!!!\r\n0\r\n\r\n";
let s = "HTTP/1.1 200 OK\r\n\
Transfer-Encoding: Chunked\r\n\
\r\n\
3\r\n\
hel\r\n\
b\r\n\
lo world!!!\r\n\
0\r\n\
\r\n";
let resp = s.parse::<Response>().unwrap();
assert_eq!("hello world!!!", resp.into_string().unwrap());
}
@@ -617,15 +664,17 @@ mod tests {
#[test]
#[cfg(feature = "json")]
fn parse_simple_json() {
let s = format!("HTTP/1.1 200 OK\r\n\r\n{{\"hello\":\"world\"}}");
let s = format!(
"HTTP/1.1 200 OK\r\n\
\r\n\
{{\"hello\":\"world\"}}"
);
let resp = s.parse::<Response>().unwrap();
let v = resp.into_json().unwrap();
assert_eq!(
v,
"{\"hello\":\"world\"}"
.parse::<serde_json::Value>()
.unwrap()
);
let compare = "{\"hello\":\"world\"}"
.parse::<serde_json::Value>()
.unwrap();
assert_eq!(v, compare);
}
#[test]

View File

@@ -23,7 +23,8 @@ fn agent_pool() {
let agent = agent();
// req 1
let resp = agent.get("https://s3.amazonaws.com/foosrvr/bbb.mp4")
let resp = agent
.get("https://s3.amazonaws.com/foosrvr/bbb.mp4")
.set("Range", "bytes=1000-1999")
.call();
assert_eq!(resp.status(), 206);
@@ -43,7 +44,8 @@ fn agent_pool() {
}
// req 2 should be done with a reused connection
let resp = agent.get("https://s3.amazonaws.com/foosrvr/bbb.mp4")
let resp = agent
.get("https://s3.amazonaws.com/foosrvr/bbb.mp4")
.set("Range", "bytes=5000-6999")
.call();
assert_eq!(resp.status(), 206);

View File

@@ -4,7 +4,11 @@ use stream::{connect_http, connect_https, connect_test, Stream};
use url::Url;
//
use pool::DEFAULT_HOST;
/// It's a "unit of work". Maybe a bad name for it?
///
/// *Internal API*
#[derive(Debug)]
pub struct Unit {
pub agent: Arc<Mutex<Option<AgentState>>>,
@@ -35,7 +39,7 @@ impl Unit {
let is_head = req.method.eq_ignore_ascii_case("head");
let hostname = url.host_str().unwrap_or("localhost").to_string();
let hostname = url.host_str().unwrap_or(DEFAULT_HOST).to_string();
let query_string = combine_query(&url, &req.query);
@@ -98,6 +102,7 @@ impl Unit {
}
}
/// Perform a connection. Used recursively for redirects.
pub fn connect(
mut unit: Unit,
method: &str,
@@ -202,6 +207,7 @@ fn match_cookies<'a>(jar: &'a CookieJar, domain: &str, path: &str, is_secure: bo
.collect()
}
/// Combine the query of the url and the query options set on the request object.
fn combine_query(url: &Url, query: &QString) -> String {
match (url.query(), query.len() > 0) {
(Some(urlq), true) => format!("?{}&{}", urlq, query),
@@ -211,6 +217,7 @@ fn combine_query(url: &Url, query: &QString) -> String {
}
}
/// Connect the socket, either by using the pool or grab a new one.
fn connect_socket(unit: &Unit, use_pooled: bool) -> Result<(Stream, bool), Error> {
if use_pooled {
let state = &mut unit.agent.lock().unwrap();
@@ -229,9 +236,14 @@ fn connect_socket(unit: &Unit, use_pooled: bool) -> Result<(Stream, bool), Error
Ok((stream?, false))
}
/// Send request line + headers (all up until the body).
fn send_prelude(unit: &Unit, method: &str, stream: &mut Stream) -> IoResult<()> {
// send the request start + headers
//
// build into a buffer and send in one go.
let mut prelude: Vec<u8> = vec![];
// request line
write!(
prelude,
"{} {}{} HTTP/1.1\r\n",
@@ -239,19 +251,27 @@ fn send_prelude(unit: &Unit, method: &str, stream: &mut Stream) -> IoResult<()>
unit.url.path(),
&unit.query_string
)?;
// host header if not set by user.
if !has_header(&unit.headers, "host") {
write!(prelude, "Host: {}\r\n", unit.url.host().unwrap())?;
}
// other headers
for header in &unit.headers {
write!(prelude, "{}: {}\r\n", header.name(), header.value())?;
}
// finish
write!(prelude, "\r\n")?;
// write all to the wire
stream.write_all(&mut prelude[..])?;
Ok(())
}
/// Investigate a response for "Set-Cookie" headers.
fn save_cookies(unit: &Unit, resp: &Response) {
//
@@ -260,6 +280,7 @@ fn save_cookies(unit: &Unit, resp: &Response) {
return;
}
// only lock if we know there is something to process
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() {