From 7956bd52c965cf318a9cb997a0607f08fd23a32b Mon Sep 17 00:00:00 2001 From: cn-kali-team Date: Mon, 20 May 2024 23:13:27 +0800 Subject: [PATCH] update --- src/client.rs | 64 ++++++++++++++++++++++++++++----------- src/connector.rs | 79 ++++++++++++++++++++++-------------------------- src/lib.rs | 2 +- src/proxy.rs | 11 ++----- tests/tests.rs | 6 ++-- 5 files changed, 89 insertions(+), 73 deletions(-) diff --git a/src/client.rs b/src/client.rs index baae753..f868f4b 100644 --- a/src/client.rs +++ b/src/client.rs @@ -3,10 +3,55 @@ use crate::record::{HTTPRecord, InternalRequest, InternalResponse}; // use crate::redirect::{Action, Attempt, Policy}; use crate::{ConnectorBuilder, redirect, Result}; use bytes::Bytes; -use http::HeaderMap; +use http::{HeaderMap, HeaderValue}; use std::io::{BufReader, Write}; use std::sync::Arc; use std::time::Duration; +use http::header::{ACCEPT, USER_AGENT}; + +#[derive(Clone)] +pub struct Client { + inner: Arc, +} + +#[must_use] +pub struct ClientBuilder { + config: Config, +} +impl Default for ClientBuilder { + fn default() -> Self { + Self::new() + } +} +impl ClientBuilder { + pub fn new() -> ClientBuilder { + let mut headers: HeaderMap = HeaderMap::with_capacity(2); + headers.insert(ACCEPT, HeaderValue::from_static("*/*")); + ClientBuilder { + config: Config { + connect_timeout: None, + headers, + referer: false, + max_size: None, + proxy: None, + }, + } + } + pub fn build(self) -> crate::Result { + let config = self.config; + let mut connector =ConnectorBuilder::default(); + Ok(Client { + inner: Arc::new(ClientRef { + headers: Default::default(), + redirect_policy: Default::default(), + referer: false, + request_timeout: None, + read_timeout: None, + proxies: Arc::new(vec![]), + }), + }) + } +} #[derive(Clone, Debug, PartialEq)] pub struct Config { @@ -30,6 +75,7 @@ impl Default for Config { } } } + struct ClientRef { // cookie_store: Option>, headers: HeaderMap, @@ -39,10 +85,6 @@ struct ClientRef { read_timeout: Option, proxies: Arc>, } -#[derive(Debug, Default, Clone, PartialEq)] -pub struct Client { - pub config: Config, -} // fn same_origin(attempt: Attempt) -> Action { // if let Some(p) = attempt.previous().last() { @@ -63,18 +105,6 @@ pub struct Client { // } // } -impl Client { - pub fn new(config: Config) -> Result { - Ok(Self { config }) - } -} - -impl Client { - pub fn request(&self, request: InternalRequest) -> Result { - http_request(request) - } -} - pub fn http_request(request: InternalRequest) -> Result { let mut record = HTTPRecord::default(); let mut socket = ConnectorBuilder::default() diff --git a/src/connector.rs b/src/connector.rs index f33ab4c..2b91739 100644 --- a/src/connector.rs +++ b/src/connector.rs @@ -16,10 +16,6 @@ pub struct ConnectorBuilder { write_timeout: Option, connect_timeout: Option, nodelay: bool, - addr: Option, - domain: Option, - socket_type: Option, - protocol: Option, proxy: Option, } @@ -49,54 +45,35 @@ impl ConnectorBuilder { self.connect_timeout = Some(timeout); self } - pub fn addr(mut self, addr: SocketAddr) -> ConnectorBuilder { - self.addr = Some(addr); - self - } pub fn proxy(mut self, addr: Proxy) -> ConnectorBuilder { self.proxy = Some(addr); self } - pub fn domain(mut self, addr: Option) -> ConnectorBuilder { - self.domain = addr; - self - } - pub fn socket_type(mut self, socket_type: Option) -> ConnectorBuilder { - self.socket_type = socket_type; - self - } - pub fn protocol(mut self, protocol: Option) -> ConnectorBuilder { - self.protocol = protocol; - self - } } impl ConnectorBuilder { - pub fn connect_with_target(self, target: &http::Uri) -> Result { - ProxySocket::new(target, &self.proxy).conn_with_connector_builder(self) - } pub fn build(&self) -> Result { - let tcp = match self.addr { - None => Socket::new( - self.domain.unwrap_or(Domain::IPV4), - self.socket_type.unwrap_or(Type::STREAM), - self.protocol, - )?, - Some(addr) => Socket::new(Domain::for_address(addr), Type::STREAM, Some(Protocol::TCP))?, - }; + let (ipv4, ipv6) = ( + Socket::new(Domain::IPV4, Type::STREAM, Some(Protocol::TCP))?, + Socket::new(Domain::IPV6, Type::STREAM, Some(Protocol::TCP))?, + ); let tls = TlsConnector::builder() .danger_accept_invalid_hostnames(!self.hostname_verification) .danger_accept_invalid_certs(!self.certs_verification) .build()?; if self.nodelay { - tcp.set_nodelay(self.nodelay)?; + ipv4.set_nodelay(self.nodelay)?; + ipv6.set_nodelay(self.nodelay)?; } - tcp.set_read_timeout(self.read_timeout)?; - tcp.set_write_timeout(self.write_timeout)?; + ipv4.set_read_timeout(self.read_timeout)?; + ipv4.set_write_timeout(self.write_timeout)?; + ipv6.set_read_timeout(self.read_timeout)?; + ipv6.set_write_timeout(self.write_timeout)?; let conn = Connector { connect_timeout: self.connect_timeout, proxy: self.proxy.clone(), - tcp: InternalSocket::TCP(tcp), + ipv4: InternalSocket::TCP(ipv4), + ipv6: InternalSocket::TCP(ipv6), tls, }; Ok(conn) @@ -107,7 +84,8 @@ impl ConnectorBuilder { pub struct Connector { connect_timeout: Option, proxy: Option, - tcp: InternalSocket, + ipv4: InternalSocket, + ipv6: InternalSocket, tls: TlsConnector, } @@ -119,15 +97,30 @@ impl PartialEq for Connector { impl Connector { pub fn connect_with_addr(&self, addr: SocketAddr) -> Result { - match self.connect_timeout { - None => { - self.tcp.connect(&addr.into())?; + return match addr { + SocketAddr::V4(_) => { + match self.connect_timeout { + None => { + self.ipv4.connect(&addr.into())?; + } + Some(timeout) => { + self.ipv4.connect_timeout(&addr.into(), timeout)?; + } + } + Ok(InternalSocket::TCP(self.ipv4.try_clone().unwrap())) } - Some(timeout) => { - self.tcp.connect_timeout(&addr.into(), timeout)?; + SocketAddr::V6(_) => { + match self.connect_timeout { + None => { + self.ipv6.connect(&addr.into())?; + } + Some(timeout) => { + self.ipv6.connect_timeout(&addr.into(), timeout)?; + } + } + Ok(InternalSocket::TCP(self.ipv6.try_clone().unwrap())) } - } - Ok(InternalSocket::TCP(self.tcp.try_clone().unwrap())) + }; } pub fn connect_with_uri(&self, target: &http::Uri) -> Result { ProxySocket::new(target, &self.proxy).conn_with_connector(self) diff --git a/src/lib.rs b/src/lib.rs index 1ace471..8379c28 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -9,7 +9,7 @@ mod record; mod redirect; mod socket; -pub use client::{Client, Config}; +pub use client::{Client, Config, http_request}; pub use connector::{Connector, ConnectorBuilder}; pub use errors::{EngineError, ReplyError, Result}; pub use http::{HeaderMap, HeaderName, HeaderValue, Method}; diff --git a/src/proxy.rs b/src/proxy.rs index aa9dba2..0315730 100644 --- a/src/proxy.rs +++ b/src/proxy.rs @@ -1,4 +1,4 @@ -use crate::connector::{Connector, ConnectorBuilder}; +use crate::connector::Connector; use crate::{EngineError, InternalRequest, InternalResponse, ReplyError, Result}; use bytes::Bytes; use http::HeaderValue; @@ -247,13 +247,6 @@ impl ProxySocket { } }; } - pub fn conn_with_connector_builder( - self, - connector_builder: ConnectorBuilder, - ) -> Result { - let conn = connector_builder.addr(self.get_conn_addr()?).build()?; - self.conn_with_connector(&conn) - } } fn default_port(uri: &http::Uri) -> Option { @@ -582,7 +575,7 @@ impl TargetAddr { port, )) } - } + }; } pub fn to_be_bytes(&self, cmd: Socks5Command) -> Result<(usize, Vec)> { let mut packet = [0u8; consts::MAX_ADDR_LEN + 3]; diff --git a/tests/tests.rs b/tests/tests.rs index 6982120..28b9c15 100644 --- a/tests/tests.rs +++ b/tests/tests.rs @@ -1,7 +1,7 @@ #[cfg(test)] mod tests { use http::{HeaderName, HeaderValue}; - use rawrequest_rs::{Client, Config, InternalRequest}; + use rawrequest_rs::{Client, Config, http_request, InternalRequest}; use std::collections::HashMap; #[test] fn socket() { @@ -25,8 +25,8 @@ mod tests { }; b = b.uri("https://httpbin.org/ip"); let r = b.body(None).unwrap(); - let client = Client::new(Config::default()).unwrap(); - let recode = client.request(r.into()).expect("err"); + // let client = Client::new(Config::default()).unwrap(); + let recode = http_request(r.into()).expect("err"); println!("{:?}", recode); } }