use std::collections::{HashMap, VecDeque};
use std::io::{self, Read, Write};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use go_lib::context::{with_timeout, Context};
use go_lib::net::TcpStream;
use url::Url;
use crate::cookie::{Cookie, CookieJar};
use crate::error::HttpError;
use crate::header::Header;
use crate::parse::response::{read_response, ParsedResponse};
use crate::parse::transfer::Body;
use crate::request::Request;
use crate::response::Response;
pub trait RoundTripper: Send + Sync {
fn round_trip(&self, req: Request) -> Result<Response, HttpError>;
}
struct IdleConn {
stream: TcpStream,
}
pub type ProxyFn = Arc<dyn Fn(&Request) -> Result<Option<Url>, HttpError> + Send + Sync>;
pub struct Transport {
pub max_idle_conns_per_host: usize,
pub idle_conn_timeout: Option<Duration>,
pub dial_timeout: Option<Duration>,
pub tls_config: Option<Arc<rustls::ClientConfig>>,
pub proxy: Option<ProxyFn>,
pool: Mutex<HashMap<String, VecDeque<IdleConn>>>,
}
impl Transport {
pub fn new() -> Self {
Self {
max_idle_conns_per_host: 10,
idle_conn_timeout: Some(Duration::from_secs(90)),
dial_timeout: Some(Duration::from_secs(30)),
tls_config: None,
proxy: None,
pool: Mutex::new(HashMap::new()),
}
}
fn acquire(&self, host_port: &str) -> io::Result<TcpStream> {
if let Some(conn) = self
.pool
.lock()
.unwrap()
.get_mut(host_port)
.and_then(|q| q.pop_front())
{
return Ok(conn.stream);
}
TcpStream::connect(host_port)
}
fn release(&self, host_port: &str, stream: TcpStream) {
let mut pool = self.pool.lock().unwrap();
let queue = pool.entry(host_port.to_owned()).or_default();
if queue.len() < self.max_idle_conns_per_host {
queue.push_back(IdleConn { stream });
}
}
}
impl Default for Transport {
fn default() -> Self {
Self::new()
}
}
impl RoundTripper for Transport {
fn round_trip(&self, mut req: Request) -> Result<Response, HttpError> {
let is_https = req.url.scheme() == "https";
let host = req.url.host_str().unwrap_or("localhost").to_owned();
let port = req.url.port_or_known_default()
.unwrap_or(if is_https { 443 } else { 80 });
let target_hp = format!("{host}:{port}");
let proxy_url = match &self.proxy {
Some(f) => f(&req)?,
None => None,
};
match proxy_url {
None => {
if is_https {
self.https_round_trip(req, &host, &target_hp, None)
} else {
self.http_round_trip(req, &target_hp, false)
}
}
Some(pu) => {
let proxy_hp = proxy_host_port(&pu)?;
let auth = proxy_auth_header(&pu);
if is_https {
self.https_round_trip(req, &host, &target_hp, Some((proxy_hp, auth)))
} else {
if let Some(a) = auth {
req.header.set("Proxy-Authorization", a);
}
self.http_round_trip(req, &proxy_hp, true)
}
}
}
}
}
impl Transport {
fn http_round_trip(
&self,
mut req: Request,
dial_hp: &str,
absolute: bool,
) -> Result<Response, HttpError> {
let mut stream = self.acquire(dial_hp).map_err(HttpError::Io)?;
send_request(&mut stream, &mut req, absolute)?;
let mut parsed = read_response(
stream.try_clone().map_err(HttpError::Io)?,
Some(req.method.as_str()),
crate::parse::request::DEFAULT_MAX_HEADER_BYTES,
)?;
let keep_alive = is_keep_alive_parsed(&parsed, req.proto_minor);
if keep_alive {
let bytes = parsed.body.read_to_vec().map_err(|_| HttpError::BodyRead)?;
parsed.body = Body::Unbounded(Box::new(io::Cursor::new(bytes)));
self.release(dial_hp, stream);
}
Ok(parsed_response_to_response(parsed))
}
fn https_round_trip(
&self,
mut req: Request,
sni_host: &str,
target_hp: &str,
proxy: Option<(String, Option<String>)>,
) -> Result<Response, HttpError> {
let stream = match proxy {
Some((proxy_hp, auth)) => {
let mut s = TcpStream::connect(proxy_hp.as_str()).map_err(HttpError::Io)?;
connect_tunnel(&mut s, target_hp, auth.as_deref())?;
s
}
None => TcpStream::connect(target_hp).map_err(HttpError::Io)?,
};
let tls_cfg = match &self.tls_config {
Some(c) => Arc::clone(c),
None => crate::tls::default_client_config(),
};
let server_name = rustls::pki_types::ServerName::try_from(sni_host.to_owned())
.map_err(|e| HttpError::Tls(e.to_string()))?;
let client_conn = rustls::ClientConnection::new(tls_cfg, server_name)
.map_err(|e| HttpError::Tls(e.to_string()))?;
let mut tls = rustls::StreamOwned::new(client_conn, stream);
send_request(&mut tls, &mut req, false)?;
let read_ptr: *mut dyn Read = &mut tls as &mut dyn Read as *mut dyn Read;
let parsed = read_response(
RawRead(read_ptr),
Some(req.method.as_str()),
crate::parse::request::DEFAULT_MAX_HEADER_BYTES,
)?;
Ok(parsed_response_to_response(parsed))
}
}
struct RawRead(*mut dyn Read);
unsafe impl Send for RawRead {}
impl Read for RawRead {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
unsafe { (*self.0).read(buf) }
}
}
pub enum RedirectPolicy {
Follow,
UseLastResponse,
}
pub type CheckRedirect =
Arc<dyn Fn(&Request, &[Request]) -> Result<RedirectPolicy, HttpError> + Send + Sync>;
pub struct Client {
pub transport: Arc<dyn RoundTripper>,
pub timeout: Option<Duration>,
pub max_redirects: usize,
pub check_redirect: Option<CheckRedirect>,
pub jar: Option<Arc<dyn CookieJar>>,
}
impl Client {
pub fn new() -> Self {
Self {
transport: Arc::new(Transport::new()),
timeout: None,
max_redirects: 10,
check_redirect: None,
jar: None,
}
}
pub fn get(&self, url: &str) -> Result<Response, HttpError> {
let req = Request::new("GET", url, None)?;
self.do_request(req)
}
pub fn post(
&self,
url: &str,
content_type: &str,
body: Body,
) -> Result<Response, HttpError> {
let mut req = Request::new("POST", url, Some(body))?;
req.header.set("Content-Type", content_type);
self.do_request(req)
}
pub fn post_form(
&self,
url: &str,
values: &[(&str, &str)],
) -> Result<Response, HttpError> {
let encoded = url_encode(values);
let body = Body::Unbounded(Box::new(io::Cursor::new(encoded.into_bytes())));
self.post(url, "application/x-www-form-urlencoded", body)
}
pub fn head(&self, url: &str) -> Result<Response, HttpError> {
let req = Request::new("HEAD", url, None)?;
self.do_request(req)
}
pub fn do_request(&self, req: Request) -> Result<Response, HttpError> {
let (_cancel, ctx) = match self.timeout {
Some(d) => {
let (c, cancel) = with_timeout(req.context(), d);
(Some(cancel), c)
}
None => (None, req.context().clone()),
};
let mut req = req;
if let Some(jar) = &self.jar {
attach_cookies(&mut req, jar.as_ref());
}
let mut via: Vec<Request> = Vec::new();
loop {
let method = req.method.clone();
let url = req.url.clone();
let via_entry = snapshot_request(&req);
let mut resp = self.execute_round_trip(req, &ctx)?;
via.push(via_entry);
if let Some(jar) = &self.jar {
store_cookies(&url, &resp.header, jar.as_ref());
}
let status = resp.status;
if !is_redirect(status) {
return Ok(resp);
}
let location = resp
.header
.get("Location")
.ok_or_else(|| HttpError::InvalidUrl("redirect with no Location".into()))?
.to_owned();
let new_url = resolve_url(&url, &location)?;
let new_method = match status {
301..=303 => {
if method == "POST" { "GET".to_owned() } else { method }
}
_ => method,
};
let mut new_req = Request::new(&new_method, new_url.as_str(), None)?;
forward_headers(&mut new_req.header, &resp.header, same_origin(&url, &new_url));
if let Some(jar) = &self.jar {
attach_cookies(&mut new_req, jar.as_ref());
}
match &self.check_redirect {
Some(policy) => match policy(&new_req, &via)? {
RedirectPolicy::Follow => {}
RedirectPolicy::UseLastResponse => return Ok(resp),
},
None => {
if via.len() > self.max_redirects {
return Err(HttpError::TooManyRedirects);
}
}
}
let _ = resp.body_bytes();
req = new_req;
}
}
fn execute_round_trip(&self, req: Request, ctx: &Context) -> Result<Response, HttpError> {
if ctx.deadline().is_none() {
return self.transport.round_trip(req);
}
if ctx.is_done() {
return Err(HttpError::Timeout);
}
let (tx, rx) = go_lib::chan::chan::<Result<Response, HttpError>>(1);
let transport = Arc::clone(&self.transport);
go_lib::go!(move || {
let result = transport.round_trip(req);
let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| tx.send(result)));
});
go_lib::select! {
recv(ctx.done()) -> _sig => { Err(HttpError::Timeout) }
recv(rx) -> result => { result.unwrap_or_else(|| Err(HttpError::Timeout)) }
}
}
}
impl Default for Client {
fn default() -> Self {
Self::new()
}
}
fn default_client() -> &'static Client {
use std::sync::OnceLock;
static DEFAULT: OnceLock<Client> = OnceLock::new();
DEFAULT.get_or_init(Client::new)
}
pub fn get(url: &str) -> Result<Response, HttpError> {
default_client().get(url)
}
pub fn post(url: &str, content_type: &str, body: Body) -> Result<Response, HttpError> {
default_client().post(url, content_type, body)
}
pub fn post_form(url: &str, values: &[(&str, &str)]) -> Result<Response, HttpError> {
default_client().post_form(url, values)
}
pub fn head(url: &str) -> Result<Response, HttpError> {
default_client().head(url)
}
fn send_request(w: &mut impl Write, req: &mut Request, absolute: bool) -> Result<(), HttpError> {
if absolute {
req.write_header_absolute_to(w)?;
} else {
req.write_header_to(w)?;
}
if let Some(body) = req.body.take() {
use std::io::Read;
let mut body = body;
let mut buf = [0u8; 8192];
loop {
let n = body.read(&mut buf).map_err(|_| HttpError::BodyRead)?;
if n == 0 { break; }
w.write_all(&buf[..n])?;
}
}
Ok(())
}
fn is_keep_alive_parsed(resp: &ParsedResponse, req_minor: u8) -> bool {
let conn = resp.header.get("Connection").unwrap_or("").to_ascii_lowercase();
if conn.contains("close") { return false; }
if req_minor == 0 { conn.contains("keep-alive") } else { true }
}
fn parsed_response_to_response(p: ParsedResponse) -> Response {
Response {
status: p.status,
status_text: p.status_text,
proto: p.proto,
proto_major: p.proto_major,
proto_minor: p.proto_minor,
header: p.header,
body: match p.body {
Body::Empty => None,
other => Some(other),
},
content_length: p.content_length,
transfer_encoding: p.transfer_encoding,
trailer: Header::new(),
}
}
pub fn proxy_from_environment() -> ProxyFn {
Arc::new(|req: &Request| -> Result<Option<Url>, HttpError> {
let is_https = req.url.scheme() == "https";
let host = req.url.host_str().unwrap_or("");
if let Some(no) = env_first(&["NO_PROXY", "no_proxy"])
&& host_matches_no_proxy(host, &no)
{
return Ok(None);
}
let names: &[&str] = if is_https {
&["HTTPS_PROXY", "https_proxy"]
} else {
&["HTTP_PROXY", "http_proxy"]
};
match env_first(names) {
None => Ok(None),
Some(v) if v.trim().is_empty() => Ok(None),
Some(v) => {
let raw = v.trim();
let normalized = if raw.contains("://") {
raw.to_owned()
} else {
format!("http://{raw}")
};
let url = Url::parse(&normalized)
.map_err(|e| HttpError::Proxy(format!("bad proxy URL {raw:?}: {e}")))?;
Ok(Some(url))
}
}
})
}
fn env_first(names: &[&str]) -> Option<String> {
names.iter().find_map(|n| std::env::var(n).ok())
}
fn host_matches_no_proxy(host: &str, no_proxy: &str) -> bool {
let host = host.trim_start_matches('.').to_ascii_lowercase();
for entry in no_proxy.split(',') {
let e = entry.trim().to_ascii_lowercase();
if e.is_empty() {
continue;
}
if e == "*" {
return true;
}
let suffix = e.trim_start_matches('.');
if host == suffix || host.ends_with(&format!(".{suffix}")) {
return true;
}
}
false
}
fn proxy_host_port(pu: &Url) -> Result<String, HttpError> {
let h = pu
.host_str()
.ok_or_else(|| HttpError::Proxy("proxy URL has no host".into()))?;
let p = pu.port_or_known_default().unwrap_or(80);
Ok(format!("{h}:{p}"))
}
fn proxy_auth_header(pu: &Url) -> Option<String> {
let user = pu.username();
if user.is_empty() {
return None;
}
let pass = pu.password().unwrap_or("");
let creds = format!("{user}:{pass}");
let b64 = base64::Engine::encode(&base64::engine::general_purpose::STANDARD, creds);
Some(format!("Basic {b64}"))
}
fn connect_tunnel<S: Read + Write>(
stream: &mut S,
target_hp: &str,
auth: Option<&str>,
) -> Result<(), HttpError> {
let mut req = format!("CONNECT {target_hp} HTTP/1.1\r\nHost: {target_hp}\r\n");
if let Some(a) = auth {
req.push_str(&format!("Proxy-Authorization: {a}\r\n"));
}
req.push_str("\r\n");
stream.write_all(req.as_bytes()).map_err(HttpError::Io)?;
let status = read_connect_response(stream)?;
if !(200..300).contains(&status) {
return Err(HttpError::Proxy(format!(
"CONNECT to {target_hp} failed with status {status}"
)));
}
Ok(())
}
fn read_connect_response<R: Read>(stream: &mut R) -> Result<u16, HttpError> {
let mut buf = Vec::with_capacity(128);
let mut byte = [0u8; 1];
loop {
let n = stream.read(&mut byte).map_err(HttpError::Io)?;
if n == 0 {
return Err(HttpError::Proxy("proxy closed before CONNECT response".into()));
}
buf.push(byte[0]);
if buf.ends_with(b"\r\n\r\n") {
break;
}
if buf.len() > 8192 {
return Err(HttpError::Proxy("CONNECT response headers too large".into()));
}
}
let text = String::from_utf8_lossy(&buf);
let first = text.lines().next().unwrap_or("");
first
.split_whitespace()
.nth(1)
.and_then(|s| s.parse::<u16>().ok())
.ok_or_else(|| HttpError::Proxy(format!("malformed CONNECT status line: {first:?}")))
}
fn snapshot_request(r: &Request) -> Request {
let mut s = Request::new(&r.method, r.url.as_str(), None)
.unwrap_or_else(|_| Request::new("GET", "http://invalid.invalid/", None).unwrap());
s.header = r.header.clone();
s.host = r.host.clone();
s.proto = r.proto.clone();
s
}
fn forward_headers(dst: &mut Header, src: &Header, same_origin: bool) {
for (name, values) in src.iter() {
let lower = name.to_ascii_lowercase();
if matches!(
lower.as_str(),
"connection" | "keep-alive" | "proxy-authenticate"
| "proxy-authorization" | "te" | "trailers"
| "transfer-encoding" | "upgrade"
) {
continue;
}
if !same_origin && lower == "authorization" {
continue;
}
for v in values {
dst.add(name, v.as_str());
}
}
}
fn is_redirect(status: u16) -> bool {
matches!(status, 301 | 302 | 303 | 307 | 308)
}
fn resolve_url(base: &Url, location: &str) -> Result<Url, HttpError> {
if location.starts_with("http://") || location.starts_with("https://") {
Url::parse(location).map_err(|e| HttpError::InvalidUrl(e.to_string()))
} else {
base.join(location).map_err(|e| HttpError::InvalidUrl(e.to_string()))
}
}
fn same_origin(a: &Url, b: &Url) -> bool {
a.scheme() == b.scheme()
&& a.host_str() == b.host_str()
&& a.port() == b.port()
}
fn attach_cookies(req: &mut Request, jar: &dyn CookieJar) {
let cookies = jar.cookies(&req.url);
if !cookies.is_empty() {
let pairs: Vec<String> = cookies
.iter()
.map(|c| format!("{}={}", c.name, c.value))
.collect();
req.header.set("Cookie", pairs.join("; "));
}
}
fn store_cookies(url: &Url, header: &Header, jar: &dyn CookieJar) {
let cookies: Vec<Cookie> = header
.values("Set-Cookie")
.iter()
.filter_map(|v| {
let eq = v.find('=')?;
let name = v[..eq].trim().to_owned();
let rest = &v[eq + 1..];
let value = rest.split(';').next().unwrap_or("").trim().to_owned();
Some(Cookie::new(name, value))
})
.collect();
if !cookies.is_empty() {
jar.set_cookies(url, &cookies);
}
}
fn url_encode(values: &[(&str, &str)]) -> String {
values
.iter()
.map(|(k, v)| format!("{}={}", encode_form(k), encode_form(v)))
.collect::<Vec<_>>()
.join("&")
}
fn encode_form(s: &str) -> String {
let mut out = String::with_capacity(s.len());
for b in s.bytes() {
match b {
b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9'
| b'-' | b'_' | b'.' | b'~' => out.push(b as char),
b' ' => out.push('+'),
_ => {
out.push('%');
out.push_str(&format!("{b:02X}"));
}
}
}
out
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn url_encode_basic() {
let pairs = [("q", "hello world"), ("lang", "rust")];
assert_eq!(url_encode(&pairs), "q=hello+world&lang=rust");
}
#[test]
fn url_encode_special_chars() {
let pairs = [("a", "b&c=d")];
assert_eq!(url_encode(&pairs), "a=b%26c%3Dd");
}
#[test]
fn resolve_url_absolute() {
let base = Url::parse("http://example.com/foo").unwrap();
let resolved = resolve_url(&base, "http://other.com/bar").unwrap();
assert_eq!(resolved.as_str(), "http://other.com/bar");
}
#[test]
fn resolve_url_relative() {
let base = Url::parse("http://example.com/a/b").unwrap();
let resolved = resolve_url(&base, "/c").unwrap();
assert_eq!(resolved.as_str(), "http://example.com/c");
}
#[test]
fn is_redirect_codes() {
for code in [301u16, 302, 303, 307, 308] {
assert!(is_redirect(code), "{code} should be redirect");
}
for code in [200u16, 404, 500] {
assert!(!is_redirect(code), "{code} should not be redirect");
}
}
#[test]
fn same_origin_check() {
let a = Url::parse("http://example.com/foo").unwrap();
let b = Url::parse("http://example.com/bar").unwrap();
let c = Url::parse("https://example.com/foo").unwrap();
let d = Url::parse("http://other.com/foo").unwrap();
assert!(same_origin(&a, &b));
assert!(!same_origin(&a, &c)); assert!(!same_origin(&a, &d)); }
#[test]
fn transport_pool_reuse() {
let t = Transport::new();
assert_eq!(t.max_idle_conns_per_host, 10);
assert!(t.pool.lock().unwrap().is_empty());
}
#[test]
fn proxy_auth_header_from_userinfo() {
use base64::Engine;
let with = Url::parse("http://user:pass@proxy.local:3128").unwrap();
let expected = format!(
"Basic {}",
base64::engine::general_purpose::STANDARD.encode("user:pass")
);
assert_eq!(proxy_auth_header(&with), Some(expected));
let without = Url::parse("http://proxy.local:3128").unwrap();
assert_eq!(proxy_auth_header(&without), None);
}
#[test]
fn proxy_host_port_defaults_port_80() {
let u = Url::parse("http://proxy.local").unwrap();
assert_eq!(proxy_host_port(&u).unwrap(), "proxy.local:80");
let u2 = Url::parse("http://proxy.local:8080").unwrap();
assert_eq!(proxy_host_port(&u2).unwrap(), "proxy.local:8080");
}
#[test]
fn no_proxy_matching() {
assert!(host_matches_no_proxy("example.com", "example.com"));
assert!(host_matches_no_proxy("api.example.com", ".example.com"));
assert!(host_matches_no_proxy("api.example.com", "example.com"));
assert!(host_matches_no_proxy("anything", "*"));
assert!(host_matches_no_proxy("b.internal", "foo.com, .internal"));
assert!(!host_matches_no_proxy("example.org", "example.com"));
assert!(!host_matches_no_proxy("notexample.com", ".example.com"));
}
#[test]
fn connect_response_parsing() {
use std::io::Cursor;
let mut ok = Cursor::new(b"HTTP/1.1 200 Connection established\r\n\r\n".to_vec());
assert_eq!(read_connect_response(&mut ok).unwrap(), 200);
let mut denied = Cursor::new(b"HTTP/1.1 407 Proxy Auth Required\r\n\r\n".to_vec());
assert_eq!(read_connect_response(&mut denied).unwrap(), 407);
}
struct MockConn {
to_read: std::io::Cursor<Vec<u8>>,
written: Vec<u8>,
}
impl Read for MockConn {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
self.to_read.read(buf)
}
}
impl Write for MockConn {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
self.written.extend_from_slice(buf);
Ok(buf.len())
}
fn flush(&mut self) -> io::Result<()> {
Ok(())
}
}
#[test]
fn connect_tunnel_sends_request_and_accepts_200() {
let mut conn = MockConn {
to_read: std::io::Cursor::new(b"HTTP/1.1 200 OK\r\n\r\n".to_vec()),
written: Vec::new(),
};
connect_tunnel(&mut conn, "example.com:443", Some("Basic Zm9v")).unwrap();
let sent = String::from_utf8(conn.written.clone()).unwrap();
assert!(sent.starts_with("CONNECT example.com:443 HTTP/1.1\r\n"), "got: {sent:?}");
assert!(sent.contains("Host: example.com:443\r\n"));
assert!(sent.contains("Proxy-Authorization: Basic Zm9v\r\n"));
assert!(sent.ends_with("\r\n\r\n"));
}
#[test]
fn connect_tunnel_errors_on_non_2xx() {
let mut conn = MockConn {
to_read: std::io::Cursor::new(b"HTTP/1.1 403 Forbidden\r\n\r\n".to_vec()),
written: Vec::new(),
};
let err = connect_tunnel(&mut conn, "example.com:443", None).unwrap_err();
assert!(matches!(err, HttpError::Proxy(_)), "got: {err:?}");
}
}