proxy_server/
lib.rs

1//! Low level proxy server.
2//! To implement request proxying, only standard [`TcpStream`] was used without additional libraries
3//! # Examples
4//! With default params:
5//! ```no_run
6//! use proxy_server::Builder;                                                                       
7//!      
8//! fn main() {
9//!     Builder::new().bind(None).expect("Error in proxy");
10//! }
11//! ```
12//! With custom params:
13//! ```no_run
14//! use proxy_server::{log::LogLevel, Builder};
15//!
16//! fn main() {
17//!     Builder::new()
18//!         .with_address("127.0.0.1:3000")
19//!         .with_target("127.0.0.1:3001")
20//!         .with_log_level(LogLevel::Warn)
21//!         .with_threads(4)
22//!         .bind(None)
23//!         .expect("Error in proxy");
24//! }
25//! ```
26//! With check and change target if needed on every request
27//! ```no_run
28//!
29//! use proxy_server::{Builder, ChangeTarget};                  
30//! fn get_actual_target(old: &str) -> &'static str {     
31//!     let target1 = "127.0.0.1:3001";                   
32//!     let target2 = "127.0.0.1:3003";
33//!     let res = match old {                                 
34//!         "127.0.0.1:3001" => target2,
35//!         "127.0.0.1:3003" => target1,
36//!         _ => target1,
37//!         };
38//!         res
39//! }
40//!
41//! fn main() {
42//!     let cb: ChangeTarget = |old| get_actual_target(old);                                          
43//!     Builder::new()                          
44//!         .bind(Some(cb))
45//!         .expect("Error in proxy");
46//! }
47//! ```
48
49use std::{
50    convert::Infallible,
51    io::{Error, ErrorKind, Result, Write},
52    net::{TcpListener, TcpStream},
53    str,
54    thread::sleep,
55    time::Duration,
56};
57mod thread_pool;
58use thread_pool::ThreadPool;
59pub mod http;
60use http::Http;
61
62pub mod log;
63use log::{Log, LogLevel, LOG_LEVEL};
64pub mod prelude;
65use prelude::constants::*;
66
67use crate::http::{
68    headers::Headers,
69    request::{Request, Socket},
70    status::Status,
71};
72
73#[cfg(test)]
74mod tests;
75
76/// Callback function for change target on fly.
77/// Use only fast method because this function if it provieded then run every request again.
78pub type ChangeTarget = fn(&'static str) -> &'static str;
79
80/// Structure for proxy server configuration
81#[derive(Clone, Copy, Debug)]
82pub struct Builder {
83    pub address: &'static str,
84    pub target: &'static str,
85    pub log_level: LogLevel,
86    pub threads: usize,
87}
88
89impl Builder {
90    /// Create new proxy server builder
91    pub fn new() -> Self {
92        Self {
93            address: PROXY_ADDRESS,
94            target: TARGET_ADDRESS,
95            log_level: LOG_LEVEL,
96            threads: THREADS,
97        }
98    }
99
100    /// Set proxy server address
101    pub fn with_address(mut self, address: &'static str) -> Self {
102        self.address = address;
103        self
104    }
105
106    /// Set proxy server target address
107    pub fn with_target(mut self, target: &'static str) -> Self {
108        self.target = target;
109        self
110    }
111
112    /// Set log level of proxy server
113    pub fn with_log_level(mut self, log_level: LogLevel) -> Self {
114        self.log_level = log_level;
115        self
116    }
117
118    /// Set proxy server count of used threads
119    pub fn with_threads(mut self, threads: usize) -> Self {
120        self.threads = threads;
121        self
122    }
123
124    /// Proxy server listener releasing [`std::net::TcpListener`] via thread pool
125    pub fn bind(mut self, cb: Option<ChangeTarget>) -> Result<Infallible> {
126        let listener = TcpListener::bind(&self.address)?;
127
128        let _log = Log::new(&self.log_level);
129        println!(
130            "Listening: {}; Target: {}; Chunk size: {}KB; Log level: {:?}",
131            &self.address, &self.target, CHUNK_SIZE, &self.log_level
132        );
133
134        let pool = ThreadPool::new(self.threads);
135        for stream in listener.incoming() {
136            if let Err(err) = stream {
137                println!("Error in incoming stream {:?}", err);
138                continue;
139            }
140
141            if let Some(func) = cb {
142                self.target = func(&self.target);
143            }
144            let cl = Handler::new(self);
145            pool.execute(|| {
146                let stream = stream.unwrap();
147                let res = cl.handle_proxy(stream);
148                if let Err(err) = res {
149                    println!("Error in handle proxy {:?}", err);
150                }
151            });
152        }
153        Err(Error::new(ErrorKind::Interrupted, "main thread crashed"))
154    }
155}
156
157struct Handler {
158    config: Builder,
159}
160
161impl Handler {
162    fn new(config: Builder) -> Self {
163        Self { config }
164    }
165
166    fn handle_proxy(self, client: TcpStream) -> Result<()> {
167        const TAG: &str = "Handle proxy";
168        let _log = Log::new(&self.config.log_level);
169
170        _log.println(LogLevel::Info, TAG, "client", &client);
171
172        let mut client = Http::from(client);
173
174        let head_client_buf = client.read_headers()?;
175
176        let error = client.socket.take_error().unwrap();
177        let error = match error {
178            None => "".to_string(),
179            Some(val) => val.to_string(),
180        };
181
182        let mut req_client = Request::new(
183            Socket {
184                host: client.socket.local_addr().unwrap().to_string(),
185                peer_addr: client.socket.peer_addr().unwrap().to_string(),
186                ttl: client.socket.ttl().unwrap(),
187                error,
188            },
189            head_client_buf,
190        )?;
191
192        _log.println(LogLevel::Info, TAG, "client request", &req_client);
193        req_client.change_host(&self.config.target)?;
194
195        let http = Http::connect(&self.config.target);
196        if let Err(e) = &http {
197            _log.println(LogLevel::Warn, TAG, "Failed proxy", e);
198            client.write(
199                Headers::new_response(&Status::new(502), vec![])
200                    .raw
201                    .as_bytes(),
202            )?;
203            client.flush()?;
204            sleep(Duration::from_millis(100));
205            return Ok(());
206        }
207        let mut http = http?;
208
209        http.write(req_client.headers.raw.as_bytes())?;
210
211        if req_client.content_length != 0 {
212            let body = client.read_body(&req_client)?;
213            _log.println(
214                LogLevel::Info,
215                TAG,
216                "request body",
217                str::from_utf8(&body).unwrap(),
218            );
219            http.write(&body)?;
220        }
221
222        let h = http.read_headers()?;
223
224        let error = client.socket.take_error().unwrap();
225        let error = match error {
226            None => "".to_string(),
227            Some(val) => val.to_string(),
228        };
229
230        let req_http = Request::new(
231            Socket {
232                host: client.socket.local_addr().unwrap().to_string(),
233                peer_addr: client.socket.peer_addr().unwrap().to_string(),
234                ttl: client.socket.ttl().unwrap(),
235                error,
236            },
237            h.clone(),
238        )?;
239        _log.println(LogLevel::Info, TAG, "target response", &req_http);
240        client.write(req_http.headers.raw.as_bytes())?;
241        client.tunnel(&mut http, &_log)?;
242
243        Ok(())
244    }
245}
proxy_server/lib.rs

proxy_server/
lib.rs
