From 68d11956bc4e77029edb42b89ad0297f3b958fa0 Mon Sep 17 00:00:00 2001
From: Alex Grinman <alex@krypt.co>
Date: Fri, 24 Jul 2020 20:44:19 -0400
Subject: [PATCH] fix query param forwarding

---
 .gitignore                     |  3 ++-
 Cargo.lock                     |  2 +-
 Cargo.toml                     |  6 +++++-
 tunnelto/Cargo.toml            |  2 +-
 tunnelto/src/introspect/mod.rs | 38 ++++++++++++++++++++++++++++++++--
 tunnelto/templates/detail.html |  2 +-
 tunnelto/templates/index.html  |  2 +-
 7 files changed, 47 insertions(+), 8 deletions(-)

diff --git a/.gitignore b/.gitignore
index 1d9b32b..fdcc697 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,4 +1,5 @@
 /target
 tunnelto_server
 !tunnelto_server/
-.env
\ No newline at end of file
+.env
+tunnelto_proxy/
\ No newline at end of file
diff --git a/Cargo.lock b/Cargo.lock
index 24b44ad..4a04fe9 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2194,7 +2194,7 @@ dependencies = [
 
 [[package]]
 name = "tunnelto"
-version = "0.1.11"
+version = "0.1.12"
 dependencies = [
  "askama",
  "bytes",
diff --git a/Cargo.toml b/Cargo.toml
index a379ff9..6821024 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -3,5 +3,9 @@
 members = [
     "tunnelto_lib",
     "tunnelto",
-    "tunnelto_server"
+    "tunnelto_server",
+]
+
+exclude = [
+    "tunnelto_proxy"
 ]
\ No newline at end of file
diff --git a/tunnelto/Cargo.toml b/tunnelto/Cargo.toml
index a1003aa..99e60b2 100644
--- a/tunnelto/Cargo.toml
+++ b/tunnelto/Cargo.toml
@@ -1,7 +1,7 @@
 [package]
 name = "tunnelto"
 description = "expose your local web server to the internet with a public url"
-version = "0.1.11"
+version = "0.1.12"
 authors = ["Alex Grinman <alex@tunnelto.dev>"]
 edition = "2018"
 license = "MIT"
diff --git a/tunnelto/src/introspect/mod.rs b/tunnelto/src/introspect/mod.rs
index 657345b..da10cc6 100644
--- a/tunnelto/src/introspect/mod.rs
+++ b/tunnelto/src/introspect/mod.rs
@@ -10,6 +10,7 @@ use futures::{Stream, StreamExt};
 use bytes::Buf;
 use uuid::Uuid;
 use http_body::Body;
+use std::convert::Infallible;
 
 type HttpClient = hyper::Client<hyper::client::HttpConnector>;
 
@@ -19,6 +20,7 @@ pub struct Request {
     status: u16,
     is_replay: bool,
     path: String,
+    query: Option<String>,
     method: Method,
     headers: HashMap<String, Vec<String>>,
     body_data: Vec<u8>,
@@ -28,6 +30,16 @@ pub struct Request {
     completed: chrono::NaiveDateTime,
 }
 
+impl Request {
+    pub fn path_and_query(&self) -> String {
+        if let Some(query) = self.query.as_ref() {
+            format!("{}?{}", self.path, query)
+        } else {
+            self.path.clone()
+        }
+    }
+}
+
 impl Request {
     pub fn elapsed(&self) -> String {
         let duration = self.completed - self.started;
@@ -72,6 +84,7 @@ pub fn start_introspection_server(config: Config) -> IntrospectionAddrs {
         .and(warp::any().map(move || local_addr.clone()))
         .and(warp::method())
         .and(warp::path::full())
+        .and(opt_raw_query())
         .and(warp::header::headers_cloned())
         .and(warp::body::stream())
         .and(get_client())
@@ -124,6 +137,7 @@ pub fn start_introspection_server(config: Config) -> IntrospectionAddrs {
 async fn forward(local_addr: String,
                  method: Method,
                  path: FullPath,
+                 query: Option<String>,
                  headers: HeaderMap,
                  mut body: impl Stream<Item = Result<impl Buf, warp::Error>> + Send + Sync + Unpin + 'static,
                  client: HttpClient) -> Result<Box<dyn warp::Reply>, warp::reject::Rejection>
@@ -147,7 +161,14 @@ async fn forward(local_addr: String,
         collected.extend_from_slice(chunk.as_ref())
     }
 
-    let url = format!("http://{}{}", local_addr, path.as_str());
+    let query_str = if let Some(query) = query.as_ref() {
+        format!("?{}", query)
+    } else {
+        String::new()
+    };
+
+    let url = format!("http://{}{}{}", local_addr, path.as_str(), query_str);
+    log::debug!("forwarding to: {}", &url);
 
     let mut request = hyper::Request::builder()
         .method(method.clone())
@@ -196,6 +217,7 @@ async fn forward(local_addr: String,
         id: Uuid::new_v4().to_string(),
         status: parts.status.as_u16(),
         path: path.as_str().to_owned(),
+        query,
         method,
         headers: request_headers,
         body_data: collected,
@@ -294,7 +316,13 @@ async fn replay_request(rid: String, client: HttpClient, addr: SocketAddr) -> Re
         None => return Err(warp::reject::not_found())
     };
 
-    let url = format!("http://localhost:{}{}", addr.port(), &request.path);
+    let query_str = if let Some(query) = request.query.as_ref() {
+        format!("?{}", query)
+    } else {
+        String::new()
+    };
+
+    let url = format!("http://localhost:{}{}{}", addr.port(), &request.path, query_str);
 
     let mut new_request = hyper::Request::builder()
         .method(request.method)
@@ -339,4 +367,10 @@ impl <T> warp::reply::Reply for Page<T> where T:askama::Template + Send + 'stati
             "text/html",
         ).body(res.into()).unwrap()
     }
+}
+
+fn opt_raw_query() -> impl Filter<Extract = (Option<String>,), Error = Infallible> + Copy {
+    warp::filters::query::raw().map(|q| Some(q))
+        .or(warp::any().map(|| None))
+        .unify()
 }
\ No newline at end of file
diff --git a/tunnelto/templates/detail.html b/tunnelto/templates/detail.html
index 315c4cf..71fe47f 100644
--- a/tunnelto/templates/detail.html
+++ b/tunnelto/templates/detail.html
@@ -47,7 +47,7 @@
                     <span class="has-text-weight-bold">{{request.method}}</span>
                 </td>
                 <td>
-                    <span class="is-family-code">{{request.path}}</span>
+                    <span class="is-family-code">{{request.path_and_query()}}</span>
                 </td>
                 <td class="is-narrow">
                     <span class="">{{request.body_data.len()/1024}} KB</span>
diff --git a/tunnelto/templates/index.html b/tunnelto/templates/index.html
index 1de811a..603ec86 100644
--- a/tunnelto/templates/index.html
+++ b/tunnelto/templates/index.html
@@ -50,7 +50,7 @@
                     <span class="has-text-weight-bold">{{r.method}}</span>
                 </td>
                 <td>
-                    <span class="is-family-code">{{r.path}}</span>
+                    <span class="is-family-code">{{r.path_and_query()}}</span>
                 </td>
                 <td class="is-narrow">
                     <span class="">{{r.body_data.len()/1024}} KB</span>