diff --git a/lib/srv/desktop/rdp/rdpclient/src/rdpdr/mod.rs b/lib/srv/desktop/rdp/rdpclient/src/rdpdr/mod.rs index 6dfe4aaa2c91d..673f8a0747224 100644 --- a/lib/srv/desktop/rdp/rdpclient/src/rdpdr/mod.rs +++ b/lib/srv/desktop/rdp/rdpclient/src/rdpdr/mod.rs @@ -2572,6 +2572,36 @@ impl ClientDriveQueryDirectoryResponse { io_status: NTSTATUS, buffer: Option, ) -> RdpResult { + // This match block ensures that the passed parameters are in a configuration that's + // explicitly supported by the length calculation (below) and the self.encode() method. + match io_status { + NTSTATUS::STATUS_SUCCESS => { + if buffer.is_none() { + return Err(invalid_data_error( + "a ClientDriveQueryDirectoryResponse with NTSTATUS::STATUS_SUCCESS \ + should have Some(FsInformationClass) buffer, got None", + )); + } + } + NTSTATUS::STATUS_NOT_SUPPORTED + | NTSTATUS::STATUS_NO_MORE_FILES + | NTSTATUS::STATUS_UNSUCCESSFUL => { + if buffer.is_some() { + return Err(invalid_data_error(&format!( + "a ClientDriveQueryDirectoryResponse with NTSTATUS = {:?} \ + should have a None buffer, got {:?}", + io_status, buffer, + ))); + } + } + _ => { + return Err(invalid_data_error(&format!( + "received unsupported io_status for ClientDriveQueryDirectoryResponse: {:?}", + io_status + ))) + } + } + let length = match buffer { Some(ref fs_information_class) => match fs_information_class { FsInformationClass::FileBothDirectoryInformation(fs_info_class) => { @@ -2597,40 +2627,18 @@ impl ClientDriveQueryDirectoryResponse { }) } - // TODO(isaiah): make this logic more sane fn encode(&self) -> RdpResult> { let mut w = vec![]; w.extend_from_slice(&self.device_io_reply.encode()?); - - if self.device_io_reply.io_status == NTSTATUS::to_u32(&NTSTATUS::STATUS_SUCCESS).unwrap() { - w.write_u32::(self.length)?; - w.extend_from_slice( - &self - .buffer.as_ref() - .ok_or_else(|| invalid_data_error( - "ClientDriveQueryDirectoryResponse with NTSTATUS::STATUS_SUCCESS expects a FsInformationClass" - ))? - .encode()?, - ); - } else if self.device_io_reply.io_status + w.write_u32::(self.length)?; + if let Some(buffer) = &self.buffer { + w.extend_from_slice(&buffer.encode()?); + } + if self.device_io_reply.io_status == NTSTATUS::to_u32(&NTSTATUS::STATUS_NO_MORE_FILES).unwrap() { - // https://github.com/FreeRDP/FreeRDP/blob/511444a65e7aa2f537c5e531fa68157a50c1bd4d/channels/drive/client/drive_file.c#L935-L937 - w.write_u32::(0)?; + // https://github.com/FreeRDP/FreeRDP/blob/511444a65e7aa2f537c5e531fa68157a50c1bd4d/channels/drive/client/drive_file.c#L937 w.write_u8(0)?; - } else if self.device_io_reply.io_status - == NTSTATUS::to_u32(&NTSTATUS::STATUS_NOT_SUPPORTED).unwrap() - || self.device_io_reply.io_status - == NTSTATUS::to_u32(&NTSTATUS::STATUS_UNSUCCESSFUL).unwrap() - { - // https://github.com/FreeRDP/FreeRDP/blob/511444a65e7aa2f537c5e531fa68157a50c1bd4d/channels/drive/client/drive_main.c#L665 - // https://github.com/FreeRDP/FreeRDP/blob/511444a65e7aa2f537c5e531fa68157a50c1bd4d/channels/drive/client/drive_main.c#L634 - w.write_u32::(self.length)?; - } else { - return Err(invalid_data_error(&format!( - "Found ClientDriveQueryDirectoryResponse with invalid or unhandled NTSTATUS: {:?}", - self.device_io_reply.io_status - ))); } Ok(w)