Add support for Connect method, allowing https proxying

Author: Az107

examples/proxy_con.rs

@@ -0,0 +1,18 @@
+use hteapot::{Hteapot, HttpMethod, HttpRequest, HttpResponse, TunnelResponse, headers};
+
+fn main() {
+    let server = Hteapot::new("0.0.0.0", 8081);
+    server.listen(move |req: HttpRequest| {
+        println!("New request to {} {}!", req.method.to_str(), &req.path);
+        if req.method == HttpMethod::CONNECT {
+            TunnelResponse::new(&req.path)
+        } else {
+            let secure_path = req.path.replace("http", "https");
+            HttpResponse::new(
+                hteapot::HttpStatus::MovedPermanently,
+                "Moved",
+                headers! {"location" => secure_path},
+            )
+        }
+    });
+}

src/hteapot/engine.rs

@@ -279,7 +279,7 @@ impl Hteapot {
                 response
                     .base()
                     .headers
-                    .entry("Connection".to_string())
+                    .entry("connection".to_string())
                     .or_insert("keep-alive".to_string());
                 response.base().headers.insert(
                     "Keep-Alive".to_string(),
@@ -293,8 +293,10 @@ impl Hteapot {
             }
             status.write = true;
             status.response = response;
+            status.response.set_stream(&socket_data.stream);
         }
 
+        // Seting the stream in case is needed for the response, (example: streaming)
         // Write the response to the client in chunks
         loop {
             match status.response.peek() {

src/hteapot/mod.rs

@@ -31,7 +31,7 @@ pub use self::request::HttpRequest;
 pub use engine::Hteapot;
 pub use methods::HttpMethod;
 
-pub use response::{HttpResponse, StreamedResponse};
+pub use response::{HttpResponse, StreamedResponse, TunnelResponse};
 pub use status::HttpStatus;
 
 /// Crate version as set by `Cargo.toml`.

src/hteapot/response.rs

@@ -7,14 +7,19 @@
 //!
 //! All response types implement the [`HttpResponseCommon`] trait.
 
+use crate::headers;
+
 use super::HttpStatus;
 use super::{BUFFER_SIZE, VERSION};
 use std::collections::{HashMap, VecDeque};
+use std::io::Write;
+use std::net::TcpStream;
 use std::sync::Arc;
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::mpsc::{self, Receiver, SendError, Sender, TryRecvError};
-use std::thread;
 use std::thread::JoinHandle;
+use std::time::Duration;
+use std::{io, thread};
 
 /// Basic HTTP status line + headers.
 pub struct BaseResponse {
@@ -62,6 +67,8 @@ pub trait HttpResponseCommon {
 
     /// Advances and returns the next chunk of the response body.
     fn peek(&mut self) -> Result<Vec<u8>, IterError>;
+
+    fn set_stream(&mut self, stream: &TcpStream);
 }
 
 /// Error returned during response iteration.
@@ -165,6 +172,10 @@ impl HttpResponseCommon for HttpResponse {
         let byte_chunk = raw.next().ok_or(IterError::Finished)?.to_vec();
         return Ok(byte_chunk);
     }
+
+    fn set_stream(&mut self, _: &TcpStream) {
+        ()
+    }
 }
 
 /// Dummy response used when nothing needs to be returned.
@@ -183,6 +194,10 @@ impl HttpResponseCommon for EmptyHttpResponse {
     fn peek(&mut self) -> Result<Vec<u8>, IterError> {
         Err(IterError::Finished)
     }
+
+    fn set_stream(&mut self, _: &TcpStream) {
+        ()
+    }
 }
 
 /// Sends response chunks in a `Transfer-Encoding: chunked` format.
@@ -292,4 +307,80 @@ impl HttpResponseCommon for StreamedResponse {
             self.queue.pop_front().ok_or(IterError::WouldBlock)
         }
     }
+
+    fn set_stream(&mut self, _: &TcpStream) {
+        ()
+    }
+}
+
+pub struct TunnelResponse {
+    base: BaseResponse,
+    addr: String,
+    has_end: Arc<AtomicBool>,
+    stream_in: Option<TcpStream>, // In as Stream from the client *in* this server
+    stream_out: Option<TcpStream>, // Out as Stream from the server *to* this server
+}
+
+impl TunnelResponse {
+    pub fn new(addr: &str) -> Box<Self> {
+        return Box::new(TunnelResponse {
+            base: BaseResponse {
+                status: HttpStatus::OK,
+                headers: HashMap::new(),
+                // headers: headers! {"connection" => "keep-alive"}.unwrap(),
+            },
+            addr: addr.to_string(),
+            has_end: Arc::new(AtomicBool::new(false)),
+            stream_in: None,
+            stream_out: None,
+        });
+    }
+}
+
+impl HttpResponseCommon for TunnelResponse {
+    fn base(&mut self) -> &mut BaseResponse {
+        &mut self.base
+    }
+
+    fn next(&mut self) -> Result<Vec<u8>, IterError> {
+        self.peek()
+    }
+
+    fn peek(&mut self) -> Result<Vec<u8>, IterError> {
+        if self.has_end.load(Ordering::SeqCst) {
+            return Err(IterError::Finished);
+        }
+        let mut buf = [0; 1];
+        let r = self.stream_in.as_ref().unwrap().peek(&mut buf);
+
+        return Err(IterError::WouldBlock);
+    }
+
+    fn set_stream(&mut self, stream: &TcpStream) {
+        let mut client_stream = stream.try_clone().expect("clone failed...");
+        self.stream_in = Some(client_stream.try_clone().expect("clone failed..."));
+        let server_stream = TcpStream::connect(&self.addr);
+        if server_stream.is_err() {
+            println!("Error connecting");
+        }
+        let mut server_stream = server_stream.unwrap();
+        let _ = client_stream.set_nonblocking(false);
+        let _ = client_stream.set_nodelay(false);
+        let _ = client_stream.set_read_timeout(Some(Duration::from_secs(3)));
+        let _ = client_stream.set_write_timeout(Some(Duration::from_secs(3)));
+        let _ = client_stream.write_all(&self.base.to_bytes());
+        let mut server_stream_1 = server_stream.try_clone().expect("Error cloning");
+        let mut client_stream_1 = client_stream.try_clone().expect("clone failed...");
+        let has_ended = self.has_end.clone();
+        thread::spawn(move || {
+            let _ = io::copy(&mut client_stream_1, &mut server_stream_1);
+            has_ended.store(true, Ordering::SeqCst);
+        });
+
+        let has_ended = self.has_end.clone();
+        thread::spawn(move || {
+            let _ = io::copy(&mut server_stream, &mut client_stream);
+            has_ended.store(true, Ordering::SeqCst);
+        });
+    }
 }