x-plane-backend-id header

jamsocket · Sep 24, 2024 · 034e312 · 034e312
1 parent f0c82a9
commit 034e312
Show file tree

Hide file tree

Showing 2 changed files with 48 additions and 1 deletion.
diff --git a/plane/plane-tests/tests/proxy.rs b/plane/plane-tests/tests/proxy.rs
@@ -280,3 +280,45 @@ async fn proxy_backend_passes_forwarded_headers(env: TestEnvironment) {
     assert_eq!(headers.get("x-forwarded-for").unwrap(), "127.0.0.1");
     assert_eq!(headers.get("x-forwarded-proto").unwrap(), "http");
 }
+
+#[plane_test]
+async fn proxy_returns_backend_id_in_header(env: TestEnvironment) {
+    let server = SimpleAxumServer::new().await;
+
+    let mut proxy = MockProxy::new().await;
+    let port = proxy.port();
+    let cluster = ClusterName::from_str(&format!("plane.test:{}", port)).unwrap();
+    let url = format!("http://plane.test:{port}/abc123/");
+    let client = localhost_client();
+    let handle = tokio::spawn(client.get(url).send());
+
+    let backend_id = BackendName::new_random();
+    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: backend_id.clone(),
+                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);
+    let headers = response.headers();
+    assert_eq!(
+        headers.get("x-plane-backend-id").unwrap().to_str().unwrap(),
+        &backend_id.to_string()
+    );
+}
diff --git a/plane/src/proxy/proxy_server.rs b/plane/src/proxy/proxy_server.rs
@@ -99,19 +99,24 @@ impl Service<Request<Incoming>> for ProxyState {
                     .lock()
                     .expect("Monitor lock poisoned")
                     .inc_connection(&route_info.backend_id);
+                let backend_id = route_info.backend_id.clone();
                 tokio::spawn(async move {
                     upgrade_handler.run().await.unwrap();
 
                     monitor
                         .lock()
                         .expect("Monitor lock poisoned")
-                        .dec_connection(&route_info.backend_id);
+                        .dec_connection(&backend_id);
                 });
             } else {
                 inner.monitor.touch_backend(&route_info.backend_id);
             }
 
             apply_cors_headers(&mut res);
+            res.headers_mut().insert(
+                "x-plane-backend-id",
+                HeaderValue::from_str(&route_info.backend_id.to_string()).unwrap(),
+            );
 
             Ok(res)
         })