websocket proxy works

1 files changed, 15 insertions, 14 deletions

diff --git a/src/main.rs b/src/main.rs
index 8b4d4d5..976cad5 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -2,11 +2,15 @@ pub mod config;
 
 use crate::config::Config;
 use anyhow::{bail, Context, Result};
-use bytes::Bytes;
-use http_body_util::{combinators::BoxBody, BodyExt, Empty};
+use http_body_util::{combinators::BoxBody, BodyExt};
 use hyper::{
-    body::Incoming, header::UPGRADE, http::uri::PathAndQuery, server::conn::http1,
-    service::service_fn, upgrade::OnUpgrade, Method, Request, Response, StatusCode, Uri,
+    body::Incoming,
+    header::UPGRADE,
+    http::{uri::PathAndQuery, HeaderValue},
+    server::conn::http1,
+    service::service_fn,
+    upgrade::OnUpgrade,
+    Request, Response, Uri,
 };
 use log::{debug, error, info, warn};
 use std::{fs::File, io::BufReader, path::Path, sync::Arc};
@@ -87,7 +91,7 @@ pub fn serve_stream<T: AsyncRead + AsyncWrite + Unpin + Send + 'static>(
                             Ok(r) => Ok(r),
                             Err(ServiceError::Hyper(e)) => Err(e),
                             Err(other) => Ok(Response::new(format!(
-                                "the reverse proxy encountered an error: {other:?}"
+                                "the reverse proxy encountered an error: gnix-{other:?}"
                             ))
                             .map(|b| b.map_err(|e| match e {}).boxed())),
                         }
@@ -151,7 +155,7 @@ async fn service(
         ))
         .ok_or(ServiceError::NoHost)?;
 
-    let upgrade_header = req.headers().get(UPGRADE).cloned();
+    let do_upgrade = req.headers().contains_key(UPGRADE);
     let on_upgrade_downstream = req.extensions_mut().remove::<OnUpgrade>();
 
     let mut resp = {
@@ -173,7 +177,10 @@ async fn service(
             .map_err(ServiceError::Hyper)?
     };
 
-    if let Some(proto) = upgrade_header {
+    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");
@@ -198,14 +205,8 @@ async fn service(
                 (a, b) => eprintln!("upgrade error: upstream={a:?} downstream={b:?}"),
             }
         });
-
-        let mut resp = Response::new(Empty::<Bytes>::new());
-        *resp.status_mut() = StatusCode::SWITCHING_PROTOCOLS;
-        resp.headers_mut().insert(UPGRADE, proto);
-        Ok(resp.map(|b| b.map_err(|e| match e {}).boxed()))
-    } else {
-        Ok(resp.map(|b| b.boxed()))
     }
+    Ok(resp.map(|b| b.boxed()))
 }
 
 pub fn remove_port(s: &str) -> &str {