From 36da11c02ed380b77ff7ad96900544e7981573a3 Mon Sep 17 00:00:00 2001
From: Zac Bergquist <zac.bergquist@goteleport.com>
Date: Tue, 19 Nov 2024 18:44:40 -0700
Subject: [PATCH] feat(connector): allow clients to specify the x224 nego
 request data (#580)

The previous code would (correctly) set a cookie containining the
username, but only when using username/password credentials. When
smart card credentials are used, the cookie would always contain
the empty string.
---
 crates/ironrdp-client/src/config.rs        | 1 +
 crates/ironrdp-connector/src/connection.rs | 8 +++++---
 crates/ironrdp-connector/src/lib.rs        | 4 ++++
 crates/ironrdp-web/src/session.rs          | 1 +
 crates/ironrdp/examples/screenshot.rs      | 1 +
 ffi/src/connector/config.rs                | 1 +
 6 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/crates/ironrdp-client/src/config.rs b/crates/ironrdp-client/src/config.rs
index 084ae21a..e05af8f3 100644
--- a/crates/ironrdp-client/src/config.rs
+++ b/crates/ironrdp-client/src/config.rs
@@ -318,6 +318,7 @@ impl Config {
             },
             no_server_pointer: args.no_server_pointer,
             autologon: args.autologon,
+            request_data: None,
             pointer_software_rendering: true,
             performance_flags: PerformanceFlags::default(),
         };
diff --git a/crates/ironrdp-connector/src/connection.rs b/crates/ironrdp-connector/src/connection.rs
index 317e79ce..dbf76e3e 100644
--- a/crates/ironrdp-connector/src/connection.rs
+++ b/crates/ironrdp-connector/src/connection.rs
@@ -251,9 +251,11 @@ impl Sequence for ClientConnector {
                 }
 
                 let connection_request = nego::ConnectionRequest {
-                    nego_data: Some(nego::NegoRequestData::cookie(
-                        self.config.credentials.username().to_owned(),
-                    )),
+                    nego_data: self.config.request_data.clone().or_else(|| {
+                        Some(nego::NegoRequestData::cookie(
+                            self.config.credentials.username().to_owned(),
+                        ))
+                    }),
                     flags: nego::RequestFlags::empty(),
                     protocol: security_protocol,
                 };
diff --git a/crates/ironrdp-connector/src/lib.rs b/crates/ironrdp-connector/src/lib.rs
index eb24bf6c..2e7c59e0 100644
--- a/crates/ironrdp-connector/src/lib.rs
+++ b/crates/ironrdp-connector/src/lib.rs
@@ -26,6 +26,7 @@ pub use channel_connection::{ChannelConnectionSequence, ChannelConnectionState};
 pub use connection::{encode_send_data_request, ClientConnector, ClientConnectorState, ConnectionResult};
 pub use connection_finalization::{ConnectionFinalizationSequence, ConnectionFinalizationState};
 use ironrdp_core::{encode_buf, encode_vec, Encode, WriteBuf};
+use ironrdp_pdu::nego::NegoRequestData;
 use ironrdp_pdu::rdp::capability_sets;
 use ironrdp_pdu::rdp::client_info::PerformanceFlags;
 use ironrdp_pdu::x224::X224;
@@ -162,6 +163,9 @@ pub struct Config {
     pub dig_product_id: String,
     pub client_dir: String,
     pub platform: capability_sets::MajorPlatformType,
+    /// Optional data for the x224 connection request.
+    /// Defaults to a cookie containing the username if unspecified.
+    pub request_data: Option<NegoRequestData>,
     /// If true, the INFO_AUTOLOGON flag is set in the [`ClientInfoPdu`](ironrdp_pdu::rdp::ClientInfoPdu)
     pub autologon: bool,
 
diff --git a/crates/ironrdp-web/src/session.rs b/crates/ironrdp-web/src/session.rs
index 63879ef7..f4416c1c 100644
--- a/crates/ironrdp-web/src/session.rs
+++ b/crates/ironrdp-web/src/session.rs
@@ -860,6 +860,7 @@ fn build_config(
         platform: ironrdp::pdu::rdp::capability_sets::MajorPlatformType::UNSPECIFIED,
         no_server_pointer: false,
         autologon: false,
+        request_data: None,
         pointer_software_rendering: false,
         performance_flags: PerformanceFlags::default(),
         desktop_scale_factor: 0,
diff --git a/crates/ironrdp/examples/screenshot.rs b/crates/ironrdp/examples/screenshot.rs
index b398268a..68bc59b1 100644
--- a/crates/ironrdp/examples/screenshot.rs
+++ b/crates/ironrdp/examples/screenshot.rs
@@ -223,6 +223,7 @@ fn build_config(username: String, password: String, domain: Option<String>) -> c
 
         // Disable custom pointers (there is no user interaction anyway)
         no_server_pointer: true,
+        request_data: None,
         autologon: false,
         pointer_software_rendering: true,
         performance_flags: PerformanceFlags::default(),
diff --git a/ffi/src/connector/config.rs b/ffi/src/connector/config.rs
index 317eed8d..d7bc1e7a 100644
--- a/ffi/src/connector/config.rs
+++ b/ffi/src/connector/config.rs
@@ -192,6 +192,7 @@ pub mod ffi {
 
                 no_server_pointer: self.no_server_pointer.unwrap_or(false),
                 autologon: self.autologon.unwrap_or(false),
+                request_data: None,
                 pointer_software_rendering: self.pointer_software_rendering.unwrap_or(false),
                 performance_flags: self.performance_flags.ok_or("performance flag is missing")?,
                 desktop_scale_factor: 0,