refactor (pool 3/N): extract host-call construction/classification into host_call.rs

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/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<HostCall, PwshCoreError> {
+    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,
+    }
+}

crates/ironposh-client-core/src/runspace_pool/mod.rs

@@ -1,7 +1,9 @@
 pub mod creator;
+mod crypto;
 pub mod enums;
 pub mod expect_shell_connected;
 pub mod expect_shell_created;
+mod host_call;
 pub mod pool;
 pub mod types;