1 files changed, 26 insertions, 102 deletions

diff --git a/src/main.rs b/src/main.rs
index 2edbe3d..c325e61 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,39 +1,35 @@
+#![feature(try_trait_v2)]
+
 pub mod config;
+pub mod error;
+pub mod files;
+pub mod proxy;
 
-use crate::config::Config;
+use crate::{
+    config::{Config, HostConfig},
+    files::serve_files,
+    proxy::proxy_request,
+};
 use anyhow::{anyhow, bail, Context, Result};
+use error::ServiceError;
 use http_body_util::{combinators::BoxBody, BodyExt};
 use hyper::{
     body::Incoming,
-    header::{HOST, UPGRADE},
-    http::{
-        uri::{PathAndQuery, Scheme},
-        HeaderValue,
-    },
+    header::{CONTENT_TYPE, HOST},
+    http::HeaderValue,
     server::conn::http1,
     service::service_fn,
-    upgrade::OnUpgrade,
-    Request, Response, StatusCode, Uri,
+    Request, Response, StatusCode,
 };
-use log::{debug, error, info, warn};
+use log::{debug, info, warn};
 use std::{fs::File, io::BufReader, net::SocketAddr, path::Path, sync::Arc};
 use tokio::{
     io::{AsyncRead, AsyncWrite},
-    net::{TcpListener, TcpStream},
+    net::TcpListener,
     signal::ctrl_c,
 };
 use tokio_rustls::TlsAcceptor;
 
-#[derive(Debug, thiserror::Error)]
-enum ServiceError {
-    #[error("hyper error")]
-    Hyper(hyper::Error),
-    #[error("unknown host")]
-    NoHost,
-    #[error("can't connect to the backend")]
-    CantConnect,
-}
-
 #[tokio::main]
 async fn main() -> anyhow::Result<()> {
     env_logger::init_from_env("LOG");
@@ -60,7 +56,7 @@ async fn serve_http(config: Arc<Config>) -> Result<()> {
         let (stream, addr) = listener.accept().await.context("accepting connection")?;
         debug!("connection from {addr}");
         let config = config.clone();
-        tokio::spawn(async move { serve_stream(config, stream, addr) });
+        tokio::spawn(async move { serve_stream(config, stream, addr).await });
     }
 }
 async fn serve_https(config: Arc<Config>) -> Result<()> {
@@ -114,10 +110,11 @@ pub async fn serve_stream<T: AsyncRead + AsyncWrite + Unpin + Send + 'static>(
                         Ok(r) => Ok(r),
                         Err(ServiceError::Hyper(e)) => Err(e),
                         Err(error) => Ok({
-                            let mut resp = Response::new(format!(
-                                "the reverse proxy encountered an issue: {error}"
-                            ));
+                            let mut resp =
+                                Response::new(format!("gnix encountered an issue: {error}"));
                             *resp.status_mut() = StatusCode::BAD_REQUEST;
+                            resp.headers_mut()
+                                .insert(CONTENT_TYPE, HeaderValue::from_static("text/plain"));
                             resp
                         }
                         .map(|b| b.map_err(|e| match e {}).boxed())),
@@ -148,34 +145,10 @@ fn load_private_key(path: &Path) -> anyhow::Result<rustls::PrivateKey> {
 
 async fn service(
     config: Arc<Config>,
-    mut req: Request<Incoming>,
+    req: Request<Incoming>,
     addr: SocketAddr,
-) -> Result<hyper::Response<BoxBody<bytes::Bytes, hyper::Error>>, ServiceError> {
-    let scheme_secure = req.uri().scheme() == Some(&Scheme::HTTPS);
+) -> Result<hyper::Response<BoxBody<bytes::Bytes, ServiceError>>, ServiceError> {
     debug!("{addr} ~> {:?} {}", req.headers().get(HOST), req.uri());
-    *req.uri_mut() = Uri::builder()
-        .path_and_query(
-            req.uri()
-                .clone()
-                .path_and_query()
-                .cloned()
-                .unwrap_or(PathAndQuery::from_static("/")),
-        )
-        .build()
-        .unwrap();
-
-    req.headers_mut().insert(
-        "x-forwarded-for",
-        HeaderValue::from_str(&format!("{addr}")).unwrap(),
-    );
-    req.headers_mut().insert(
-        "x-forwarded-proto",
-        if scheme_secure {
-            HeaderValue::from_static("https")
-        } else {
-            HeaderValue::from_static("http")
-        },
-    );
 
     let route = config
         .hosts
@@ -188,59 +161,10 @@ async fn service(
         ))
         .ok_or(ServiceError::NoHost)?;
 
-    let do_upgrade = req.headers().contains_key(UPGRADE);
-    let on_upgrade_downstream = req.extensions_mut().remove::<OnUpgrade>();
-
-    debug!("\tforwarding to {}", route.backend);
-    let mut resp = {
-        let client_stream = TcpStream::connect(&route.backend)
-            .await
-            .map_err(|_| ServiceError::CantConnect)?;
-
-        let (mut sender, conn) = hyper::client::conn::http1::handshake(client_stream)
-            .await
-            .map_err(ServiceError::Hyper)?;
-        tokio::task::spawn(async move {
-            if let Err(err) = conn.await {
-                warn!("connection failed: {:?}", err);
-            }
-        });
-        sender
-            .send_request(req)
-            .await
-            .map_err(ServiceError::Hyper)?
-    };
-
-    resp.headers_mut()
-        .insert("server", HeaderValue::from_static("gnix"));
-
-    if do_upgrade {
-        let on_upgrade_upstream = resp.extensions_mut().remove::<OnUpgrade>();
-        tokio::task::spawn(async move {
-            debug!("about upgrading connection, sending empty response");
-            match (
-                on_upgrade_upstream.unwrap().await,
-                on_upgrade_downstream.unwrap().await,
-            ) {
-                (Ok(mut upgraded_upstream), Ok(mut upgraded_downstream)) => {
-                    debug!("upgrade successful");
-                    match tokio::io::copy_bidirectional(
-                        &mut upgraded_downstream,
-                        &mut upgraded_upstream,
-                    )
-                    .await
-                    {
-                        Ok((from_client, from_server)) => {
-                            debug!("proxy socket terminated: {from_server} sent, {from_client} received")
-                        }
-                        Err(e) => warn!("proxy socket error: {e}"),
-                    }
-                }
-                (a, b) => eprintln!("upgrade error: upstream={a:?} downstream={b:?}"),
-            }
-        });
+    match route {
+        HostConfig::Backend { backend } => proxy_request(req, addr, backend).await,
+        HostConfig::Files { files } => serve_files(req, files).await,
     }
-    Ok(resp.map(|b| b.boxed()))
 }
 
 pub fn remove_port(s: &str) -> &str {