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