Skip to content

Instantly share code, notes, and snippets.

@gauravssnl
Last active August 30, 2022 17:12
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save gauravssnl/9a03df7f650a206d67f9a562ed2be79e to your computer and use it in GitHub Desktop.
Save gauravssnl/9a03df7f650a206d67f9a562ed2be79e to your computer and use it in GitHub Desktop.
Proxy server in progress
[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"
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