Last active
August 30, 2022 17:12
-
-
Save gauravssnl/9a03df7f650a206d67f9a562ed2be79e to your computer and use it in GitHub Desktop.
Proxy server in progress
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
[package] | |
name = "tokio-server" | |
version = "0.1.0" | |
edition = "2021" | |
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html | |
[dependencies] | |
tokio = { version = "1.20.0", features = ["full"] } | |
tokio-stream = "0.1.9" | |
tokio-util = { version = "0.7.0", features = ["codec"] } | |
futures = "0.3.22" | |
bytes = "1.2.1" | |
urlparse = "0.7.3" |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
use std::{collections::HashMap, fmt::format, io::BufRead, ops::Index}; | |
use tokio::{ | |
io::{self, AsyncReadExt, AsyncWriteExt, Interest}, | |
net::{TcpListener, TcpStream}, | |
}; | |
const HOST: &'static str = "Host"; | |
#[tokio::main] | |
async fn main() -> Result<(), Box<dyn std::error::Error>> { | |
let address = "127.0.01:8081"; | |
let listener = TcpListener::bind(address).await?; | |
println!("Serving on {}", address); | |
loop { | |
let mut client = listener.accept().await?.0; | |
tokio::spawn(async move { handle_client(&mut client).await }); | |
} | |
} | |
async fn handle_client( | |
mut client: &mut TcpStream, | |
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> { | |
// println!("Handling client request"); | |
let request_buffer = read_stream(&mut client).await?; | |
println!( | |
"******************* Request Received: *****************\n{}", | |
String::from_utf8_lossy(&request_buffer) | |
); | |
let request = Request::from(request_buffer); | |
println!("request: {:?}", request); | |
read_and_write_data_from_remote_server(client, request).await?; | |
Ok(()) | |
} | |
async fn read_stream( | |
stream: &mut TcpStream, | |
) -> Result<Vec<u8>, Box<dyn std::error::Error + Send + Sync>> { | |
let mut buffer: Vec<u8> = Vec::new(); | |
let ready = stream.ready(Interest::READABLE).await?; | |
println!("Stream: {:?}", stream); | |
println!("Ready to read: {:?}", ready); | |
if ready.is_readable() { | |
let buffer_size: usize = 1024; | |
loop { | |
let mut fixed_buffer = vec![0; buffer_size]; | |
match stream.read(&mut fixed_buffer).await { | |
Ok(n) if n == 0 => break, | |
Ok(n) if n < buffer_size => { | |
buffer.append(&mut fixed_buffer[..n].to_vec()); | |
break; | |
} | |
Ok(_) => { | |
buffer.append(&mut fixed_buffer); | |
} | |
Err(e) => { | |
println!("Error in reading stram data: {}", e); | |
break; | |
} | |
} | |
} | |
} | |
Ok(buffer) | |
} | |
#[derive(Debug)] | |
enum Body { | |
Empty, | |
} | |
#[derive(Debug)] | |
struct Request { | |
method: String, | |
version: String, | |
headers: HashMap<String, String>, | |
url_raw: String, | |
url: urlparse::Url, | |
body: Body, | |
raw_data: Vec<u8>, | |
host: String, | |
port: i32, | |
} | |
impl Request { | |
fn new() -> Self { | |
Request { | |
method: String::new(), | |
version: String::new(), | |
headers: HashMap::new(), | |
url_raw: String::new(), | |
url: urlparse::Url::new(), | |
body: Body::Empty, | |
raw_data: vec![], | |
host: String::new(), | |
port: 0, | |
} | |
} | |
} | |
impl From<Vec<u8>> for Request { | |
fn from(buffer: Vec<u8>) -> Self { | |
// let reader = buffer.reader(); | |
// let mut lines = reader.lines(); | |
let cursor = std::io::Cursor::new(&buffer); | |
let mut lines_iter = cursor.lines().map(|l| l.unwrap()); | |
let first_line = lines_iter.next().unwrap(); | |
let first_line_vec: Vec<&str> = first_line.split(" ").collect(); | |
let method = first_line_vec.get(0).unwrap().to_string(); | |
let url_raw = first_line_vec.get(1).unwrap().to_string(); | |
let version = first_line_vec.get(2).unwrap().to_string(); | |
let headers = get_request_headers(&mut lines_iter); | |
let body = Body::Empty; | |
let url = urlparse::urlparse(&url_raw); | |
let mut host: String = headers.get(HOST).unwrap().into(); | |
let mut port = if method == "CONNECT" { 443 } else { 80 }; | |
if host.contains(":") { | |
let index = host.find(":").unwrap(); | |
host = host[..index].to_string(); | |
port = host[index + 1..].parse().unwrap(); | |
} | |
Request { | |
method, | |
url_raw, | |
url, | |
version, | |
headers, | |
body, | |
raw_data: buffer, | |
host, | |
port, | |
} | |
} | |
} | |
async fn read_and_write_data_from_remote_server( | |
mut client: &mut TcpStream, | |
request: Request, | |
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> { | |
let address = format!("{}:{}", request.host, request.port); | |
println!("Connecting to ({})", address); | |
let mut remote = TcpStream::connect(address).await?; | |
if request.method == "CONNECT" { | |
let empty_response = format!( | |
"{} 200 Connection Established\r\nProxy-Agent: SimpleProxy/1.1\r\n\r\n", | |
request.version | |
); | |
println!( | |
"***** Sending empty Response to client: ********\n{}", | |
empty_response | |
); | |
client.write_all(empty_response.as_bytes()).await?; | |
match io::copy_bidirectional(&mut client, &mut remote).await { | |
Ok((from_client, to_client)) => { | |
println!( | |
"Client wrote {} bytes and received {} bytes", | |
from_client, to_client | |
); | |
} | |
Err(e) => println!("Error in copying bi-directionally: {}", e), | |
} | |
// remote.flush().await?; | |
// client.flush().await?; | |
// println!("Response sent for CONNECT"); | |
} else { | |
println!("Handling non-HTTPS request"); | |
match remote.write(&request.raw_data).await { | |
Ok(n) => println!( | |
"Wrote {} bytes and data to remote: {:?}", | |
n, | |
String::from_utf8_lossy(&request.raw_data) | |
), | |
Err(e) => println!("Write error in remote: {}", e), | |
} | |
match read_stream(&mut remote).await { | |
Ok(response) => { | |
println!( | |
"Read {} bytes and data from server: {:?}", | |
response.len(), | |
String::from_utf8_lossy(&response) | |
); | |
match client.write(&response).await { | |
Ok(n) => println!( | |
"Wrote {} bytes and data to client: {:?}", | |
n, | |
String::from_utf8_lossy(&response) | |
), | |
Err(e) => println!("Write error in client: {}", e), | |
} | |
} | |
Err(e) => println!("Write error in client: {}", e), | |
} | |
} | |
// remote.flush().await?; | |
// client.flush().await?; | |
println!("Response sent successfully"); | |
Ok(()) | |
} | |
fn get_request_headers(lines_iter: &mut impl Iterator<Item = String>) -> HashMap<String, String> { | |
let mut headers = HashMap::new(); | |
loop { | |
match lines_iter.next() { | |
None => break, | |
Some(line) => { | |
// println!("line: {:?}", line); | |
if line == "" { | |
break; | |
} | |
let header_line_vec: Vec<String> = line.split(":").map(|x| x.to_string()).collect(); | |
let header_key = header_line_vec.get(0).unwrap().to_string(); | |
let header_value = header_line_vec.get(1).unwrap().trim().to_string(); | |
headers.insert(header_key, header_value); | |
} | |
} | |
} | |
headers | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment