1use std::collections::HashMap;
2use std::net::ToSocketAddrs;
3use std::sync::{Arc, Mutex};
4use url::Url;
5
6use crate::client::connector::{connect, QuikConnection};
7use crate::client::proxy::Proxy;
8use crate::client::request::inject_chrome_headers;
9use crate::client::response::Response;
10use crate::error::{Error, Result};
11use crate::profile::ChromeProfile;
12
13use bytes::Bytes;
14use cookie_store::CookieStore;
15use std::sync::RwLock;
16
17#[derive(Clone)]
34pub struct Client {
35 pool: Arc<Mutex<HashMap<String, QuikConnection>>>,
37 profile: ChromeProfile,
39 proxy: Option<Proxy>,
41 pub cookie_store: Arc<RwLock<CookieStore>>,
46}
47
48impl Default for Client {
49 fn default() -> Self {
50 Self::new()
51 }
52}
53
54impl Client {
55 pub fn new() -> Self {
61 Self::builder().build().unwrap_or_else(|_| Client {
62 pool: Arc::new(Mutex::new(HashMap::new())),
63 profile: crate::profile::chrome_134::profile_auto(),
64 proxy: None,
65 cookie_store: Arc::new(RwLock::new(CookieStore::default())),
66 })
67 }
68
69 pub fn builder() -> ClientBuilder {
71 ClientBuilder::default()
72 }
73
74 pub async fn get(&self, url: &str) -> Result<Response> {
76 self.execute_with_redirects("GET", url, None).await
77 }
78
79 pub async fn post(&self, url: &str, body: Bytes) -> Result<Response> {
81 self.execute_with_redirects("POST", url, Some(body)).await
82 }
83
84 async fn execute_with_redirects(
97 &self,
98 initial_method: &str,
99 initial_url: &str,
100 initial_body: Option<Bytes>,
101 ) -> Result<Response> {
102 let mut current_url_str = initial_url.to_string();
103 let mut current_method = initial_method.to_string();
104 let mut current_body = initial_body;
105
106 let mut sec_fetch_site = "none".to_string();
107 let mut is_cross_site = false;
108
109 for hop in 0..10 {
110 let parsed_url =
111 Url::parse(¤t_url_str).map_err(|e| Error::InvalidUrl(e.to_string()))?;
112 let authority = parsed_url
113 .host_str()
114 .ok_or_else(|| Error::InvalidUrl("missing host".to_string()))?;
115 let port = parsed_url.port().unwrap_or(443);
116
117 let proxy_prefix = self
119 .proxy
120 .as_ref()
121 .map(|p| match p {
122 Proxy::Http(a) => format!("http://{}@", a),
123 Proxy::Socks5(a) => format!("socks5://{}@", a),
124 })
125 .unwrap_or_default();
126
127 let key = format!("{}{}:{}", proxy_prefix, authority, port);
128
129 let cookie_header = {
131 let store = self
132 .cookie_store
133 .read()
134 .map_err(|_| Error::Connect(std::io::Error::other("cookie store poisoned")))?;
135 let cookies: Vec<_> = store
136 .matches(&parsed_url)
137 .iter()
138 .map(|c| format!("{}={}", c.name(), c.value()))
139 .collect();
140 if cookies.is_empty() {
141 None
142 } else {
143 Some(cookies.join("; "))
144 }
145 };
146
147 let mut request = http::Request::builder()
148 .method(current_method.as_str())
149 .uri(parsed_url.as_str())
150 .body(())
151 .map_err(|e| Error::InvalidUrl(e.to_string()))?;
152
153 if let Some(c) = cookie_header.as_deref() {
154 if let Ok(val) = http::header::HeaderValue::from_str(c) {
155 request.headers_mut().insert("cookie", val);
156 }
157 }
158
159 if current_method == "POST" || current_method == "PUT" || current_method == "PATCH" {
162 if let Ok(val) =
163 http::header::HeaderValue::from_str(&parsed_url.origin().ascii_serialization())
164 {
165 request.headers_mut().insert("origin", val);
166 }
167 }
168
169 let is_initial = hop == 0;
171 inject_chrome_headers(
172 request.headers_mut(),
173 &self.profile,
174 &sec_fetch_site,
175 is_initial,
176 );
177
178 let conn = {
180 let mut pool = self.pool.lock().map_err(|_| {
181 Error::Connect(std::io::Error::other("connection pool poisoned"))
182 })?;
183 pool.remove(&key)
184 };
185
186 let mut h2_client = if let Some(mut c) = conn {
187 match c.h2.ready().await {
189 Ok(h2) => {
190 c.h2 = h2;
191 c
192 }
193 Err(_) => self.dial(authority, port, &self.profile).await?,
194 }
195 } else {
196 self.dial(authority, port, &self.profile).await?
197 };
198
199 let mut response = h2_client.send(request, current_body.clone()).await?;
200
201 if let Ok(mut pool) = self.pool.lock() {
203 pool.insert(key, h2_client);
204 }
205
206 self.store_cookies(&response, &parsed_url);
207
208 let status = response.status();
209 if status.is_redirection() {
210 if let Some(location) = response.headers().get("location") {
211 let loc_str = location.to_str().unwrap_or("");
212 let next_url = parsed_url
213 .join(loc_str)
214 .map_err(|e| Error::InvalidUrl(e.to_string()))?;
215
216 if status == http::StatusCode::MOVED_PERMANENTLY
218 || status == http::StatusCode::FOUND
219 || status == http::StatusCode::SEE_OTHER
220 {
221 current_method = "GET".to_string();
222 current_body = None;
223 }
224
225 if !is_cross_site {
227 if next_url.origin() == parsed_url.origin() {
228 sec_fetch_site = "same-origin".to_string();
229 } else if next_url.domain() == parsed_url.domain() {
230 sec_fetch_site = "same-site".to_string();
231 } else {
232 sec_fetch_site = "cross-site".to_string();
233 is_cross_site = true;
234 }
235 }
236
237 current_url_str = next_url.to_string();
238 continue;
239 }
240 }
241
242 response.set_url(current_url_str);
243 return Ok(response);
244 }
245
246 Err(Error::Connect(std::io::Error::other(
247 "Redirect limit exceeded (max 10)",
248 )))
249 }
250
251 async fn dial(
253 &self,
254 authority: &str,
255 port: u16,
256 profile: &ChromeProfile,
257 ) -> Result<QuikConnection> {
258 let addr_str = format!("{}:{}", authority, port);
259 let addr = addr_str.to_socket_addrs()?.next().ok_or_else(|| {
260 std::io::Error::new(std::io::ErrorKind::NotFound, "could not resolve host")
261 })?;
262
263 connect(authority, port, addr, profile, self.proxy.as_ref()).await
264 }
265
266 fn store_cookies(&self, resp: &Response, url: &Url) {
268 if let Ok(mut store) = self.cookie_store.write() {
269 for v in resp.headers().get_all("set-cookie").iter() {
270 if let Ok(cookie_str) = v.to_str() {
271 let _ = store.parse(cookie_str, url);
272 }
273 }
274 }
275 }
276}
277
278#[derive(Default)]
280pub struct ClientBuilder {
281 profile: Option<ChromeProfile>,
282 proxy: Option<Proxy>,
283 cookie_store: Option<Arc<RwLock<CookieStore>>>,
284}
285
286impl ClientBuilder {
287 pub fn profile(mut self, profile: ChromeProfile) -> Self {
289 self.profile = Some(profile);
290 self
291 }
292
293 pub fn proxy(mut self, proxy: Proxy) -> Self {
295 self.proxy = Some(proxy);
296 self
297 }
298
299 pub fn cookie_store(mut self, store: Arc<RwLock<CookieStore>>) -> Self {
301 self.cookie_store = Some(store);
302 self
303 }
304
305 pub fn build(self) -> Result<Client> {
307 let profile = self
308 .profile
309 .unwrap_or_else(crate::profile::chrome_134::profile_auto);
310
311 Ok(Client {
312 pool: Arc::new(Mutex::new(HashMap::new())),
313 profile,
314 proxy: self.proxy,
315 cookie_store: self
316 .cookie_store
317 .unwrap_or_else(|| Arc::new(RwLock::new(CookieStore::default()))),
318 })
319 }
320}