diff options
Diffstat (limited to 'src/client/endpoint_commands.rs')
| -rw-r--r-- | src/client/endpoint_commands.rs | 224 |
1 files changed, 224 insertions, 0 deletions
diff --git a/src/client/endpoint_commands.rs b/src/client/endpoint_commands.rs new file mode 100644 index 00000000..decb8522 --- /dev/null +++ b/src/client/endpoint_commands.rs @@ -0,0 +1,224 @@ +use std::collections::VecDeque; +use std::io; + +use crate::api::client::ApiClientError; +use crate::api::schema::{Request, ResponseResult}; +use crate::ipc::LocalStream; +use crate::protocol::ClientMessage; + +use super::shell::ClientShellEndpointError; + +struct InFlightCommand { + boot_id: String, + request_id: String, + response: Vec<u8>, +} + +pub(super) struct EndpointCommandResult { + pub(super) boot_id: String, + pub(super) request_id: String, + pub(super) result: Result<ResponseResult, ClientShellEndpointError>, +} + +#[derive(Default)] +pub(super) struct EndpointCommands { + queued: VecDeque<(String, Box<Request>)>, + in_flight: Option<InFlightCommand>, +} + +impl EndpointCommands { + pub(super) fn enqueue(&mut self, boot_id: String, request: Box<Request>) { + self.queued.push_back((boot_id, request)); + } + + pub(super) fn send_next(&mut self, stream: &mut LocalStream) -> io::Result<()> { + if self.in_flight.is_some() { + return Ok(()); + } + let Some((boot_id, request)) = self.queued.pop_front() else { + return Ok(()); + }; + let request_id = request.id.clone(); + let request = serde_json::to_string(&request) + .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?; + super::write_to_server( + stream, + &ClientMessage::ClientShellEndpointRequest { + boot_id: boot_id.clone(), + request, + }, + )?; + self.in_flight = Some(InFlightCommand { + boot_id, + request_id, + response: Vec::new(), + }); + Ok(()) + } + + pub(super) fn receive_chunk( + &mut self, + response_boot_id: &str, + response_request_id: &str, + final_chunk: bool, + data: Vec<u8>, + ) -> io::Result<Option<EndpointCommandResult>> { + let Some(in_flight) = self.in_flight.as_mut() else { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "endpoint response arrived without an in-flight command", + )); + }; + if response_boot_id != in_flight.boot_id || response_request_id != in_flight.request_id { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "endpoint response correlation did not match the in-flight command", + )); + } + in_flight.response.extend(data); + if !final_chunk { + return Ok(None); + } + + let in_flight = self.in_flight.take().expect("checked in-flight command"); + let response = String::from_utf8(in_flight.response) + .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?; + let result = parse_response(&in_flight.request_id, &response); + Ok(Some(EndpointCommandResult { + boot_id: in_flight.boot_id, + request_id: in_flight.request_id, + result, + })) + } +} + +fn parse_response( + expected_id: &str, + response: &str, +) -> Result<ResponseResult, ClientShellEndpointError> { + let value = serde_json::from_str(response).map_err(|error| ClientShellEndpointError { + code: None, + message: format!("invalid endpoint response: {error}"), + })?; + match crate::api::client::parse_response_value(value) { + Ok(response) if response.id == expected_id => Ok(response.result), + Ok(response) => Err(ClientShellEndpointError { + code: None, + message: format!( + "endpoint response id {:?} did not match {expected_id:?}", + response.id + ), + }), + Err(ApiClientError::ErrorResponse(response)) if response.id == expected_id => { + Err(ClientShellEndpointError { + code: Some(response.error.code), + message: response.error.message, + }) + } + Err(ApiClientError::ErrorResponse(response)) => Err(ClientShellEndpointError { + code: None, + message: format!( + "endpoint error id {:?} did not match {expected_id:?}", + response.id + ), + }), + Err(error) => Err(ClientShellEndpointError { + code: None, + message: error.to_string(), + }), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::api::schema::{ResponseResult, SuccessResponse}; + + fn commands_with_in_flight() -> EndpointCommands { + EndpointCommands { + in_flight: Some(InFlightCommand { + boot_id: "boot-a".into(), + request_id: "request-a".into(), + response: Vec::new(), + }), + ..EndpointCommands::default() + } + } + + #[test] + fn chunked_response_completion_is_correlated_and_clears_the_lane() { + let mut commands = commands_with_in_flight(); + let response = serde_json::to_string(&SuccessResponse { + id: "request-a".into(), + result: ResponseResult::Ok {}, + }) + .unwrap(); + let split = response.len() / 2; + + assert!(commands + .receive_chunk( + "boot-a", + "request-a", + false, + response.as_bytes()[..split].to_vec(), + ) + .unwrap() + .is_none()); + let completed = commands + .receive_chunk( + "boot-a", + "request-a", + true, + response.as_bytes()[split..].to_vec(), + ) + .unwrap() + .unwrap(); + + assert_eq!(completed.boot_id, "boot-a"); + assert_eq!(completed.request_id, "request-a"); + assert!(matches!(completed.result, Ok(ResponseResult::Ok {}))); + assert!(commands.in_flight.is_none()); + } + + #[test] + fn large_selection_response_reassembles_without_truncation() { + let mut commands = commands_with_in_flight(); + let selection = "selected".repeat(160_000); + let response = serde_json::to_vec(&SuccessResponse { + id: "request-a".into(), + result: ResponseResult::PaneSelection { + pane_id: "w1:p1".into(), + text: selection.clone(), + }, + }) + .unwrap(); + let chunk_count = response.len().div_ceil(128 * 1024); + let mut completed = None; + for (index, chunk) in response.chunks(128 * 1024).enumerate() { + completed = commands + .receive_chunk( + "boot-a", + "request-a", + index + 1 == chunk_count, + chunk.to_vec(), + ) + .unwrap(); + } + + assert!(matches!( + completed.expect("final selection response").result, + Ok(ResponseResult::PaneSelection { text, .. }) if text == selection + )); + } + + #[test] + fn response_from_another_boot_is_rejected() { + let mut commands = commands_with_in_flight(); + + let Err(error) = commands.receive_chunk("boot-b", "request-a", true, b"{}".to_vec()) else { + panic!("mismatched boot should fail"); + }; + + assert_eq!(error.kind(), io::ErrorKind::InvalidData); + } +} |