atlas_http/
client.rs

1#![allow(clippy::large_enum_variant)]
2
3use super::{
4    HttpBody, HttpClientConfig, HttpRequest, HttpResponse, HttpSyncClient, ProxyType,
5};
6use crate::client_builder::HttpClientBuilder;
7use crate::error::{Error, FileNotCreatedError, InvalidResponseError};
8use crate::socks5;
9use rustls::pki_types::ServerName;
10use std::fs::File;
11use std::io::{BufRead, BufReader, Read, Write};
12use std::net::{TcpStream, ToSocketAddrs};
13use std::path::Path;
14use std::sync::Arc;
15use std::time::Duration;
16use url::Url;
17
18#[derive(Debug, Clone)]
19pub struct HttpClient {
20    pub config: HttpClientConfig,
21}
22
23impl HttpClient {
24    pub fn new(config: &HttpClientConfig) -> Self {
25        Self {
26            config: config.clone(),
27        }
28    }
29
30    /// Instantiate HTTP client builder
31    pub fn builder() -> HttpClientBuilder {
32        HttpClientBuilder::new()
33    }
34
35    /// Send HTTP request, and return response
36    pub async fn send(&mut self, req: &HttpRequest) -> Result<HttpResponse, Error> {
37        self.send_request(req, &String::new()).await
38    }
39
40    /// Download a file
41    pub async fn download(&mut self, url: &str, dest_file: &str) -> Result<HttpResponse, Error> {
42        let req = HttpRequest::new("GET", url, &vec![], &HttpBody::empty());
43        self.send_request(&req, &dest_file.to_string()).await
44    }
45
46    /// Send GET request
47    pub async fn get(&mut self, url: &str) -> Result<HttpResponse, Error> {
48        let req = HttpRequest::new("GET", url, &Vec::new(), &HttpBody::empty());
49        self.send_request(&req, &String::new()).await
50    }
51
52    /// Send POST request
53    pub async fn post(&mut self, url: &str, body: &HttpBody) -> Result<HttpResponse, Error> {
54        let req = HttpRequest::new("POST", url, &Vec::new(), body);
55        self.send_request(&req, &String::new()).await
56    }
57
58    /// Send PUT request
59    pub async fn put(&mut self, url: &str, data: &[u8]) -> Result<HttpResponse, Error> {
60        let req = HttpRequest::new("PUT", url, &Vec::new(), &HttpBody::from_raw(data));
61        self.send_request(&req, &String::new()).await
62    }
63
64    /// Send DELETE request
65    pub async fn delete(&mut self, url: &str) -> Result<HttpResponse, Error> {
66        let req = HttpRequest::new("DELETE", url, &Vec::new(), &HttpBody::empty());
67        self.send_request(&req, &String::new()).await
68    }
69
70    /// Send OPTIONS request
71    pub async fn options(&mut self, url: &str) -> Result<HttpResponse, Error> {
72        let req = HttpRequest::new("OPTIONS", url, &Vec::new(), &HttpBody::empty());
73        self.send_request(&req, &String::new()).await
74    }
75
76    /// Send HEAD request
77    pub async fn head(&mut self, url: &str) -> Result<HttpResponse, Error> {
78        let req = HttpRequest::new("HEAD", url, &Vec::new(), &HttpBody::empty());
79        self.send_request(&req, &String::new()).await
80    }
81
82    // Send request, used internally by the other methods.
83    async fn send_request(
84        &mut self,
85        req: &HttpRequest,
86        dest_file: &String,
87    ) -> Result<HttpResponse, Error> {
88        // Prepare uri and http message
89        let (uri, port, message) = req.prepare(&self.config)?;
90
91        // Connect
92        let mut reader = self.connect(&uri, &port, &message).await?;
93
94        // Read header
95        let mut res = HttpResponse::read_header(&mut reader, req, dest_file)?;
96        self.config.cookie.update_jar(&res.headers());
97
98        // Check follow location
99        if self.config.follow_location && res.headers().has_lower("location") {
100            res = self.follow(&res, dest_file)?;
101        }
102
103        // Return if not downloading a file
104        if dest_file.is_empty() {
105            return Ok(res);
106        }
107
108        // Save output file
109        let dest_path = Path::new(&dest_file);
110        let mut fh = match File::create(dest_path) {
111            Ok(r) => r,
112            Err(e) => {
113                return Err(Error::FileNotCreated(FileNotCreatedError {
114                    filename: dest_file.to_string(),
115                    error: e.to_string(),
116                }));
117            }
118        };
119
120        // Save file
121        let mut buffer = [0u8; 2048];
122        loop {
123            let bytes_read = match reader.read(&mut buffer) {
124                Ok(r) => r,
125                Err(e) => {
126                    return Err(Error::NoRead(InvalidResponseError {
127                        url: req.url.clone(),
128                        response: e.to_string(),
129                    }));
130                }
131            };
132
133            if bytes_read == 0 {
134                break;
135            }
136            fh.write_all(&buffer).unwrap();
137        }
138
139        Ok(res)
140    }
141
142    /// Check redirect if follow_location enabled
143    fn follow(&self, res: &HttpResponse, dest_file: &String) -> Result<HttpResponse, Error> {
144        let redirect_url = res.headers().get_lower("location").unwrap();
145        let mut rhttp = HttpSyncClient::new(&self.config.clone());
146
147        let next_res = if dest_file.is_empty() {
148            rhttp.get(&redirect_url.clone())?
149        } else {
150            rhttp.download(&redirect_url.clone(), dest_file)?
151        };
152
153        Ok(next_res)
154    }
155
156    // Connect to remote server
157    pub async fn connect(
158        &self,
159        uri: &Url,
160        port: &u16,
161        message: &[u8],
162    ) -> Result<Box<dyn BufRead>, Error> {
163        // Prepare uri
164        let hostname =
165            if self.config.proxy_type != ProxyType::None && !self.config.proxy_host.is_empty() {
166                format!("{}:{}", self.config.proxy_host, self.config.proxy_port)
167            } else {
168                format!("{}:{}", &uri.host_str().unwrap(), port)
169            };
170        let mut address = hostname.to_socket_addrs().unwrap();
171        let addr = address.next().unwrap();
172
173        // Open tcp stream
174        let mut sock =
175            match TcpStream::connect_timeout(&addr, Duration::from_secs(self.config.timeout)) {
176                Ok(r) => r,
177                Err(_e) => {
178                    return Err(Error::NoConnect(hostname.clone()));
179                }
180            };
181        sock.set_nodelay(true).unwrap();
182
183        // SOCKs5 connection, if needed
184        if self.config.proxy_type == ProxyType::SOCKS5 {
185            socks5::connect(&mut sock, &self.config, uri, port);
186        }
187
188        // Connect over SSL, if needed
189        if uri.scheme() == "https" && self.config.proxy_type != ProxyType::HTTP {
190            let dns_name = ServerName::try_from(uri.host_str().unwrap())
191                .unwrap()
192                .to_owned();
193            let conn = rustls::ClientConnection::new(Arc::clone(&self.config.tls_config), dns_name)
194                .unwrap();
195
196            let mut tls_stream = rustls::StreamOwned::new(conn, sock);
197            tls_stream.flush().unwrap();
198            tls_stream.write_all(message).unwrap();
199
200            let reader = BufReader::with_capacity(2048, tls_stream);
201            return Ok(Box::new(reader));
202        }
203
204        // Get reader
205        sock.write_all(message).unwrap();
206        let reader = BufReader::with_capacity(2048, sock);
207
208        Ok(Box::new(reader))
209    }
210}
atlas_http/client.rs

atlas_http/
client.rs
