From ce54a1fe1c1a46962a1a563c345b02713b0d71a7 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 23 Jun 2026 13:34:31 -0400 Subject: [PATCH 1/3] refactor(runspace_pool): extract PSRP session crypto into crypto.rs Pure, behavior-preserving move of the SecureString in-place encryption walk, the AES-256-CBC helper, and KeyExchangeState out of the ~1900-line pool.rs into a dedicated runspace_pool/crypto.rs module (private, items pub(super)). First step of decomposing the RunspacePool god-object; no logic or wire-format change. --- .../src/runspace_pool/crypto.rs | 120 +++++++++++++++++ .../src/runspace_pool/mod.rs | 1 + .../src/runspace_pool/pool.rs | 122 +----------------- 3 files changed, 127 insertions(+), 116 deletions(-) create mode 100644 crates/ironposh-client-core/src/runspace_pool/crypto.rs diff --git a/crates/ironposh-client-core/src/runspace_pool/crypto.rs b/crates/ironposh-client-core/src/runspace_pool/crypto.rs new file mode 100644 index 0000000..07ab9b8 --- /dev/null +++ b/crates/ironposh-client-core/src/runspace_pool/crypto.rs @@ -0,0 +1,120 @@ +//! PSRP session crypto: SecureString encryption and the key-exchange state. +//! +//! MS-PSRP encrypts `SecureString` values with a session key negotiated via +//! `PUBLIC_KEY` / `ENCRYPTED_SESSION_KEY` (AES-256-CBC, zero IV). This module +//! holds that state and the in-place encryption walk; the runspace pool drives +//! the key exchange and calls in here when serializing values. + +use aes::Aes256; +use cipher::block_padding::Pkcs7; +use cipher::{BlockModeEncrypt, KeyIvInit}; +use tracing::debug; + +#[derive(Debug)] +pub(super) struct KeyExchangeState { + pub(super) private_key: rsa::RsaPrivateKey, + pub(super) session_key: Option>, +} + +pub(super) fn encrypt_secure_strings_in_value_rec( + value: &mut ironposh_psrp::PsValue, + session_key: Option<&[u8]>, +) -> Result<(), crate::PwshCoreError> { + use ironposh_psrp::{ComplexObjectContent, Container, PsPrimitiveValue, PsValue}; + + match value { + PsValue::Primitive(PsPrimitiveValue::SecureString(bytes)) => { + let Some(session_key) = session_key else { + return Err(crate::PwshCoreError::InvalidResponse( + "SecureString encountered but PSRP session key is not established".into(), + )); + }; + encrypt_secure_string_bytes_in_place(bytes, session_key)?; + } + PsValue::Primitive(_) => {} + PsValue::Object(obj) => { + for value in obj.properties.values_mut() { + encrypt_secure_strings_in_value_rec(value, session_key)?; + } + + match &mut obj.content { + ComplexObjectContent::ExtendedPrimitive(p) => { + if let PsPrimitiveValue::SecureString(bytes) = p { + let Some(session_key) = session_key else { + return Err(crate::PwshCoreError::InvalidResponse( + "SecureString encountered but PSRP session key is not established" + .into(), + )); + }; + encrypt_secure_string_bytes_in_place(bytes, session_key)?; + } + } + ComplexObjectContent::Container( + Container::Stack(items) | Container::Queue(items) | Container::List(items), + ) => { + for item in items.iter_mut() { + encrypt_secure_strings_in_value_rec(item, session_key)?; + } + } + ComplexObjectContent::Container(Container::Dictionary(dict)) => { + for (_k, v) in dict.iter_mut() { + encrypt_secure_strings_in_value_rec(v, session_key)?; + } + } + ComplexObjectContent::Standard | ComplexObjectContent::PsEnums(_) => {} + } + } + } + + Ok(()) +} + +fn encrypt_secure_string_bytes_in_place( + bytes: &mut Vec, + session_key: &[u8], +) -> Result<(), crate::PwshCoreError> { + if session_key.len() != 32 { + return Err(crate::PwshCoreError::InvalidResponse( + format!( + "PSRP SecureString encryption requires 32-byte session key; got {}", + session_key.len() + ) + .into(), + )); + } + + // PowerShell's PSRP SecureString encryption uses AES-256-CBC with a zero IV. + // The payload is the ciphertext bytes only (base64 encoded). + let iv = [0u8; 16]; + + let encryptor = cbc::Encryptor::::new_from_slices(session_key, &iv).map_err(|e| { + crate::PwshCoreError::InvalidResponse( + format!("Failed to initialize AES encryptor: {e}").into(), + ) + })?; + + // MS-PSRP SecureString payload is UTF-16LE plaintext encrypted with AES-256-CBC. + let msg_len = bytes.len(); + let pad = 16 - (msg_len % 16); + let mut buf = bytes.clone(); + buf.resize(msg_len + pad, 0); + let ciphertext = encryptor + .encrypt_padded::(&mut buf, msg_len) + .map_err(|e| { + crate::PwshCoreError::InvalidResponse( + format!("Failed to encrypt SecureString (padding): {e}").into(), + ) + })?; + + let out = ciphertext.to_vec(); + + debug!( + session_key_len = session_key.len(), + plaintext_len = msg_len, + encrypted_len = out.len(), + "encrypted SecureString payload" + ); + + *bytes = out; + Ok(()) +} diff --git a/crates/ironposh-client-core/src/runspace_pool/mod.rs b/crates/ironposh-client-core/src/runspace_pool/mod.rs index 445f409..50b9058 100644 --- a/crates/ironposh-client-core/src/runspace_pool/mod.rs +++ b/crates/ironposh-client-core/src/runspace_pool/mod.rs @@ -1,4 +1,5 @@ pub mod creator; +mod crypto; pub mod enums; pub mod expect_shell_connected; pub mod expect_shell_created; diff --git a/crates/ironposh-client-core/src/runspace_pool/pool.rs b/crates/ironposh-client-core/src/runspace_pool/pool.rs index a8bd4c1..80af60b 100644 --- a/crates/ironposh-client-core/src/runspace_pool/pool.rs +++ b/crates/ironposh-client-core/src/runspace_pool/pool.rs @@ -18,9 +18,6 @@ use rsa::traits::PublicKeyParts; use rsa::{RsaPrivateKey, pkcs1v15::Pkcs1v15Encrypt}; use tracing::{debug, error, info, instrument, trace, warn}; -use aes::Aes256; -use cipher::block_padding::Pkcs7; -use cipher::{BlockModeEncrypt, KeyIvInit}; use uuid::Uuid; use crate::{ @@ -116,12 +113,6 @@ pub enum AcceptResponsResult { }, } -#[derive(Debug)] -pub(super) struct KeyExchangeState { - private_key: rsa::RsaPrivateKey, - session_key: Option>, -} - #[derive(Debug)] pub struct RunspacePool { pub(super) id: uuid::Uuid, @@ -140,114 +131,11 @@ pub struct RunspacePool { pub(super) pipelines: HashMap, pub(super) fragmenter: fragmentation::Fragmenter, pub(super) desired_stream_is_pooling: bool, - pub(super) key_exchange: Option, + pub(super) key_exchange: Option, pub(super) psrp_key_exchange_pending: bool, pub(super) pending_host_calls: VecDeque, } -fn encrypt_secure_strings_in_value_rec( - value: &mut ironposh_psrp::PsValue, - session_key: Option<&[u8]>, -) -> Result<(), crate::PwshCoreError> { - use ironposh_psrp::{ComplexObjectContent, Container, PsPrimitiveValue, PsValue}; - - match value { - PsValue::Primitive(PsPrimitiveValue::SecureString(bytes)) => { - let Some(session_key) = session_key else { - return Err(crate::PwshCoreError::InvalidResponse( - "SecureString encountered but PSRP session key is not established".into(), - )); - }; - encrypt_secure_string_bytes_in_place(bytes, session_key)?; - } - PsValue::Primitive(_) => {} - PsValue::Object(obj) => { - for value in obj.properties.values_mut() { - encrypt_secure_strings_in_value_rec(value, session_key)?; - } - - match &mut obj.content { - ComplexObjectContent::ExtendedPrimitive(p) => { - if let PsPrimitiveValue::SecureString(bytes) = p { - let Some(session_key) = session_key else { - return Err(crate::PwshCoreError::InvalidResponse( - "SecureString encountered but PSRP session key is not established" - .into(), - )); - }; - encrypt_secure_string_bytes_in_place(bytes, session_key)?; - } - } - ComplexObjectContent::Container( - Container::Stack(items) | Container::Queue(items) | Container::List(items), - ) => { - for item in items.iter_mut() { - encrypt_secure_strings_in_value_rec(item, session_key)?; - } - } - ComplexObjectContent::Container(Container::Dictionary(dict)) => { - for (_k, v) in dict.iter_mut() { - encrypt_secure_strings_in_value_rec(v, session_key)?; - } - } - ComplexObjectContent::Standard | ComplexObjectContent::PsEnums(_) => {} - } - } - } - - Ok(()) -} - -fn encrypt_secure_string_bytes_in_place( - bytes: &mut Vec, - session_key: &[u8], -) -> Result<(), crate::PwshCoreError> { - if session_key.len() != 32 { - return Err(crate::PwshCoreError::InvalidResponse( - format!( - "PSRP SecureString encryption requires 32-byte session key; got {}", - session_key.len() - ) - .into(), - )); - } - - // PowerShell's PSRP SecureString encryption uses AES-256-CBC with a zero IV. - // The payload is the ciphertext bytes only (base64 encoded). - let iv = [0u8; 16]; - - let encryptor = cbc::Encryptor::::new_from_slices(session_key, &iv).map_err(|e| { - crate::PwshCoreError::InvalidResponse( - format!("Failed to initialize AES encryptor: {e}").into(), - ) - })?; - - // MS-PSRP SecureString payload is UTF-16LE plaintext encrypted with AES-256-CBC. - let msg_len = bytes.len(); - let pad = 16 - (msg_len % 16); - let mut buf = bytes.clone(); - buf.resize(msg_len + pad, 0); - let ciphertext = encryptor - .encrypt_padded::(&mut buf, msg_len) - .map_err(|e| { - crate::PwshCoreError::InvalidResponse( - format!("Failed to encrypt SecureString (padding): {e}").into(), - ) - })?; - - let out = ciphertext.to_vec(); - - debug!( - session_key_len = session_key.len(), - plaintext_len = msg_len, - encrypted_len = out.len(), - "encrypted SecureString payload" - ); - - *bytes = out; - Ok(()) -} - impl RunspacePool { pub fn encrypt_secure_strings_in_value( &self, @@ -257,7 +145,7 @@ impl RunspacePool { .key_exchange .as_ref() .and_then(|s| s.session_key.as_deref()); - encrypt_secure_strings_in_value_rec(value, session_key) + super::crypto::encrypt_secure_strings_in_value_rec(value, session_key) } /// Build the negotiation payload shared by [`Self::open`] and @@ -1816,13 +1704,15 @@ impl RunspacePool { Ok(xml) } - fn ensure_key_exchange_state(&mut self) -> Result<&mut KeyExchangeState, PwshCoreError> { + fn ensure_key_exchange_state( + &mut self, + ) -> Result<&mut super::crypto::KeyExchangeState, PwshCoreError> { if self.key_exchange.is_none() { let mut rng = rand::thread_rng(); let private_key = RsaPrivateKey::new(&mut rng, 2048).map_err(|e| { PwshCoreError::InternalError(format!("failed to generate RSA keypair: {e}")) })?; - self.key_exchange = Some(KeyExchangeState { + self.key_exchange = Some(super::crypto::KeyExchangeState { private_key, session_key: None, }); From 25ea785ba515f9088f4b306dc0d37814436bcee4 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 23 Jun 2026 13:41:40 -0400 Subject: [PATCH 2/3] refactor(pipeline): Pipeline owns its PSInvocationState Make the `state` field private and route all transitions through `set_state()` / reads through `state()` / terminal checks through `is_terminal()`, instead of the RunspacePool poking the field directly. Behavior-preserving; no wire/logic change. --- crates/ironposh-client-core/src/pipeline.rs | 21 ++++++++++++++++++- .../src/runspace_pool/pool.rs | 15 ++++++------- 2 files changed, 26 insertions(+), 10 deletions(-) diff --git a/crates/ironposh-client-core/src/pipeline.rs b/crates/ironposh-client-core/src/pipeline.rs index 4cb4999..ab352e5 100644 --- a/crates/ironposh-client-core/src/pipeline.rs +++ b/crates/ironposh-client-core/src/pipeline.rs @@ -73,7 +73,7 @@ pub struct ExecutionResult { /// This is owned and managed by the `RunspacePool`. #[derive(Debug, Clone)] pub struct Pipeline { - pub(crate) state: PsInvocationState, + state: PsInvocationState, pub(crate) commands: Vec, pub(crate) results: ExecutionResult, } @@ -98,6 +98,25 @@ impl Pipeline { pub(crate) fn add_command(&mut self, command: PipelineCommand) { self.commands.push(command); } + + /// Returns the current invocation state of the pipeline. + pub(crate) fn state(&self) -> PsInvocationState { + self.state + } + + /// Sets the invocation state of the pipeline. + pub(crate) fn set_state(&mut self, state: PsInvocationState) { + self.state = state; + } + + /// Returns `true` when the pipeline has reached a terminal state + /// (`Completed`, `Failed`, or `Stopped`). + pub(crate) fn is_terminal(&self) -> bool { + matches!( + self.state, + PsInvocationState::Completed | PsInvocationState::Failed | PsInvocationState::Stopped + ) + } } impl Pipeline { diff --git a/crates/ironposh-client-core/src/runspace_pool/pool.rs b/crates/ironposh-client-core/src/runspace_pool/pool.rs index 80af60b..fa94e75 100644 --- a/crates/ironposh-client-core/src/runspace_pool/pool.rs +++ b/crates/ironposh-client-core/src/runspace_pool/pool.rs @@ -611,7 +611,7 @@ impl RunspacePool { "Pipeline not found for command response".into(), ) })? - .state = PsInvocationState::Running; + .set_state(PsInvocationState::Running); result.push(AcceptResponsResult::ReceiveResponse { desired_streams: vec![DesiredStream::stdout_for_command(pipeline_id)], @@ -1443,7 +1443,7 @@ impl RunspacePool { PwshCoreError::InvalidResponse("Pipeline not found for command_id".into()) })?; // Update the pipeline state - pipeline.state = PsInvocationState::from(pipeline_state.pipeline_state); + pipeline.set_state(PsInvocationState::from(pipeline_state.pipeline_state)); Ok(()) } @@ -1459,7 +1459,7 @@ impl RunspacePool { .ok_or(PwshCoreError::InvalidState("Pipeline handle not found"))?; // Set pipeline state to Running - pipeline.state = PsInvocationState::Running; + pipeline.set_state(PsInvocationState::Running); info!(pipeline_id = %handle.id(), "Invoking pipeline"); // Convert business pipeline to protocol pipeline and build CreatePipeline message @@ -1503,17 +1503,14 @@ impl RunspacePool { error!(pipeline_id = ?&handle.id(), "Pipeline handle not found "); })?; - if pipeline.state == PsInvocationState::Stopped - || pipeline.state == PsInvocationState::Completed - || pipeline.state == PsInvocationState::Failed - { + if pipeline.is_terminal() { return Err(PwshCoreError::InvalidState( "Cannot kill a pipeline that is already stopped, completed, or failed", )); } // Set pipeline state to Stopping - pipeline.state = PsInvocationState::Stopping; + pipeline.set_state(PsInvocationState::Stopping); info!(pipeline_id = %handle.id(), "Killing pipeline"); let request = self @@ -1789,7 +1786,7 @@ impl RunspacePool { .get_mut(&powershell.id()) .ok_or(PwshCoreError::InvalidState("Pipeline handle not found"))?; - if pipeline.state != PsInvocationState::NotStarted { + if pipeline.state() != PsInvocationState::NotStarted { return Err(PwshCoreError::InvalidState( "Cannot add to a pipeline that has already been started", )); From ad6494dc25f3fb1789148391ef0048bfc8ce1ff6 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 23 Jun 2026 13:47:05 -0400 Subject: [PATCH 3/3] refactor(runspace_pool): extract host-call construction/classification into host_call.rs Move the pure PipelineHostCall parsing and the needs-session-key classification out of pool.rs into a new private runspace_pool::host_call module (named to avoid collision with crate::host). Fragmentation-coupled send + the pending_host_calls queue stay in the pool. Behavior-preserving. --- .../src/runspace_pool/host_call.rs | 75 +++++++++++++++++++ .../src/runspace_pool/mod.rs | 1 + .../src/runspace_pool/pool.rs | 47 +----------- 3 files changed, 79 insertions(+), 44 deletions(-) create mode 100644 crates/ironposh-client-core/src/runspace_pool/host_call.rs diff --git a/crates/ironposh-client-core/src/runspace_pool/host_call.rs b/crates/ironposh-client-core/src/runspace_pool/host_call.rs new file mode 100644 index 0000000..44d139a --- /dev/null +++ b/crates/ironposh-client-core/src/runspace_pool/host_call.rs @@ -0,0 +1,75 @@ +//! Pure host-call construction and classification helpers extracted from +//! `runspace_pool::pool`. +//! +//! These functions are behavior-preserving and free of any `RunspacePool` +//! state: they only parse a `PsValue` into a [`HostCall`] and classify an +//! existing [`HostCall`]. Anything that needs pool internals (fragmenter, id, +//! shell, the `pending_host_calls` queue) stays in `pool.rs`. + +use tracing::debug; +use uuid::Uuid; + +use crate::{ + PwshCoreError, + host::{HostCall, HostCallScope}, +}; +use ironposh_psrp::PsValue; + +/// Parse a `PipelineHostCall` `PsValue` into a [`HostCall`] scoped to a pipeline. +/// +/// This is the pure parsing body previously inlined in +/// `RunspacePool::handle_pipeline_host_call`; it does not touch pool state. +pub(super) fn pipeline_host_call_from( + ps_value: PsValue, + stream_name: &str, + command_id: Option<&Uuid>, +) -> Result { + let PsValue::Object(pipeline_host_call) = ps_value else { + return Err(PwshCoreError::InvalidResponse( + "Expected PipelineHostCall as PsValue::Object".into(), + )); + }; + + let pipeline_host_call = ironposh_psrp::PipelineHostCall::try_from(pipeline_host_call)?; + + debug!( + ?pipeline_host_call, + stream_name = stream_name, + command_id = ?command_id, + method = ?pipeline_host_call.method, + parameters = ?pipeline_host_call.parameters, + "Received PipelineHostCall" + ); + + // Question: Can we have a Optional command id here? + let Some(command_id) = command_id else { + return Err(PwshCoreError::InvalidResponse( + "Expected command_id to be Some".into(), + )); + }; + + let scope = HostCallScope::Pipeline { + command_id: command_id.to_owned(), + }; + + HostCall::try_from_pipeline(scope, pipeline_host_call).map_err(|e| { + PwshCoreError::InvalidResponse(format!("Failed to parse host call: {e}").into()) + }) +} + +/// Classify whether a host call requires a PSRP session key to be established +/// before it can be answered (because it transports secure-string data). +pub(super) fn needs_session_key(host_call: &HostCall) -> bool { + match host_call { + HostCall::ReadLineAsSecureString { .. } + | HostCall::PromptForCredential1 { .. } + | HostCall::PromptForCredential2 { .. } => true, + HostCall::Prompt { transport } => { + let (_, _, fields) = &transport.params; + fields + .iter() + .any(|f| f.parameter_type.contains("SecureString")) + } + _ => false, + } +} diff --git a/crates/ironposh-client-core/src/runspace_pool/mod.rs b/crates/ironposh-client-core/src/runspace_pool/mod.rs index 50b9058..abf5543 100644 --- a/crates/ironposh-client-core/src/runspace_pool/mod.rs +++ b/crates/ironposh-client-core/src/runspace_pool/mod.rs @@ -3,6 +3,7 @@ mod crypto; pub mod enums; pub mod expect_shell_connected; pub mod expect_shell_created; +mod host_call; pub mod pool; pub mod types; diff --git a/crates/ironposh-client-core/src/runspace_pool/pool.rs b/crates/ironposh-client-core/src/runspace_pool/pool.rs index fa94e75..821a6b3 100644 --- a/crates/ironposh-client-core/src/runspace_pool/pool.rs +++ b/crates/ironposh-client-core/src/runspace_pool/pool.rs @@ -22,7 +22,7 @@ use uuid::Uuid; use crate::{ PwshCoreError, - host::{HostCall, HostCallScope}, + host::HostCall, pipeline::{Pipeline, PipelineCommand, PipelineSpec}, powershell::PipelineHandle, runspace::win_rs::WinRunspace, @@ -1085,18 +1085,7 @@ impl RunspacePool { })?; debug!(target: "host_call", host_call = ?host_call, "successfully created host call"); - let needs_session_key = match &host_call { - HostCall::ReadLineAsSecureString { .. } - | HostCall::PromptForCredential1 { .. } - | HostCall::PromptForCredential2 { .. } => true, - HostCall::Prompt { transport } => { - let (_, _, fields) = &transport.params; - fields - .iter() - .any(|f| f.parameter_type.contains("SecureString")) - } - _ => false, - }; + let needs_session_key = super::host_call::needs_session_key(&host_call); let has_session_key = self .key_exchange @@ -1527,37 +1516,7 @@ impl RunspacePool { stream_name: &str, command_id: Option<&Uuid>, ) -> Result { - let PsValue::Object(pipeline_host_call) = ps_value else { - return Err(PwshCoreError::InvalidResponse( - "Expected PipelineHostCall as PsValue::Object".into(), - )); - }; - - let pipeline_host_call = ironposh_psrp::PipelineHostCall::try_from(pipeline_host_call)?; - - debug!( - ?pipeline_host_call, - stream_name = stream_name, - command_id = ?command_id, - method = ?pipeline_host_call.method, - parameters = ?pipeline_host_call.parameters, - "Received PipelineHostCall" - ); - - // Question: Can we have a Optional command id here? - let Some(command_id) = command_id else { - return Err(PwshCoreError::InvalidResponse( - "Expected command_id to be Some".into(), - )); - }; - - let scope = HostCallScope::Pipeline { - command_id: command_id.to_owned(), - }; - - HostCall::try_from_pipeline(scope, pipeline_host_call).map_err(|e| { - crate::PwshCoreError::InvalidResponse(format!("Failed to parse host call: {e}").into()) - }) + super::host_call::pipeline_host_call_from(ps_value, stream_name, command_id) } /// Send a pipeline host response to the server