From c06e4e4945fc9cf13f2d22e7241a9fbd2f4a4ab4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=CE=A3rebe=20-=20Romain=20GERARD?= <erebe@erebe.eu>
Date: Sat, 20 Jan 2024 11:17:31 +0100
Subject: [PATCH] Improvements http2 transport protocol

---
 src/tunnel/transport/http2.rs     |  8 +++++---
 src/tunnel/transport/websocket.rs | 15 ++++++++-------
 2 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/src/tunnel/transport/http2.rs b/src/tunnel/transport/http2.rs
index 99809fd..2b61749 100644
--- a/src/tunnel/transport/http2.rs
+++ b/src/tunnel/transport/http2.rs
@@ -119,12 +119,14 @@ pub async fn connect(
         .header(CONTENT_TYPE, "application/json")
         .version(hyper::Version::HTTP_2);
 
+    let headers = req.headers_mut().unwrap();
     for (k, v) in &client_cfg.http_headers {
-        let _ = req.headers_mut().unwrap().remove(k);
-        req = req.header(k, v);
+        let _ = headers.remove(k);
+        headers.append(k, v.clone());
     }
     if let Some(auth) = &client_cfg.http_upgrade_credentials {
-        req = req.header(AUTHORIZATION, auth);
+        let _ = headers.remove(AUTHORIZATION);
+        headers.append(AUTHORIZATION, auth.clone());
     }
 
     let (tx, rx) = mpsc::channel::<Bytes>(1024);
diff --git a/src/tunnel/transport/websocket.rs b/src/tunnel/transport/websocket.rs
index b2ed814..3201661 100644
--- a/src/tunnel/transport/websocket.rs
+++ b/src/tunnel/transport/websocket.rs
@@ -17,7 +17,7 @@ use std::io;
 use std::io::ErrorKind;
 use std::ops::DerefMut;
 use tokio::io::{AsyncWrite, AsyncWriteExt, ReadHalf, WriteHalf};
-use tracing::{error, trace};
+use tracing::trace;
 use uuid::Uuid;
 
 pub struct WebsocketTunnelWrite {
@@ -59,7 +59,6 @@ impl TunnelWrite for WebsocketTunnelWrite {
         if buf.capacity() == read_len {
             let new_size = buf.capacity() + (buf.capacity() / 4); // grow buffer by 1.25 %
             buf.reserve(new_size);
-            buf.resize(buf.capacity(), 0);
             trace!(
                 "Buffer {} Mb {} {} {}",
                 buf.capacity() as f64 / 1024.0 / 1024.0,
@@ -103,8 +102,8 @@ impl WebsocketTunnelRead {
     }
 }
 
-fn frame_reader(x: Frame<'_>) -> futures_util::future::Ready<anyhow::Result<()>> {
-    error!("frame {:?} {:?}", x.opcode, x.payload);
+fn frame_reader(_: Frame<'_>) -> futures_util::future::Ready<anyhow::Result<()>> {
+    //error!("frame {:?} {:?}", x.opcode, x.payload);
     futures_util::future::ready(anyhow::Ok(()))
 }
 
@@ -156,12 +155,14 @@ pub async fn connect(
         )
         .version(hyper::Version::HTTP_11);
 
+    let headers = req.headers_mut().unwrap();
     for (k, v) in &client_cfg.http_headers {
-        let _ = req.headers_mut().unwrap().remove(k);
-        req = req.header(k, v);
+        let _ = headers.remove(k);
+        headers.append(k, v.clone());
     }
     if let Some(auth) = &client_cfg.http_upgrade_credentials {
-        req = req.header(AUTHORIZATION, auth);
+        let _ = headers.remove(AUTHORIZATION);
+        headers.append(AUTHORIZATION, auth.clone());
     }
 
     let req = req.body(Empty::<Bytes>::new()).with_context(|| {