From 6da1708922b2233cd2048f8ae2728f785e565c42 Mon Sep 17 00:00:00 2001
From: Paul Butler <paul@driftingin.space>
Date: Tue, 24 Sep 2024 08:59:45 -0400
Subject: [PATCH] valid response has CORS headers

---
 plane/plane-tests/tests/proxy_cors.rs | 82 ++++++++++++++++++++++++++-
 plane/src/proxy/proxy_server.rs       |  4 +-
 2 files changed, 82 insertions(+), 4 deletions(-)

diff --git a/plane/plane-tests/tests/proxy_cors.rs b/plane/plane-tests/tests/proxy_cors.rs
index 523365ce3..d12dc35e7 100644
--- a/plane/plane-tests/tests/proxy_cors.rs
+++ b/plane/plane-tests/tests/proxy_cors.rs
@@ -2,16 +2,20 @@ use common::{
     localhost_resolver::localhost_client, proxy_mock::MockProxy,
     simple_axum_server::SimpleAxumServer, test_env::TestEnvironment,
 };
-use plane::{protocol::RouteInfoResponse, types::BearerToken};
+use plane::{
+    log_types::BackendAddr,
+    names::{BackendName, Name},
+    protocol::{RouteInfo, RouteInfoResponse},
+    types::{BearerToken, ClusterName, SecretToken},
+};
 use plane_test_macro::plane_test;
 use reqwest::StatusCode;
+use std::str::FromStr;
 
 mod common;
 
 #[plane_test]
 async fn proxy_forbidden_request_has_cors_headers(env: TestEnvironment) {
-    let _server = SimpleAxumServer::new().await;
-
     let mut proxy = MockProxy::new().await;
     let port = proxy.port();
     let url = format!("http://plane.test:{port}/abc123/");
@@ -71,3 +75,75 @@ async fn proxy_forbidden_request_has_cors_headers(env: TestEnvironment) {
         "true"
     );
 }
+
+#[plane_test]
+async fn proxy_valid_request_has_cors_headers(env: TestEnvironment) {
+    let server = SimpleAxumServer::new().await;
+    let mut proxy = MockProxy::new().await;
+    let cluster = ClusterName::from_str(&format!("plane.test:{}", proxy.port())).unwrap();
+    let port = proxy.port();
+    let url = format!("http://plane.test:{port}/abc123/");
+    let client = localhost_client();
+
+    let handle = tokio::spawn(client.get(url).send());
+
+    let route_info_request = proxy.recv_route_info_request().await;
+    assert_eq!(
+        route_info_request.token,
+        BearerToken::from("abc123".to_string())
+    );
+
+    proxy
+        .send_route_info_response(RouteInfoResponse {
+            token: BearerToken::from("abc123".to_string()),
+            route_info: Some(RouteInfo {
+                backend_id: BackendName::new_random(),
+                address: BackendAddr(server.addr()),
+                secret_token: SecretToken::from("secret".to_string()),
+                cluster,
+                user: None,
+                user_data: None,
+                subdomain: None,
+            }),
+        })
+        .await;
+
+    let response = handle.await.unwrap().unwrap();
+    assert_eq!(response.status(), StatusCode::OK);
+    assert_eq!(
+        response
+            .headers()
+            .get("access-control-allow-origin")
+            .unwrap()
+            .to_str()
+            .unwrap(),
+        "*"
+    );
+    assert_eq!(
+        response
+            .headers()
+            .get("access-control-allow-methods")
+            .unwrap()
+            .to_str()
+            .unwrap(),
+        "*"
+    );
+    assert_eq!(
+        response
+            .headers()
+            .get("access-control-allow-headers")
+            .unwrap()
+            .to_str()
+            .unwrap(),
+        "*"
+    );
+    assert_eq!(
+        response
+            .headers()
+            .get("access-control-allow-credentials")
+            .unwrap()
+            .to_str()
+            .unwrap(),
+        "true"
+    );
+}
diff --git a/plane/src/proxy/proxy_server.rs b/plane/src/proxy/proxy_server.rs
index 569804969..49376e049 100644
--- a/plane/src/proxy/proxy_server.rs
+++ b/plane/src/proxy/proxy_server.rs
@@ -91,7 +91,7 @@ impl Service<Request<Incoming>> for ProxyState {
 
             let request = request.into_request_with_simple_body();
 
-            let (res, upgrade_handler) = inner.proxy_client.request(request).await.unwrap();
+            let (mut res, upgrade_handler) = inner.proxy_client.request(request).await.unwrap();
 
             if let Some(upgrade_handler) = upgrade_handler {
                 let monitor = inner.monitor.monitor();
@@ -111,6 +111,8 @@ impl Service<Request<Incoming>> for ProxyState {
                 inner.monitor.touch_backend(&route_info.backend_id);
             }
 
+            apply_cors_headers(&mut res);
+
             Ok(res)
         })
     }