From a315c7ac7f68c969760abefc84a6858d69e89eae Mon Sep 17 00:00:00 2001 From: cn-kali-team Date: Wed, 29 Oct 2025 21:57:46 +0800 Subject: [PATCH] =?UTF-8?q?=E6=B7=BB=E5=8A=A0http2?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- Cargo.toml | 7 +- examples/async.rs | 11 ++- examples/cookies.rs | 2 +- examples/http2.rs | 29 ++++++++ src/client.rs | 117 ++++++++++++++++++++--------- src/connector.rs | 177 ++++++++++++++++++++++++++------------------ src/cookies.rs | 2 +- src/h2_client.rs | 159 +++++++++++++++++++++++++++++++++++++++ src/lib.rs | 6 +- src/proxy.rs | 2 +- src/redirect.rs | 12 ++- src/response.rs | 8 +- src/socket.rs | 29 +++++++- 13 files changed, 437 insertions(+), 124 deletions(-) create mode 100644 examples/http2.rs create mode 100644 src/h2_client.rs diff --git a/Cargo.toml b/Cargo.toml index 59b263d..2afd2ed 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "slinger" #改这个 -version = "0.2.7" +version = "0.2.8" edition = "2021" description = "An HTTP Client for Rust designed for hackers." homepage = "https://github.com/emo-crab/slinger" @@ -47,6 +47,8 @@ tokio-rustls = { version = "0.26.2", optional = true, features = ["ring", "tls12 tokio = { version = "1", default-features = false, features = ["net", "io-util", "time"] } rustls-pemfile = { version = "2", optional = true } rustls-pki-types = { version = "1.1.0", features = ["alloc"], optional = true } +rustls-native-certs = { version = "0.8", optional = true } +h2 = { version = "0.4", optional = true } schemars = { version = "1.0.3", optional = true, features = ["bytes1"] } [dev-dependencies] tokio = { version = "1.0", default-features = false, features = ["macros", "rt-multi-thread"] } @@ -54,7 +56,8 @@ futures = { version = "0.3.31" } serde_json = "1" [features] default = [] -tls = ["tokio-rustls", "rustls-pemfile", "rustls-pki-types"] +http2 = ["dep:h2"] +tls = ["tokio-rustls", "rustls-pemfile", "rustls-pki-types", "rustls-native-certs"] cookie = ["dep:cookie"] charset = ["dep:encoding_rs", "dep:mime"] serde = ["dep:serde", "http-serde"] diff --git a/examples/async.rs b/examples/async.rs index ef440f5..e50e833 100644 --- a/examples/async.rs +++ b/examples/async.rs @@ -7,12 +7,17 @@ async fn main() -> Result<(), Box> { let client = ClientBuilder::new().build().unwrap(); let mut worker = FuturesUnordered::new(); let start = std::time::Instant::now(); - let t = vec![1, 2, 3, 4, 5, 6, 7, 8, 9]; - for n in t { - worker.push(client.get(format!("http://httpbin.org/delay/{}", n)).send()); + let mut t = vec![1, 2, 3, 4, 5, 6, 7, 8, 9].into_iter(); + for _ in 0..5 { + if let Some(n) = t.next() { + worker.push(client.get(format!("http://httpbin.org/delay/{}", n)).send()); + } } while let Some(resp) = worker.next().await { println!("{}", resp?.text()?); + if let Some(n) = t.next() { + worker.push(client.get(format!("http://httpbin.org/delay/{}", n)).send()); + } } println!("{:?}", std::time::Instant::now() - start); Ok(()) diff --git a/examples/cookies.rs b/examples/cookies.rs index 2f5ddeb..d203aff 100644 --- a/examples/cookies.rs +++ b/examples/cookies.rs @@ -5,7 +5,7 @@ async fn main() -> Result<(), Box> { use slinger::ClientBuilder; let client = ClientBuilder::new().cookie_store(true).build().unwrap(); let resp = client - .get("https://httpbin.org/cookies/set/key/value") + .get("http://httpbin.org/cookies/set/key/value") .send() .await?; println!("{:?}", resp.text()); diff --git a/examples/http2.rs b/examples/http2.rs new file mode 100644 index 0000000..d94ba9c --- /dev/null +++ b/examples/http2.rs @@ -0,0 +1,29 @@ +//! HTTP/2 example +//! +//! This example demonstrates HTTP/2 support with various stream types. +//! HTTP/2 can work with: +//! - Plain TCP connections (h2c - HTTP/2 over cleartext) +//! - TLS connections (with ALPN negotiation) +//! +//! Note: Most servers require TLS with ALPN to negotiate HTTP/2. +//! For testing h2c (HTTP/2 over cleartext), you need a server that supports it. + +#[tokio::main] +async fn main() -> Result<(), Box> { + #[cfg(all(feature = "http2", feature = "tls"))] + { + use slinger::ClientBuilder; + // Example 2: HTTP/2 with standard TLS (rustls) + println!("=== HTTP/2 with Rustls TLS ==="); + let client = ClientBuilder::new().enable_http2(true).build()?; + // Note: Rustls doesn't support HTTP/2 ALPN by default in this configuration + // This will likely fall back to HTTP/1.1 + let resp = client.get("https://httpbin.org/").send().await?; + println!("Status: {}", resp.status_code()); + println!("Version: {:?}", resp.version()); + let body_len = resp.body().as_ref().map(|b| b.len()).unwrap_or(0); + println!("Body length: {} bytes", body_len); + println!(); + } + Ok(()) +} diff --git a/src/client.rs b/src/client.rs index 3e57419..f3b6354 100644 --- a/src/client.rs +++ b/src/client.rs @@ -1,3 +1,4 @@ +use std::collections::btree_map::Entry; use std::collections::BTreeMap; use std::fmt::{Debug, Formatter}; use std::str::FromStr; @@ -228,30 +229,52 @@ impl Client { /// /// This method fails if there was an error while sending request, /// or redirect limit was exhausted. - pub async fn execute_request(&self, socket: &mut Socket, request: &Request) -> Result { - let raw: Bytes = request.to_raw(); - #[cfg(feature = "tls")] - let mut peer_certificate: Option = None; - #[cfg(feature = "tls")] - { - if let Some(pc) = socket.peer_certificate() { - peer_certificate = Some(pc); + pub async fn execute_request( + &self, + socket_opt: Option, + request: &Request, + ) -> Result<(Response, Option)> { + if let Some(mut socket) = socket_opt { + #[cfg(feature = "http2")] + { + // Check if HTTP/2 was negotiated via ALPN + if socket.is_http2_negotiated() { + let response = + crate::h2_client::send_h2_request(socket, request, self.inner.timeout).await?; + return Ok((response, None)); + } } - } - socket.write_all(&raw).await?; - socket.flush().await?; - let reader = tokio::io::BufReader::new(socket); - let mut irp = ResponseBuilder::new(reader, ResponseConfig::new(request, self.inner.timeout)) - .build() - .await?; - *irp.url_mut() = request.uri().clone(); - #[cfg(feature = "tls")] - { - if let Some(cert) = peer_certificate { - irp.extensions_mut().insert(cert); + // HTTP/1.1 path (existing code) + let raw: Bytes = request.to_raw(); + #[cfg(feature = "tls")] + let mut peer_certificate: Option = None; + #[cfg(feature = "tls")] + { + if let Some(pc) = socket.peer_certificate() { + peer_certificate = Some(pc); + } + } + socket.write_all(&raw).await?; + socket.flush().await?; + let reader = tokio::io::BufReader::new(socket); + let (mut irp, socket) = + ResponseBuilder::new(reader, ResponseConfig::new(request, self.inner.timeout)) + .build() + .await?; + *irp.url_mut() = request.uri().clone(); + #[cfg(feature = "tls")] + { + if let Some(cert) = peer_certificate { + irp.extensions_mut().insert(cert); + } } + Ok((irp, Some(socket))) + } else { + Err(new_io_error( + std::io::ErrorKind::NotConnected, + "socket is None", + )) } - Ok(irp) } /// Executes a `Request`. /// @@ -270,7 +293,7 @@ impl Client { let mut request = request.into(); let mut cur_uri = request.uri().clone(); let mut uris = vec![]; - let mut conn: BTreeMap = BTreeMap::new(); + let mut conn: BTreeMap> = BTreeMap::new(); // 连接一次，同一个主机地址下复用socket连接 let uniq_key = |u: &http::Uri| -> String { let scheme = u.scheme_str().unwrap_or_default(); @@ -309,21 +332,31 @@ impl Client { .insert(http::header::CONNECTION, HeaderValue::from_static("close")); } record.record_request(&request); - let socket = conn - .entry(uniq_key(&cur_uri)) - .or_insert(self.inner.connector.connect_with_uri(&cur_uri).await?); - let mut response = self.execute_request(socket, &request).await?; + let key = uniq_key(&cur_uri); + let socket = match conn.entry(key) { + Entry::Occupied(mut entry) => { + entry.get_mut().take() // take out the Option + } + Entry::Vacant(entry) => { + let new_socket = self.inner.connector.connect_with_uri(&cur_uri).await?; + entry.insert(None); // placeholder + Some(new_socket) + } + }; + let (mut response, socket) = self.execute_request(socket, &request).await?; response.extensions_mut().insert(request.clone()); - if let (Ok(remote_addr), Ok(local_addr)) = (socket.peer_addr(), socket.local_addr()) { - response.extensions_mut().insert(LocalPeerRecord { - remote_addr: remote_addr.into(), - local_addr: local_addr.into(), - }); + if let Some(socket) = socket { + if let (Ok(remote_addr), Ok(local_addr)) = (socket.peer_addr(), socket.local_addr()) { + response.extensions_mut().insert(LocalPeerRecord { + remote_addr: remote_addr.into(), + local_addr: local_addr.into(), + }); + } }; if let Some(cv) = response.headers().get(http::header::CONNECTION) { match cv.to_str().unwrap_or_default() { "keep-alive" => { - if let Some(s) = conn.get_mut(&uniq_key(&cur_uri)) { + if let Some(Some(s)) = conn.get_mut(&uniq_key(&cur_uri)) { if s.peer_addr().is_err() { *s = self.inner.connector.connect_with_uri(&cur_uri).await?; } @@ -420,8 +453,10 @@ impl Client { records.push(record); break; } - for (_key, mut socket) in conn { - socket.shutdown().await?; + for (_key, socket) in conn { + if let Some(mut s) = socket { + s.shutdown().await?; + } } let mut last_response = records .last() @@ -536,6 +571,10 @@ impl ClientBuilder { connector_builder = connector_builder.identity(identity); } } + #[cfg(feature = "http2")] + { + connector_builder = connector_builder.enable_http2(config.http2); + } let connector = connector_builder.build()?; Ok(Client { inner: Arc::new(ClientRef { @@ -831,9 +870,17 @@ impl ClientBuilder { self.config.max_tls_version = version; self } + #[cfg(feature = "http2")] + /// Enable or disable HTTP/2 support. + pub fn enable_http2(mut self, http2: bool) -> Self { + self.config.http2 = http2; + self + } } #[derive(Clone)] struct Config { + #[cfg(feature = "http2")] + http2: bool, timeout: Option, connect_timeout: Option, read_timeout: Option, @@ -878,6 +925,8 @@ impl Debug for Config { impl Default for Config { fn default() -> Self { Self { + #[cfg(feature = "http2")] + http2: false, timeout: Some(Duration::from_secs(30)), connect_timeout: Some(Duration::from_secs(10)), read_timeout: Some(Duration::from_secs(30)), diff --git a/src/connector.rs b/src/connector.rs index ddd46e6..389e255 100644 --- a/src/connector.rs +++ b/src/connector.rs @@ -14,12 +14,20 @@ use tokio::net::TcpSocket; use tokio_rustls::rustls; #[cfg(feature = "tls")] use tokio_rustls::rustls::pki_types::ServerName; + +/// TLS connector type enum to support both rustls and native-tls in the future #[cfg(feature = "tls")] -use tokio_rustls::TlsConnector; +#[derive(Clone)] +pub enum TlsConnectorType { + /// Rustls-based TLS connector + Rustls(tokio_rustls::TlsConnector), +} /// ConnectorBuilder #[derive(Clone)] pub struct ConnectorBuilder { + #[cfg(feature = "http2")] + http2: bool, hostname_verification: bool, certs_verification: bool, read_timeout: Option, @@ -43,6 +51,8 @@ pub struct ConnectorBuilder { impl Default for ConnectorBuilder { fn default() -> Self { Self { + #[cfg(feature = "http2")] + http2: false, hostname_verification: true, certs_verification: true, read_timeout: Some(Duration::from_secs(30)), @@ -66,6 +76,12 @@ impl Default for ConnectorBuilder { } impl ConnectorBuilder { + #[cfg(feature = "http2")] + /// Enable HTTP/2 support. + pub fn enable_http2(mut self, http2: bool) -> Self { + self.http2 = http2; + self + } /// Controls the use of hostname verification. /// /// Defaults to `false`. @@ -200,68 +216,84 @@ impl ConnectorBuilder { } impl ConnectorBuilder { - /// Combine the configuration of this builder with a connector to create a `Connector`. - pub fn build(&self) -> Result { - #[cfg(feature = "tls")] - let tls = { - let mut root_cert_store = rustls::RootCertStore::empty(); - for cert in self.certificate.clone() { - cert.add_to_tls(&mut root_cert_store)?; - } - let mut versions = rustls::ALL_VERSIONS.to_vec(); + #[cfg(feature = "tls")] + /// Use rustls only + fn only_rustls(&self) -> Result { + let mut root_cert_store = rustls::RootCertStore::empty(); + for cert in self.certificate.clone() { + cert.add_to_tls(&mut root_cert_store)?; + } + let certs = rustls_native_certs::load_native_certs().certs; + for cert in certs { + root_cert_store.add(cert)?; + } + let mut versions = rustls::ALL_VERSIONS.to_vec(); - if let Some(min_tls_version) = self.min_tls_version { - versions.retain(|&supported_version| { - match tls::Version::from_tls(supported_version.version) { - Some(version) => version >= min_tls_version, - // Assume it's so new we don't know about it, allow it - // (as of writing this is unreachable) - None => true, - } - }); - } + if let Some(min_tls_version) = self.min_tls_version { + versions.retain(|&supported_version| { + match tls::Version::from_tls(supported_version.version) { + Some(version) => version >= min_tls_version, + // Assume it's so new we don't know about it, allow it + // (as of writing this is unreachable) + None => true, + } + }); + } - if let Some(max_tls_version) = self.max_tls_version { - versions.retain(|&supported_version| { - match tls::Version::from_tls(supported_version.version) { - Some(version) => version <= max_tls_version, - None => false, - } - }); - } + if let Some(max_tls_version) = self.max_tls_version { + versions.retain(|&supported_version| { + match tls::Version::from_tls(supported_version.version) { + Some(version) => version <= max_tls_version, + None => false, + } + }); + } - if versions.is_empty() { - return Err(crate::errors::builder("empty supported tls versions")); + if versions.is_empty() { + return Err(crate::errors::builder("empty supported tls versions")); + } + let provider = rustls::crypto::CryptoProvider::get_default() + .cloned() + .unwrap_or_else(|| std::sync::Arc::new(rustls::crypto::ring::default_provider())); + let signature_algorithms = provider.signature_verification_algorithms; + let config_builder = rustls::ClientConfig::builder_with_provider(provider.clone()) + .with_protocol_versions(&versions) + .map_err(|_| crate::errors::builder("invalid TLS versions"))?; + let config_builder = if !self.certs_verification { + config_builder + .dangerous() + .with_custom_certificate_verifier(std::sync::Arc::new(NoVerifier)) + } else if !self.hostname_verification { + config_builder + .dangerous() + .with_custom_certificate_verifier(std::sync::Arc::new(IgnoreHostname::new( + root_cert_store, + signature_algorithms, + ))) + } else { + config_builder.with_root_certificates(root_cert_store) + }; + let rustls_config = if let Some(id) = self.identity.clone() { + id.add_to_tls(config_builder)? + } else { + config_builder.with_no_client_auth() + }; + #[cfg(feature = "http2")] + let rustls_config = { + let mut config = rustls_config; + if self.http2 { + config.alpn_protocols = vec![b"http/1.1".to_vec(), b"h2".to_vec()]; } - let provider = rustls::crypto::CryptoProvider::get_default() - .cloned() - .unwrap_or_else(|| std::sync::Arc::new(rustls::crypto::ring::default_provider())); - let signature_algorithms = provider.signature_verification_algorithms; - let config_builder = rustls::ClientConfig::builder_with_provider(provider.clone()) - .with_protocol_versions(&versions) - .map_err(|_| crate::errors::builder("invalid TLS versions"))?; - let config_builder = if !self.certs_verification { - config_builder - .dangerous() - .with_custom_certificate_verifier(std::sync::Arc::new(NoVerifier)) - } else if !self.hostname_verification { - config_builder - .dangerous() - .with_custom_certificate_verifier(std::sync::Arc::new(IgnoreHostname::new( - root_cert_store, - signature_algorithms, - ))) - } else { - config_builder.with_root_certificates(root_cert_store) - }; - let tls = if let Some(id) = self.identity.clone() { - id.add_to_tls(config_builder)? - } else { - config_builder.with_no_client_auth() - }; - - TlsConnector::from(std::sync::Arc::new(tls)) + config }; + Ok(TlsConnectorType::Rustls(tokio_rustls::TlsConnector::from( + std::sync::Arc::new(rustls_config), + ))) + } + /// Combine the configuration of this builder with a connector to create a `Connector`. + pub fn build(&self) -> Result { + #[cfg(feature = "tls")] + let tls = { self.only_rustls()? }; let conn = Connector { connect_timeout: self.connect_timeout, nodelay: self.nodelay, @@ -286,7 +318,7 @@ pub struct Connector { write_timeout: Option, proxy: Option, #[cfg(feature = "tls")] - tls: TlsConnector, + tls: TlsConnectorType, } impl PartialEq for Connector { @@ -332,18 +364,21 @@ impl Connector { #[cfg(feature = "tls")] /// A `Connector` will use transport layer security (TLS) by default to connect to destinations. pub async fn upgrade_to_tls(&self, stream: Socket, domain: &str) -> Result { - // 上面是原始socket - let domain = ServerName::try_from(domain.to_owned()) - .map_err(|e| crate::errors::Error::Other(e.to_string()))?; - let this = self.tls.clone(); - let connect_timeout = self.connect_timeout.unwrap_or(Duration::from_secs(30)); - let tls = - stream - .tls(move |t| async move { - tokio::time::timeout(connect_timeout, this.connect(domain, t)).await? - }) - .await?; - Ok(tls) + match &self.tls { + TlsConnectorType::Rustls(connector) => { + // Rustls implementation + let domain = ServerName::try_from(domain.to_owned()) + .map_err(|e| crate::errors::Error::Other(e.to_string()))?; + let this = connector.clone(); + let connect_timeout = self.connect_timeout.unwrap_or(Duration::from_secs(30)); + let tls = stream + .tls(move |t| async move { + tokio::time::timeout(connect_timeout, this.connect(domain, t)).await? + }) + .await?; + Ok(tls) + } + } } } diff --git a/src/cookies.rs b/src/cookies.rs index 9f4f9ba..9eab951 100644 --- a/src/cookies.rs +++ b/src/cookies.rs @@ -201,6 +201,6 @@ pub(crate) fn extract_response_cookie_headers( pub(crate) fn extract_response_cookies( headers: &http::HeaderMap, -) -> impl Iterator> { +) -> impl Iterator, cookie::ParseError>> { headers.get_all(SET_COOKIE).iter().map(Cookie::parse) } diff --git a/src/h2_client.rs b/src/h2_client.rs new file mode 100644 index 0000000..b8c82b7 --- /dev/null +++ b/src/h2_client.rs @@ -0,0 +1,159 @@ +//! HTTP/2 client implementation +//! +//! This module provides HTTP/2 support using the h2 library + +use crate::socket::{MaybeTlsStream, Socket}; +use crate::{Request, Response}; +use bytes::Bytes; +use h2::client; +use http::Version; +use std::pin::Pin; +use std::task::{Context, Poll}; +use std::time::Duration; +use tokio::io::{AsyncRead, AsyncWrite}; + +/// Wrapper for MaybeTlsStream that allows h2 to work with any stream type +enum StreamWrapper { + Tcp(tokio::net::TcpStream), + #[cfg(feature = "tls")] + Tls(Box>), +} + +impl AsyncRead for StreamWrapper { + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut tokio::io::ReadBuf<'_>, + ) -> Poll> { + match self.get_mut() { + StreamWrapper::Tcp(stream) => Pin::new(stream).poll_read(cx, buf), + #[cfg(feature = "tls")] + StreamWrapper::Tls(stream) => Pin::new(stream).poll_read(cx, buf), + } + } +} + +impl AsyncWrite for StreamWrapper { + fn poll_write( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + match self.get_mut() { + StreamWrapper::Tcp(stream) => Pin::new(stream).poll_write(cx, buf), + #[cfg(feature = "tls")] + StreamWrapper::Tls(stream) => Pin::new(stream).poll_write(cx, buf), + } + } + + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + match self.get_mut() { + StreamWrapper::Tcp(stream) => Pin::new(stream).poll_flush(cx), + #[cfg(feature = "tls")] + StreamWrapper::Tls(stream) => Pin::new(stream).poll_flush(cx), + } + } + + fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + match self.get_mut() { + StreamWrapper::Tcp(stream) => Pin::new(stream).poll_shutdown(cx), + #[cfg(feature = "tls")] + StreamWrapper::Tls(stream) => Pin::new(stream).poll_shutdown(cx), + } + } +} + +pub(crate) async fn send_h2_request( + socket: Socket, + request: &Request, + _timeout: Option, +) -> crate::Result { + // Extract the stream from the socket + let _read_timeout = socket.read_timeout; + let _write_timeout = socket.write_timeout; + + let io = match socket.inner { + MaybeTlsStream::Tcp(tcp_stream) => StreamWrapper::Tcp(tcp_stream), + #[cfg(feature = "tls")] + MaybeTlsStream::Tls(tls_stream) => StreamWrapper::Tls(tls_stream), + }; + + // Create an h2 client handshake + let (client, h2_conn) = client::handshake(io) + .await + .map_err(|e| crate::errors::Error::Other(format!("HTTP/2 handshake failed: {}", e)))?; + tokio::spawn(async { + h2_conn.await.unwrap_or_default(); + }); + // Wait for the client to be ready and get it back + let mut client = client + .ready() + .await + .map_err(|e| crate::errors::Error::Other(format!("HTTP/2 client not ready: {}", e)))?; + // Build the HTTP/2 request + let mut h2_request = http::Request::builder() + .method(request.method()) + .uri(request.uri()) + .version(Version::HTTP_2); + // Add headers (skip certain HTTP/1.1 specific headers) + for (key, value) in request.headers().iter() { + // Skip HTTP/1.1 specific headers that shouldn't be in HTTP/2 + if key == http::header::CONNECTION + || key == http::header::TRANSFER_ENCODING + || key == http::header::UPGRADE + || key.as_str().to_lowercase() == "keep-alive" + || key.as_str().to_lowercase() == "proxy-connection" + { + continue; + } + h2_request = h2_request.header(key, value); + } + // Create the request body + let body_bytes = request.body().map(|b| b.as_ref()).unwrap_or(&[]); + let h2_request = h2_request + .body(()) + .map_err(|e| crate::errors::Error::Other(format!("Failed to build HTTP/2 request: {}", e)))?; + // Send the request + let (response_future, mut send_stream) = + client + .send_request(h2_request, body_bytes.is_empty()) + .map_err(|e| crate::errors::Error::Other(format!("Failed to send HTTP/2 request: {}", e)))?; + // Send the body if present + if !body_bytes.is_empty() { + send_stream + .send_data(Bytes::copy_from_slice(body_bytes), true) + .map_err(|e| crate::errors::Error::Other(format!("Failed to send HTTP/2 body: {}", e)))?; + } + // Wait for the response + let h2_response = response_future.await.map_err(|e| { + crate::errors::Error::Other(format!("Failed to receive HTTP/2 response: {}", e)) + })?; + // Extract status and headers + let status = h2_response.status(); + let headers = h2_response.headers().clone(); + let mut body_stream = h2_response.into_body(); + // Read the response body + let mut body = Vec::new(); + while let Some(chunk_result) = body_stream.data().await { + let chunk = chunk_result.map_err(|e| { + crate::errors::Error::Other(format!("Failed to read HTTP/2 body chunk: {}", e)) + })?; + body.extend_from_slice(&chunk); + body_stream + .flow_control() + .release_capacity(chunk.len()) + .map_err(|e| crate::errors::Error::Other(format!("Failed to release flow control: {}", e)))?; + } + // Build a slinger Response + let mut response_builder = http::Response::builder() + .status(status) + .version(Version::HTTP_2); + for (key, value) in headers.iter() { + response_builder = response_builder.header(key, value); + } + let http_response = response_builder + .body(Bytes::from(body)) + .map_err(|e| crate::errors::Error::Other(format!("Failed to build response: {}", e)))?; + let response: Response = http_response.into(); + Ok(response) +} diff --git a/src/lib.rs b/src/lib.rs index b5117cd..0f7428d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -106,7 +106,7 @@ //! //! - **charset**: Improved support for decoding text. //! - **cookie**: Provides cookie session support. -//! - **tls**: Provides https support. +//! - **tls**: Provides https support via rustls. //! - **serde**: Provides serialization and deserialization support. //! - **gzip**: Provides response body gzip decompression. //! @@ -116,6 +116,8 @@ mod connector; #[cfg(feature = "cookie")] mod cookies; mod errors; +#[cfg(feature = "http2")] +mod h2_client; mod proxy; /// record info pub mod record; @@ -147,8 +149,6 @@ pub use proxy::Proxy; pub use request::{RawRequest, Request, RequestBuilder}; pub use response::{Response, ResponseBuilder, ResponseConfig}; pub use socket::Socket; -#[cfg(feature = "tls")] -pub use tokio_rustls; /// Shortcut method to quickly make a `GET` request. /// diff --git a/src/proxy.rs b/src/proxy.rs index 3d77c0c..54d2084 100644 --- a/src/proxy.rs +++ b/src/proxy.rs @@ -271,7 +271,7 @@ impl HttpProxy { let mut buffer = [0; 128]; proxy_socket.read(&mut buffer[..]).await?; let reader = tokio::io::BufReader::new(buffer.as_slice()); - let proxy_response = ResponseBuilder::new(reader, Default::default()) + let (proxy_response, _) = ResponseBuilder::new(reader, Default::default()) .build() .await?; if proxy_response.status_code() != http::StatusCode::OK { diff --git a/src/redirect.rs b/src/redirect.rs index 66c6d90..e1ba1fe 100644 --- a/src/redirect.rs +++ b/src/redirect.rs @@ -12,7 +12,7 @@ use std::str::FromStr; /// the allowed maximum redirect hops in a chain. /// - `none` can be used to disable all redirect behavior. /// - `custom` can be used to create a customized policy. -#[derive(Clone, Debug, PartialEq)] +#[derive(Clone, Debug)] pub enum Policy { /// - `custom` can be used to create a customized policy. see [only_same_host]. Custom(fn(Attempt) -> Action), @@ -22,6 +22,16 @@ pub enum Policy { None, } +impl PartialEq for Policy { + fn eq(&self, other: &Self) -> bool { + match (self, other) { + (Policy::Limit(a), Policy::Limit(b)) => *a == *b, + (Policy::None, Policy::None) => true, + (Policy::Custom(_), Policy::Custom(_)) => false, + _ => false, + } + } +} /// A type that holds information on the current response and previous requests /// in redirect chain. #[derive(Clone, Debug, PartialEq)] diff --git a/src/response.rs b/src/response.rs index 2bcbc70..b15dd9a 100644 --- a/src/response.rs +++ b/src/response.rs @@ -149,7 +149,7 @@ impl Response { /// This requires the optional `cookie` feature to be enabled. #[cfg(feature = "cookie")] #[cfg_attr(docsrs, doc(cfg(feature = "cookie")))] - pub fn cookies(&self) -> impl Iterator { + pub fn cookies(&self) -> impl Iterator> { cookies::extract_response_cookies(&self.headers).filter_map(|x| x.ok()) } @@ -667,7 +667,7 @@ impl ResponseBuilder { /// Build a `Response`, which can be inspected, modified and executed with /// `Client::execute()`. - pub async fn build(mut self) -> Result { + pub async fn build(mut self) -> Result<(Response, T)> { let (v, c) = self.parser_version().await?; self.builder = self.builder.version(v).status(c); let header = self.read_headers().await?; @@ -677,7 +677,9 @@ impl ResponseBuilder { *h = header; } let resp = self.builder.body(body)?; - Ok(resp.into()) + let response = resp.into(); + let socket = self.reader.into_inner(); + Ok((response, socket)) } } diff --git a/src/socket.rs b/src/socket.rs index 77a7876..91ed17d 100644 --- a/src/socket.rs +++ b/src/socket.rs @@ -13,10 +13,11 @@ use tokio_rustls::client::TlsStream; /// Socket #[derive(Debug)] pub struct Socket { - inner: MaybeTlsStream, - read_timeout: Option, - write_timeout: Option, + pub(crate) inner: MaybeTlsStream, + pub(crate) read_timeout: Option, + pub(crate) write_timeout: Option, } + impl Socket { pub(crate) fn new( maybe_tls_stream: MaybeTlsStream, @@ -29,6 +30,25 @@ impl Socket { write_timeout, } } + #[cfg(feature = "http2")] + pub(crate) fn is_http2_negotiated(&self) -> bool { + self + .http2_negotiated() + .map(|alpn| alpn == b"h2") + .unwrap_or(false) + } + #[cfg(feature = "http2")] + pub(crate) fn http2_negotiated(&self) -> Option> { + match &self.inner { + MaybeTlsStream::Tcp(_) => None, + #[cfg(feature = "tls")] + MaybeTlsStream::Tls(tls) => tls + .get_ref() + .1 + .alpn_protocol() + .map(|protocol| protocol.to_vec()), + } + } #[cfg(feature = "tls")] pub(crate) async fn tls(self, func: F) -> Result where @@ -54,7 +74,7 @@ pub enum MaybeTlsStream { /// TCP Tcp(TcpStream), #[cfg(feature = "tls")] - /// TLS + /// TLS with rustls Tls(Box>), } impl MaybeTlsStream { @@ -63,6 +83,7 @@ impl MaybeTlsStream { pub fn peer_certificate(&self) -> Option { match &self { MaybeTlsStream::Tcp(_) => None, + #[cfg(feature = "tls")] MaybeTlsStream::Tls(stream) => stream .get_ref() .1