From 739ecc6f6d39d942ec339a46907194bf7083fb98 Mon Sep 17 00:00:00 2001
From: Conrad Ludgate <conrad.ludgate@databricks.com>
Date: Mon, 30 Jun 2025 08:02:11 +0100
Subject: [PATCH] send compute IP address to regional ingress

---
 proxy/src/console_redirect_proxy.rs |  2 +-
 proxy/src/proxy/mod.rs              | 10 +++++++++-
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/proxy/src/console_redirect_proxy.rs b/proxy/src/console_redirect_proxy.rs
index fe799258a2..3dbc1709fd 100644
--- a/proxy/src/console_redirect_proxy.rs
+++ b/proxy/src/console_redirect_proxy.rs
@@ -230,7 +230,7 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin + Send>(
         .authenticate(ctx, &mut node)
         .or_else(|e| async { Err(stream.throw_error(e, Some(ctx)).await) })
         .await?;
-    send_client_greeting(ctx, &config.greetings, &mut stream);
+    send_client_greeting(ctx, &config.greetings, &mut stream, node.socket_addr);
 
     // let session = cancellation_handler.get_key();
 
diff --git a/proxy/src/proxy/mod.rs b/proxy/src/proxy/mod.rs
index 2d96b76479..0f59cec472 100644
--- a/proxy/src/proxy/mod.rs
+++ b/proxy/src/proxy/mod.rs
@@ -8,6 +8,7 @@ pub(crate) mod wake_compute;
 
 use std::collections::HashSet;
 use std::convert::Infallible;
+use std::net::SocketAddr;
 use std::sync::Arc;
 
 use futures::TryStreamExt;
@@ -112,7 +113,7 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin + Send>(
         Err(e) => Err(client.throw_error(e, Some(ctx)).await)?,
     };
 
-    send_client_greeting(ctx, &config.greetings, client);
+    send_client_greeting(ctx, &config.greetings, client, node.socket_addr);
 
     let auth::Backend::ControlPlane(_, _user_info) = backend else {
         unreachable!("ensured above");
@@ -160,6 +161,7 @@ pub(crate) fn send_client_greeting(
     ctx: &RequestContext,
     greetings: &String,
     client: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin>,
+    socket_addr: SocketAddr,
 ) {
     // Expose session_id to clients if we have a greeting message.
     if !greetings.is_empty() {
@@ -167,6 +169,12 @@ pub(crate) fn send_client_greeting(
         client.write_message(BeMessage::NoticeResponse(session_msg.as_str()));
     }
 
+    // needed for RI to know what IP to send cancellation to.
+    client.write_message(BeMessage::ParameterStatus {
+        name: "upstream_ip".as_bytes(),
+        value: socket_addr.ip().to_string().as_bytes(),
+    });
+
     // Forward recorded latencies for probing requests
     if let Some(testodrome_id) = ctx.get_testodrome_id() {
         client.write_message(BeMessage::ParameterStatus {