Refactor proxy

LICENSE

@@ -19,3 +19,27 @@ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 SOFTWARE.
+
+---
+
+Contains code from hyperium/hyper-util, licensed under the MIT license:
+
+Copyright (c) 2023 Sean McArthur
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.

dynamic-proxy/src/graceful_shutdown.rs

@@ -12,7 +12,7 @@ use std::{
 };
 use tokio::sync::watch;
 
-#[derive(Clone)]
+#[derive(Clone)] // Added in Plane
 pub struct GracefulShutdown {
     tx: watch::Sender<()>,
 }
@@ -34,6 +34,7 @@ impl GracefulShutdown {
         })
     }
 
+    // Added in Plane
     pub fn subscribe(&self) -> watch::Receiver<()> {
         self.tx.subscribe()
     }

dynamic-proxy/tests/common/websocket_echo_server.rs

@@ -48,14 +48,17 @@ async fn ws_handler(ws: WebSocketUpgrade) -> impl IntoResponse {
 
 async fn handle_socket(mut socket: WebSocket) {
     while let Some(msg) = socket.recv().await {
-        if let Ok(msg) = msg {
-            if let Message::Text(text) = msg {
-                if socket.send(Message::Text(text)).await.is_err() {
-                    break;
+        match msg {
+            Ok(msg) => {
+                if let Message::Text(text) = msg {
+                    if socket.send(Message::Text(text)).await.is_err() {
+                        break;
+                    }
                 }
             }
-        } else {
-            break;
+            Err(e) => {
+                panic!("Error receiving message: {}", e);
+            }
         }
     }
 }

dynamic-proxy/tests/graceful.rs

@@ -36,9 +36,11 @@ async fn test_graceful_shutdown() {
     // Create a client and start a POST request without finishing the body
     let client = reqwest::Client::new();
 
-    let _client = client.clone();
-    let _url = url.clone();
-    let response_handle = tokio::spawn(async move { _client.get(&_url).send().await.unwrap() });
+    let response_handle = {
+        let client = client.clone();
+        let url = url.clone();
+        tokio::spawn(async move { client.get(&url).send().await.unwrap() })
+    };
 
     tokio::time::sleep(Duration::from_millis(100)).await;
 

dynamic-proxy/tests/graceful_https.rs

@@ -46,9 +46,11 @@ async fn test_graceful_shutdown_https() {
         .build()
         .unwrap();
 
-    let _client = client.clone();
-    let _url = url.clone();
-    let response_handle = tokio::spawn(async move { _client.get(&_url).send().await.unwrap() });
+    let response_handle = {
+        let client = client.clone();
+        let url = url.clone();
+        tokio::spawn(async move { client.get(&url).send().await.unwrap() })
+    };
 
     tokio::time::sleep(Duration::from_millis(100)).await;
 

dynamic-proxy/tests/test_http_redirect.rs

@@ -28,9 +28,8 @@ async fn do_request(url: &str) -> Response {
     url.set_port(Some(port)).unwrap();
 
     // Request to http://foo.bar.baz should redirect to https://foo.bar.baz
-    let response = get_client().get(url).send().await.unwrap();
 
-    response
+    get_client().get(url).send().await.unwrap()
 }
 
 #[tokio::test]

plane/plane-tests/tests/common/mod.rs

@@ -10,10 +10,10 @@ pub mod docker;
 pub mod localhost_resolver;
 pub mod proxy_mock;
 pub mod resources;
-pub mod simple_axum_server; // NOTE: copied from dynamic-proxy (until we merge them back)
+pub mod simple_axum_server; // TODO: copied from dynamic-proxy (until we merge them back)
 pub mod test_env;
 pub mod timeout;
-pub mod websocket_echo_server; // NOTE: copied from dynamic-proxy (until we merge them back)
+pub mod websocket_echo_server; // TODO: copied from dynamic-proxy (until we merge them back)
 
 pub fn run_test<F, Fut>(name: &str, time_limit: Duration, test_function: F)
 where

plane/plane-tests/tests/common/websocket_echo_server.rs

@@ -48,14 +48,17 @@ async fn ws_handler(ws: WebSocketUpgrade) -> impl IntoResponse {
 
 async fn handle_socket(mut socket: WebSocket) {
     while let Some(msg) = socket.recv().await {
-        if let Ok(msg) = msg {
-            if let Message::Text(text) = msg {
-                if socket.send(Message::Text(text)).await.is_err() {
-                    break;
+        match msg {
+            Ok(msg) => {
+                if let Message::Text(text) = msg {
+                    if socket.send(Message::Text(text)).await.is_err() {
+                        break;
+                    }
                 }
             }
-        } else {
-            break;
+            Err(e) => {
+                panic!("Error receiving message: {}", e);
+            }
         }
     }
 }

plane/plane-tests/tests/proxy_connection_monitor.rs

@@ -60,7 +60,7 @@ async fn proxy_marks_backend_as_recently_active(env: TestEnvironment) {
         panic!("Backend entry not found");
     };
     assert_eq!(backend_entry.active_connections, 0);
-    assert_eq!(backend_entry.had_recent_connection, true);
+    assert!(backend_entry.had_recent_connection);
 }
 
 #[plane_test]

plane/src/proxy/proxy_server.rs

@@ -144,7 +144,7 @@ impl Service<Request<Incoming>> for ProxyState {
             let (mut res, upgrade_handler) = match result {
                 Ok((res, upgrade_handler)) => (res, upgrade_handler),
                 Err(e) => {
-                    tracing::error!("Error proxying request: {}", e);
+                    tracing::error!(?e, "Error proxying request");
                     return status_code_to_response(StatusCode::INTERNAL_SERVER_ERROR);
                 }
             };

plane/src/proxy/request.rs

@@ -47,7 +47,7 @@ pub fn get_and_maybe_remove_bearer_token(parts: &mut uri::Parts) -> Option<Beare
     // the full incoming path, and the path to proxy to is just `/`.
     let (token, path) = match full_path.split_once('/') {
         Some((token, path)) => (token, path),
-        None => (full_path, "/"),
+        None => (full_path, ""),
     };
 
     if token.is_empty() {
@@ -76,6 +76,8 @@ pub fn get_and_maybe_remove_bearer_token(parts: &mut uri::Parts) -> Option<Beare
 
 #[cfg(test)]
 mod tests {
+    use uri::Uri;
+
     use super::*;
     use std::str::FromStr;
 
@@ -134,4 +136,71 @@ mod tests {
         let cluster = ClusterName::from_str("myhost").unwrap();
         assert_eq!(subdomain_from_host(host, &cluster), Ok(Some("foobar")));
     }
+
+    #[test]
+    fn test_get_and_maybe_remove_bearer_token() {
+        let url = Uri::from_str("https://example.com/foo/bar").unwrap();
+        let mut parts = url.into_parts();
+        assert_eq!(
+            get_and_maybe_remove_bearer_token(&mut parts),
+            Some(BearerToken::from("foo".to_string()))
+        );
+        assert_eq!(
+            parts.path_and_query,
+            Some(PathAndQuery::from_str("/bar").unwrap())
+        );
+    }
+
+    #[test]
+    fn test_get_and_maybe_remove_bearer_token_ends_no_slash() {
+        let url = Uri::from_str("https://example.com/foo").unwrap();
+        let mut parts = url.into_parts();
+        assert_eq!(
+            get_and_maybe_remove_bearer_token(&mut parts),
+            Some(BearerToken::from("foo".to_string()))
+        );
+        assert_eq!(
+            parts.path_and_query,
+            Some(PathAndQuery::from_str("/").unwrap())
+        );
+    }
+
+    #[test]
+    fn test_get_and_maybe_remove_bearer_token_ends_in_slash() {
+        let url = Uri::from_str("https://example.com/foo/").unwrap();
+        let mut parts = url.into_parts();
+        assert_eq!(
+            get_and_maybe_remove_bearer_token(&mut parts),
+            Some(BearerToken::from("foo".to_string()))
+        );
+        assert_eq!(
+            parts.path_and_query,
+            Some(PathAndQuery::from_str("/").unwrap())
+        );
+    }
+
+    #[test]
+    fn test_get_and_maybe_remove_bearer_token_no_token() {
+        let url = Uri::from_str("https://example.com/").unwrap();
+        let mut parts = url.into_parts();
+        assert_eq!(get_and_maybe_remove_bearer_token(&mut parts), None);
+        assert_eq!(
+            parts.path_and_query,
+            Some(PathAndQuery::from_str("/").unwrap())
+        );
+    }
+
+    #[test]
+    fn test_get_and_maybe_remove_bearer_token_static_token() {
+        let url = Uri::from_str("https://example.com/s.foo/bar").unwrap();
+        let mut parts = url.into_parts();
+        assert_eq!(
+            get_and_maybe_remove_bearer_token(&mut parts),
+            Some(BearerToken::from("s.foo".to_string()))
+        );
+        assert_eq!(
+            parts.path_and_query,
+            Some(PathAndQuery::from_str("/s.foo/bar").unwrap())
+        );
+    }
 }