Skip to main content

br_web_server/
client.rs

1use crate::client_response::ClientResponse;
2use crate::stream::ClientStream;
3use crate::url::Url;
4use crate::{Method, Protocol};
5use json::JsonValue;
6use log::debug;
7use rand::distr::Alphanumeric;
8use rand::Rng;
9use rustls::RootCertStore;
10use rustls::{ClientConnection, StreamOwned};
11use std::collections::HashMap;
12use std::io::Read;
13use std::net::{IpAddr, TcpStream};
14use std::path::Path;
15use std::str::FromStr;
16use std::sync::Arc;
17use std::time::Duration;
18use std::{fs, thread};
19
20const REDIRECT_STATUS_CODES: [u16; 3] = [301, 302, 303];
21
22pub struct Client {
23    debug: bool,
24    url: String,
25    protocol: Protocol,
26    method: Method,
27    query: HashMap<String, String>,
28    headers: HashMap<String, JsonValue>,
29    cookies: HashMap<String, String>,
30    body: Vec<u8>,
31    read_timeout: u64,
32    write_timeout: u64,
33    auto_location: bool,
34    boundary: String,
35    /// 重试次数
36    retry: usize,
37    ca_pem: String,
38    pkcs12: String,
39    pkcs12_password: String,
40}
41impl Default for Client {
42    fn default() -> Self {
43        Self::new()
44    }
45}
46
47impl Client {
48    pub fn new() -> Self {
49        Self {
50            debug: false,
51            url: "".to_string(),
52            protocol: Protocol::HTTP1_1,
53            method: Method::Other("".to_string()),
54            query: HashMap::new(),
55            headers: Default::default(),
56            cookies: Default::default(),
57            body: vec![],
58            read_timeout: 10,
59            write_timeout: 10,
60            auto_location: false,
61            boundary: "".to_string(),
62            retry: 0,
63            ca_pem: "".to_string(),
64            pkcs12: "".to_string(),
65            pkcs12_password: "".to_string(),
66        }
67    }
68    /// 重置全部
69    pub fn reset(&mut self) -> &mut Self {
70        self.debug = false;
71        self.url = "".to_string();
72        self.protocol = Protocol::HTTP1_1;
73        self.method = Method::Other("".to_string());
74        self.query = HashMap::new();
75        self.headers = Default::default();
76        self.cookies = Default::default();
77        self.body = vec![];
78        self.read_timeout = 0;
79        self.write_timeout = 0;
80        self.auto_location = false;
81        self.boundary = "".to_string();
82        self.retry = 0;
83        self
84    }
85    pub fn debug(&mut self) -> &mut Self {
86        self.debug = true;
87        self
88    }
89    pub fn ca_pem(&mut self, ca_pem: &str) -> &mut Self {
90        self.ca_pem = ca_pem.to_string();
91        self
92    }
93    pub fn pkcs12(&mut self, pkcs12: &str, password: &str) -> &mut Self {
94        self.pkcs12 = pkcs12.to_string();
95        self.pkcs12_password = password.to_string();
96        self
97    }
98    pub fn write_timeout(&mut self, time: u64) -> &mut Self {
99        self.write_timeout = time;
100        self
101    }
102    pub fn read_timeout(&mut self, time: u64) -> &mut Self {
103        self.read_timeout = time;
104        self
105    }
106    pub fn http10(&mut self) -> &mut Self {
107        self.protocol = Protocol::HTTP1_0;
108        self
109    }
110    pub fn http11(&mut self) -> &mut Self {
111        self.protocol = Protocol::HTTP1_1;
112        self
113    }
114    pub fn http2(&mut self) -> &mut Self {
115        self.protocol = Protocol::HTTP2;
116        self
117    }
118    /// 设置重置
119    pub fn retry(&mut self, count: usize) -> &mut Self {
120        self.retry = count;
121        self
122    }
123    pub fn auto_location(&mut self) -> &mut Self {
124        self.auto_location = true;
125        self
126    }
127    fn generate_boundary(&mut self) -> String {
128        let rand: String = rand::rng()
129            .sample_iter(&Alphanumeric)
130            .take(16)
131            .map(char::from)
132            .collect();
133        format!("----WebKitFormBoundary{}", rand)
134    }
135    pub fn get(&mut self, url: &str) -> &mut Self {
136        self.url = url.to_string();
137        self.method = Method::GET;
138        self
139    }
140    pub fn post(&mut self, url: &str) -> &mut Self {
141        self.url = url.to_string();
142        self.method = Method::POST;
143        self
144    }
145    pub fn head(&mut self, url: &str) -> &mut Self {
146        self.url = url.to_string();
147        self.method = Method::HEAD;
148        self
149    }
150    pub fn delete(&mut self, url: &str) -> &mut Self {
151        self.url = url.to_string();
152        self.method = Method::DELETE;
153        self
154    }
155    pub fn method(&mut self, method: &str, url: &str) -> &mut Self {
156        self.url = url.to_string();
157        self.method = Method::Other(method.to_string());
158        self
159    }
160    pub fn header(&mut self, key: &str, value: JsonValue) -> &mut Self {
161        self.headers.insert(key.to_string(), value);
162        self
163    }
164    pub fn query(&mut self, key: &str, value: JsonValue) -> &mut Self {
165        let value = br_crypto::encoding::urlencoding_encode(value.to_string().as_str());
166        self.query.insert(key.to_string(), value.to_string());
167        self
168    }
169    /// 请求JSON数据
170    pub fn raw_json(&mut self, data: JsonValue) -> &mut Self {
171        self.header("Content-Type", "application/json".into());
172        self.body = data.to_string().as_bytes().to_vec();
173        self.header("Content-Length", self.body.len().into());
174        self
175    }
176    /// 请求FormData数据
177    pub fn form_data(&mut self, data: JsonValue) -> &mut Self {
178        self.boundary = self.generate_boundary();
179        self.header(
180            "Content-Type",
181            format!("multipart/form-data; boundary={}", self.boundary).into(),
182        );
183        let mut body = vec![];
184        for (key, value) in data.entries() {
185            let file = Path::new(value.as_str().unwrap_or(""));
186            if file.is_file() {
187                let mut t = vec![];
188                let r = match fs::File::open(file).and_then(|mut f| f.read(&mut t)) {
189                    Ok(r) => r,
190                    Err(_) => continue,
191                };
192                let filename = file
193                    .file_name()
194                    .and_then(|n| n.to_str())
195                    .unwrap_or("unknown")
196                    .to_string();
197                let content_type = if let Some(kind) = infer::get(&t[..r]) {
198                    kind.mime_type()
199                } else {
200                    "application/octet-stream"
201                };
202                body.extend(format!("--{}\r\n", self.boundary).as_bytes().to_vec());
203                body.extend(
204                    format!(
205                        r#"Content-Disposition: form-data; name="{}"; filename="{}"{}"#,
206                        key, filename, "\r\n"
207                    )
208                    .as_bytes()
209                    .to_vec(),
210                );
211                body.extend(
212                    format!(r#"Content-Type: {}{}"#, content_type, "\r\n\r\n")
213                        .as_bytes()
214                        .to_vec(),
215                );
216                body.extend(&t[..r]);
217                body.extend("\r\n".to_string().as_bytes().to_vec());
218            } else {
219                body.extend(format!("--{}\r\n", self.boundary).as_bytes().to_vec());
220                body.extend(
221                    format!(
222                        r#"Content-Disposition: form-data; name="{}"{}"#,
223                        key, "\r\n\r\n"
224                    )
225                    .as_bytes()
226                    .to_vec(),
227                );
228                body.extend(format!("{value}\r\n").as_bytes().to_vec());
229            }
230        }
231        body.extend(format!("--{}--\r\n", self.boundary).as_bytes().to_vec());
232        self.body = body.to_vec();
233        self.header("Content-Length", self.body.len().into());
234        self
235    }
236    /// 表单urlencoded请求
237    pub fn x_www_form_urlencoded(&mut self, data: JsonValue) -> &mut Self {
238        self.header("Content-Type", "application/x-www-form-urlencoded".into());
239        let mut params = vec![];
240        for (key, value) in data.entries() {
241            params.push(format!("{}={}", key, value));
242        }
243        let params = params.join("&");
244        self.body = params.as_bytes().to_vec();
245        self.header("Content-Length", self.body.len().into());
246        self
247    }
248    pub fn send(&mut self) -> Result<ClientResponse, String> {
249        let mut url = Url::parse(&self.url)?;
250        let stream = match TcpStream::connect(format!("{}:{}", url.host, url.port)) {
251            Ok(stream) => stream,
252            Err(e) => {
253                if self.retry > 0 {
254                    self.retry -= 1;
255                    return self.send();
256                }
257                return Err(e.to_string());
258            }
259        };
260
261        if self.read_timeout > 0 {
262            let _ = stream.set_read_timeout(Option::from(Duration::from_secs(self.read_timeout)));
263        }
264        if self.write_timeout > 0 {
265            let _ = stream.set_write_timeout(Option::from(Duration::from_secs(self.write_timeout)));
266        }
267
268        let mut stream = if url.scheme == "https" {
269            let root_store =
270                RootCertStore::from_iter(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
271            let mut config = rustls::ClientConfig::builder()
272                .with_root_certificates(root_store)
273                .with_no_client_auth();
274            // 需要小写 http/1.1
275            config.alpn_protocols = vec![self.protocol.str().to_lowercase().as_bytes().to_vec()];
276            let example_com = url
277                .host
278                .clone()
279                .try_into()
280                .map_err(|e| format!("invalid DNS name: {}", e))?;
281            let client = ClientConnection::new(Arc::new(config), example_com)
282                .map_err(|e| format!("TLS connection error: {}", e))?;
283            ClientStream::Https(Box::new(StreamOwned::new(client, stream)))
284        } else {
285            ClientStream::Http(stream)
286        };
287
288        if !self.query.is_empty() {
289            url.set_query(self.query.clone());
290        }
291
292        let request_line = format!(
293            "{} {} {}",
294            self.method.str(),
295            url.path_query.as_str(),
296            self.protocol.str()
297        );
298
299        let host = if IpAddr::from_str(url.host.as_str()).is_ok() {
300            format!("Host: {}:{}", url.host, url.port)
301        } else {
302            format!("Host: {}", url.host)
303        };
304
305        let mut header = vec![request_line, host];
306
307        // 拼接 Header 头
308        for (key, value) in self.headers.iter() {
309            header.push(format!("{}: {}", key, value));
310        }
311
312        // 拼接 Cookie 头
313        for (key, value) in self.cookies.iter() {
314            header.push(format!("Set-Cookie: {}={}", key, value));
315        }
316
317        // 消息体隔离符号
318        header.push("\r\n".to_string());
319        match stream.write_all(header.join("\r\n").as_bytes()) {
320            Ok(()) => {
321                match stream.flush() {
322                    Ok(e) => e,
323                    Err(e) => return Err(e.to_string()),
324                };
325            }
326            Err(e) => {
327                if self.retry > 0 {
328                    self.retry -= 1;
329                    return self.send();
330                }
331                return Err(e.to_string());
332            }
333        };
334
335        if self.debug {
336            debug!("\r\n==============请求头 {:?}============== \r\n{}\r\n============================",thread::current().id(), header.join("\r\n"));
337        }
338
339        if !self.body.is_empty() {
340            match stream.write_all(self.body.as_slice()) {
341                Ok(()) => {
342                    match stream.flush() {
343                        Ok(e) => e,
344                        Err(e) => return Err(e.to_string()),
345                    };
346                }
347                Err(e) => {
348                    if self.retry > 0 {
349                        self.retry -= 1;
350                        return self.send();
351                    }
352                    return Err(e.to_string());
353                }
354            }
355        }
356
357        if self.debug && self.boundary.is_empty() {
358            debug!("\r\n==============请求体 {:?}============== \r\n{}\r\n============================",thread::current().id(),String::from_utf8_lossy(self.body.as_slice()));
359        }
360
361        let mut res = ClientResponse::new(self.debug, stream).handle()?;
362        if REDIRECT_STATUS_CODES.contains(&res.status()) && !res.location().is_empty() {
363            self.url = res.location();
364            return self.send();
365        }
366        self.reset();
367        Ok(res)
368    }
369}