From 7f3248e92d039b9d09d6121379487282e845fb42 Mon Sep 17 00:00:00 2001
From: Hannu Hartikainen <hannu.hartikainen@gmail.com>
Date: Thu, 9 Jul 2020 00:58:17 +0300
Subject: [PATCH] Implement proxy_all for proxying entire domains

I wanted to add another domain behind a reverse proxy. The existing
reverse proxy functionality seemed to only work for subpaths so I
implemented this. (I'm happy with another better way if you can point one
to me.)

I also implemented streaming responses when reverse proxying. That way
the responses stay as similar as possible to the upstream server.
---
 src/config.rs   |  2 ++
 src/conn.rs     |  7 +++++++
 src/main.rs     | 19 +++++++++++++++++++
 src/revproxy.rs | 36 ++++++++++++++++++++++++++++++++++++
 4 files changed, 64 insertions(+)

diff --git a/src/config.rs b/src/config.rs
index bbe9a28..29c70db 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -29,6 +29,8 @@ pub struct Server {
     pub usrdir: Option<bool>,
     #[cfg(feature = "proxy")]
     pub proxy: Option<HashMap<String, String>>,
+    #[cfg(feature = "proxy")]
+    pub proxy_all: Option<String>,
     pub redirect: Option<HashMap<String, String>>,
     #[cfg(feature = "scgi")]
     pub scgi: Option<HashMap<String, String>>,
diff --git a/src/conn.rs b/src/conn.rs
index 4050457..f615a3e 100644
--- a/src/conn.rs
+++ b/src/conn.rs
@@ -1,5 +1,7 @@
 use std::io;
+use std::marker::Unpin;
 use std::net::SocketAddr;
+
 use tokio::net::TcpStream;
 use tokio::prelude::*;
 use tokio_openssl::SslStream;
@@ -40,4 +42,9 @@ impl Connection {
         self.stream.flush().await?;
         Ok(())
     }
+
+    pub async fn send_stream<S: AsyncRead + Unpin>(&mut self, reader: &mut S) -> Result<(), io::Error> {
+        tokio::io::copy(reader, &mut self.stream).await?;
+        Ok(())
+    }
 }
diff --git a/src/main.rs b/src/main.rs
index 995ef13..46843bd 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -238,6 +238,25 @@ async fn handle_connection(
         None => {}
     }
 
+    #[cfg(feature = "proxy")]
+    if let Some(pr) = &srv.server.proxy_all {
+        let host_port: Vec<&str> = pr.splitn(2, ':').collect();
+        let host = host_port[0];
+        let port: Option<u16>;
+        if host_port.len() == 2 {
+            port = host_port[1].parse().ok();
+        } else {
+            port = None;
+        }
+
+        let mut upstream_url = url.clone();
+        upstream_url.set_host(Some(host)).unwrap();
+        upstream_url.set_port(port).unwrap();
+
+        revproxy::proxy_all(pr, upstream_url, con).await?;
+        return Ok(());
+    }
+
     #[cfg(feature = "proxy")]
     match &srv.server.proxy {
         Some(pr) => match url.path_segments().map(|c| c.collect::<Vec<_>>()) {
diff --git a/src/revproxy.rs b/src/revproxy.rs
index 0aebd54..b3cad3f 100644
--- a/src/revproxy.rs
+++ b/src/revproxy.rs
@@ -55,3 +55,39 @@ pub async fn proxy(addr: String, u: url::Url, mut con: conn::Connection) -> Resu
     con.send_raw(&buf).await?;
     Ok(())
 }
+
+pub async fn proxy_all(addr: &str, u: url::Url, mut con: conn::Connection) -> Result<(), io::Error> {
+    let mut connector = SslConnector::builder(SslMethod::tls()).unwrap();
+    connector.set_verify(openssl::ssl::SslVerifyMode::NONE);
+    let config = connector.build().configure().unwrap();
+
+    // TCP handshake
+    let stream = match TcpStream::connect(&addr).await {
+        Ok(s) => s,
+        Err(_) => {
+            logger::logger(con.peer_addr, Status::ProxyError, u.as_str());
+            con.send_status(Status::ProxyError, None).await?;
+            return Ok(());
+        }
+    };
+    let domain = addr.splitn(2, ':').next().unwrap();
+
+    // TLS handshake with SNI
+    let mut stream = match tokio_openssl::connect(config, domain, stream).await {
+        Ok(s) => s,
+        Err(_) => {
+            logger::logger(con.peer_addr, Status::ProxyError, u.as_str());
+            con.send_status(Status::ProxyError, None).await?;
+            return Ok(());
+        }
+    };
+
+    // send request: URL + CRLF
+    stream.write_all(u.as_ref().as_bytes()).await?;
+    stream.write_all(b"\r\n").await?;
+    stream.flush().await?;
+
+    // stream to client
+    con.send_stream(&mut stream).await?;
+    Ok(())
+}