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