refactor (pool 2/N): Pipeline owns its PSInvocationState

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<PipelineCommand>,
     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 {

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<Vec<u8>>,
+}
+
+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<u8>,
+    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 <SS> payload is the ciphertext bytes only (base64 encoded).
+    let iv = [0u8; 16];
+
+    let encryptor = cbc::Encryptor::<Aes256>::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::<Pkcs7>(&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(())
+}

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;

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<Vec<u8>>,
-}
-
 #[derive(Debug)]
 pub struct RunspacePool {
     pub(super) id: uuid::Uuid,
@@ -140,114 +131,11 @@ pub struct RunspacePool {
     pub(super) pipelines: HashMap<uuid::Uuid, Pipeline>,
     pub(super) fragmenter: fragmentation::Fragmenter,
     pub(super) desired_stream_is_pooling: bool,
-    pub(super) key_exchange: Option<KeyExchangeState>,
+    pub(super) key_exchange: Option<super::crypto::KeyExchangeState>,
     pub(super) psrp_key_exchange_pending: bool,
     pub(super) pending_host_calls: VecDeque<HostCall>,
 }
 
-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<u8>,
-    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 <SS> payload is the ciphertext bytes only (base64 encoded).
-    let iv = [0u8; 16];
-
-    let encryptor = cbc::Encryptor::<Aes256>::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::<Pkcs7>(&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
@@ -723,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)],
@@ -1555,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(())
     }
@@ -1571,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
@@ -1615,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
@@ -1816,13 +1701,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,
             });
@@ -1899,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",
             ));