From 28db3eece23075e243bbe61683dff77e60c7bf35 Mon Sep 17 00:00:00 2001
From: Zac Bergquist <zac.bergquist@goteleport.com>
Date: Mon, 16 Dec 2024 17:25:59 -0700
Subject: [PATCH] RDP: put the username in the load balancing cookie (#50226)

This doesn't change the behavior of the RDP server in any way,
but load balancers sitting in front of the server may use the
cookie to alter how they route connections.

Note that this is the default behavior of IronRDP when using
username/password authentication, but we must explicitly enable
it because we use smart card authentication instead.
---
 lib/srv/desktop/rdp/rdpclient/client.go     | 7 +++++++
 lib/srv/desktop/rdp/rdpclient/src/client.rs | 7 ++++++-
 lib/srv/desktop/rdp/rdpclient/src/lib.rs    | 3 +++
 3 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/lib/srv/desktop/rdp/rdpclient/client.go b/lib/srv/desktop/rdp/rdpclient/client.go
index ca8ad121e81f9..c7be6dd702016 100644
--- a/lib/srv/desktop/rdp/rdpclient/client.go
+++ b/lib/srv/desktop/rdp/rdpclient/client.go
@@ -291,6 +291,12 @@ func (c *Client) startRustRDP(ctx context.Context) error {
 		return trace.Wrap(err)
 	}
 
+	// [username] need only be valid for the duration of
+	// C.client_run. It is copied on the Rust side and
+	// thus can be freed here.
+	username := C.CString(c.username)
+	defer C.free(unsafe.Pointer(username))
+
 	// [addr] need only be valid for the duration of
 	// C.client_run. It is copied on the Rust side and
 	// thus can be freed here.
@@ -328,6 +334,7 @@ func (c *Client) startRustRDP(ctx context.Context) error {
 		C.CGOConnectParams{
 			ad:               C.bool(c.cfg.AD),
 			nla:              C.bool(c.cfg.NLA),
+			go_username:      username,
 			go_addr:          addr,
 			go_computer_name: computerName,
 			go_kdc_addr:      kdcAddr,
diff --git a/lib/srv/desktop/rdp/rdpclient/src/client.rs b/lib/srv/desktop/rdp/rdpclient/src/client.rs
index d4e010c8e1fa9..3dae0fc453b59 100644
--- a/lib/srv/desktop/rdp/rdpclient/src/client.rs
+++ b/lib/srv/desktop/rdp/rdpclient/src/client.rs
@@ -44,6 +44,7 @@ use ironrdp_pdu::input::fast_path::{
 };
 use ironrdp_pdu::input::mouse::PointerFlags;
 use ironrdp_pdu::input::{InputEventError, MousePdu};
+use ironrdp_pdu::nego::NegoRequestData;
 use ironrdp_pdu::rdp::capability_sets::MajorPlatformType;
 use ironrdp_pdu::rdp::client_info::PerformanceFlags;
 use ironrdp_pdu::rdp::RdpError;
@@ -1442,8 +1443,11 @@ fn create_config(params: &ConnectParams, pin: String) -> Config {
         platform: MajorPlatformType::UNSPECIFIED,
         no_server_pointer: false,
         autologon: true,
-        request_data: None,
         pointer_software_rendering: false,
+        // Send the username in the request cookie, which is sent in the initial connection request.
+        // The RDP server ignores this value, but load balancers sitting in front of the server
+        // can use it to implement persistence.
+        request_data: Some(NegoRequestData::cookie(params.username.clone())),
         performance_flags: PerformanceFlags::default()
             | PerformanceFlags::DISABLE_CURSOR_SHADOW // this is required for pointer to work correctly in Windows 2019
             | if !params.show_desktop_wallpaper {
@@ -1457,6 +1461,7 @@ fn create_config(params: &ConnectParams, pin: String) -> Config {
 
 #[derive(Debug)]
 pub struct ConnectParams {
+    pub username: String,
     pub addr: String,
     pub kdc_addr: Option<String>,
     pub computer_name: Option<String>,
diff --git a/lib/srv/desktop/rdp/rdpclient/src/lib.rs b/lib/srv/desktop/rdp/rdpclient/src/lib.rs
index 2c79722e73af0..200f1b6a6186a 100644
--- a/lib/srv/desktop/rdp/rdpclient/src/lib.rs
+++ b/lib/srv/desktop/rdp/rdpclient/src/lib.rs
@@ -92,6 +92,7 @@ pub unsafe extern "C" fn free_string(ptr: *mut c_char) {
 pub unsafe extern "C" fn client_run(cgo_handle: CgoHandle, params: CGOConnectParams) -> CGOResult {
     trace!("client_run");
     // Convert from C to Rust types.
+    let username = from_c_string(params.go_username);
     let addr = from_c_string(params.go_addr);
     let cert_der = from_go_array(params.cert_der, params.cert_der_len);
     let key_der = from_go_array(params.key_der, params.key_der_len);
@@ -111,6 +112,7 @@ pub unsafe extern "C" fn client_run(cgo_handle: CgoHandle, params: CGOConnectPar
         ConnectParams {
             ad: params.ad,
             nla: params.nla,
+            username,
             addr,
             computer_name,
             cert_der,
@@ -480,6 +482,7 @@ pub unsafe extern "C" fn client_write_screen_resize(
 pub struct CGOConnectParams {
     ad: bool,
     nla: bool,
+    go_username: *const c_char,
     go_addr: *const c_char,
     go_domain: *const c_char,
     go_kdc_addr: *const c_char,