diff options
Diffstat (limited to 'src/platform/windows')
| -rw-r--r-- | src/platform/windows/clipboard_image.rs | 409 |
1 files changed, 409 insertions, 0 deletions
diff --git a/src/platform/windows/clipboard_image.rs b/src/platform/windows/clipboard_image.rs new file mode 100644 index 00000000..f8e17aed --- /dev/null +++ b/src/platform/windows/clipboard_image.rs @@ -0,0 +1,409 @@ +use std::io::{self, Cursor, Write as _}; + +use crate::protocol::MAX_CLIPBOARD_IMAGE_PAYLOAD; + +pub(super) const MAX_CLIPBOARD_ALLOCATION: usize = 64 * 1024 * 1024; + +const PNG_SIGNATURE: &[u8; 8] = b"\x89PNG\r\n\x1a\n"; +const MAX_IMAGE_DIMENSION: u32 = 16_384; +const MAX_IMAGE_PIXELS: usize = 16 * 1024 * 1024; +const BI_RGB: u32 = 0; +const BI_BITFIELDS: u32 = 3; +const BI_ALPHABITFIELDS: u32 = 6; + +pub(super) fn validated_png(bytes: &[u8]) -> Option<Vec<u8>> { + let logical_len = png_logical_len(bytes)?; + if logical_len > MAX_CLIPBOARD_IMAGE_PAYLOAD { + return None; + } + let bytes = &bytes[..logical_len]; + let mut decoder = png::Decoder::new_with_limits( + Cursor::new(bytes), + png::Limits { + bytes: MAX_CLIPBOARD_IMAGE_PAYLOAD, + }, + ); + decoder.set_ignore_text_chunk(true); + decoder.set_ignore_iccp_chunk(true); + let mut reader = decoder.read_info().ok()?; + validate_dimensions(reader.info().width, reader.info().height)?; + if reader.info().animation_control.is_some() { + return None; + } + while reader.next_interlaced_row().ok()?.is_some() {} + reader.finish().ok()?; + Some(bytes.to_vec()) +} + +fn png_logical_len(bytes: &[u8]) -> Option<usize> { + if !bytes.starts_with(PNG_SIGNATURE) { + return None; + } + let mut offset = PNG_SIGNATURE.len(); + while offset < bytes.len() { + let length = usize::try_from(read_u32_be(bytes, offset)?).ok()?; + let kind_start = offset.checked_add(4)?; + let data_start = kind_start.checked_add(4)?; + let data_end = data_start.checked_add(length)?; + let chunk_end = data_end.checked_add(4)?; + if chunk_end > bytes.len() { + return None; + } + if bytes.get(kind_start..data_start)? == b"IEND" { + return (length == 0).then_some(chunk_end); + } + offset = chunk_end; + } + None +} + +pub(super) fn dib_to_png(bytes: &[u8]) -> Option<Vec<u8>> { + Dib::parse(bytes)?.encode_png() +} + +struct Dib<'a> { + width: u32, + height: u32, + top_down: bool, + stride: usize, + pixels: &'a [u8], + format: PixelFormat, +} + +#[derive(Clone, Copy)] +enum PixelFormat { + Bgr24, + Bgrx32, + Masked32 { + red: ChannelMask, + green: ChannelMask, + blue: ChannelMask, + alpha: Option<ChannelMask>, + }, +} + +#[derive(Clone, Copy)] +struct ChannelMask { + mask: u32, + shift: u32, + maximum: u32, +} + +impl ChannelMask { + fn parse(mask: u32) -> Option<Self> { + if mask == 0 { + return None; + } + let shift = mask.trailing_zeros(); + let maximum = mask >> shift; + if maximum & maximum.wrapping_add(1) != 0 { + return None; + } + Some(Self { + mask, + shift, + maximum, + }) + } + + fn extract(self, pixel: u32) -> u8 { + let value = (pixel & self.mask) >> self.shift; + ((u64::from(value) * 255 + u64::from(self.maximum) / 2) / u64::from(self.maximum)) as u8 + } +} + +impl<'a> Dib<'a> { + fn parse(bytes: &'a [u8]) -> Option<Self> { + let header_size = usize::try_from(read_u32_le(bytes, 0)?).ok()?; + if !matches!(header_size, 40 | 52 | 56 | 108 | 124) || bytes.len() < header_size { + return None; + } + let signed_width = read_i32_le(bytes, 4)?; + let signed_height = read_i32_le(bytes, 8)?; + if signed_width <= 0 || signed_height == 0 || signed_height == i32::MIN { + return None; + } + let width = u32::try_from(signed_width).ok()?; + let height = signed_height.unsigned_abs(); + validate_dimensions(width, height)?; + if read_u16_le(bytes, 12)? != 1 { + return None; + } + + let bit_count = read_u16_le(bytes, 14)?; + let compression = read_u32_le(bytes, 16)?; + let colors_used = usize::try_from(read_u32_le(bytes, 32)?).ok()?; + let (format, external_mask_bytes) = + pixel_format(bytes, header_size, bit_count, compression)?; + let palette_bytes = colors_used.checked_mul(4)?; + let pixel_offset = header_size + .checked_add(external_mask_bytes)? + .checked_add(palette_bytes)?; + let row_bits = usize::try_from(width) + .ok()? + .checked_mul(usize::from(bit_count))?; + let stride = row_bits.checked_add(31)?.checked_div(32)?.checked_mul(4)?; + let image_size = stride.checked_mul(usize::try_from(height).ok()?)?; + let pixels = bytes.get(pixel_offset..pixel_offset.checked_add(image_size)?)?; + + Some(Self { + width, + height, + top_down: signed_height < 0, + stride, + pixels, + format, + }) + } + + fn encode_png(&self) -> Option<Vec<u8>> { + let mut output = LimitedWriter::new(MAX_CLIPBOARD_IMAGE_PAYLOAD); + { + let mut encoder = png::Encoder::new(&mut output, self.width, self.height); + encoder.set_color(png::ColorType::Rgba); + encoder.set_depth(png::BitDepth::Eight); + let mut writer = encoder.write_header().ok()?; + let mut stream = writer.stream_writer().ok()?; + let mut rgba = vec![0_u8; usize::try_from(self.width).ok()?.checked_mul(4)?]; + for output_row in 0..usize::try_from(self.height).ok()? { + let source_row = if self.top_down { + output_row + } else { + usize::try_from(self.height) + .ok()? + .checked_sub(output_row + 1)? + }; + let start = source_row.checked_mul(self.stride)?; + let source = self.pixels.get(start..start.checked_add(self.stride)?)?; + self.decode_row(source, &mut rgba)?; + stream.write_all(&rgba).ok()?; + } + stream.finish().ok()?; + writer.finish().ok()?; + } + Some(output.into_inner()) + } + + fn decode_row(&self, source: &[u8], output: &mut [u8]) -> Option<()> { + match self.format { + PixelFormat::Bgr24 => { + for (pixel, rgba) in source + .chunks_exact(3) + .take(usize::try_from(self.width).ok()?) + .zip(output.chunks_exact_mut(4)) + { + rgba.copy_from_slice(&[pixel[2], pixel[1], pixel[0], 255]); + } + } + PixelFormat::Bgrx32 => { + for (pixel, rgba) in source + .chunks_exact(4) + .take(usize::try_from(self.width).ok()?) + .zip(output.chunks_exact_mut(4)) + { + rgba.copy_from_slice(&[pixel[2], pixel[1], pixel[0], 255]); + } + } + PixelFormat::Masked32 { + red, + green, + blue, + alpha, + } => { + for (pixel, rgba) in source + .chunks_exact(4) + .take(usize::try_from(self.width).ok()?) + .zip(output.chunks_exact_mut(4)) + { + let pixel = u32::from_le_bytes(pixel.try_into().ok()?); + rgba.copy_from_slice(&[ + red.extract(pixel), + green.extract(pixel), + blue.extract(pixel), + alpha.map_or(255, |mask| mask.extract(pixel)), + ]); + } + } + } + Some(()) + } +} + +fn pixel_format( + bytes: &[u8], + header_size: usize, + bit_count: u16, + compression: u32, +) -> Option<(PixelFormat, usize)> { + match (bit_count, compression) { + (24, BI_RGB) => return Some((PixelFormat::Bgr24, 0)), + (32, BI_RGB) => return Some((PixelFormat::Bgrx32, 0)), + (32, BI_BITFIELDS | BI_ALPHABITFIELDS) => {} + _ => return None, + } + + let (offset, external_mask_bytes) = if header_size == 40 { + let count = if compression == BI_ALPHABITFIELDS { + 4 + } else { + 3 + }; + (40, count * 4) + } else { + if header_size < 52 || (compression == BI_ALPHABITFIELDS && header_size < 56) { + return None; + } + (40, 0) + }; + let red = ChannelMask::parse(read_u32_le(bytes, offset)?)?; + let green = ChannelMask::parse(read_u32_le(bytes, offset + 4)?)?; + let blue = ChannelMask::parse(read_u32_le(bytes, offset + 8)?)?; + if red.mask & green.mask != 0 || red.mask & blue.mask != 0 || green.mask & blue.mask != 0 { + return None; + } + let alpha = if compression == BI_ALPHABITFIELDS || header_size >= 56 { + match read_u32_le(bytes, offset + 12).filter(|mask| *mask != 0) { + Some(mask) => Some(ChannelMask::parse(mask)?), + None => None, + } + } else { + None + }; + if alpha.is_some_and(|alpha| alpha.mask & (red.mask | green.mask | blue.mask) != 0) { + return None; + } + Some(( + PixelFormat::Masked32 { + red, + green, + blue, + alpha, + }, + external_mask_bytes, + )) +} + +fn validate_dimensions(width: u32, height: u32) -> Option<()> { + if width == 0 || height == 0 || width > MAX_IMAGE_DIMENSION || height > MAX_IMAGE_DIMENSION { + return None; + } + let pixels = usize::try_from(width) + .ok()? + .checked_mul(usize::try_from(height).ok()?)?; + (pixels <= MAX_IMAGE_PIXELS).then_some(()) +} + +struct LimitedWriter { + bytes: Vec<u8>, + limit: usize, +} + +impl LimitedWriter { + fn new(limit: usize) -> Self { + Self { + bytes: Vec::new(), + limit, + } + } + + fn into_inner(self) -> Vec<u8> { + self.bytes + } +} + +impl io::Write for LimitedWriter { + fn write(&mut self, bytes: &[u8]) -> io::Result<usize> { + if bytes.len() > self.limit.saturating_sub(self.bytes.len()) { + return Err(io::Error::new( + io::ErrorKind::FileTooLarge, + "encoded clipboard image exceeds protocol limit", + )); + } + self.bytes.extend_from_slice(bytes); + Ok(bytes.len()) + } + + fn flush(&mut self) -> io::Result<()> { + Ok(()) + } +} + +fn read_u16_le(bytes: &[u8], offset: usize) -> Option<u16> { + Some(u16::from_le_bytes( + bytes.get(offset..offset + 2)?.try_into().ok()?, + )) +} + +fn read_u32_le(bytes: &[u8], offset: usize) -> Option<u32> { + Some(u32::from_le_bytes( + bytes.get(offset..offset + 4)?.try_into().ok()?, + )) +} + +fn read_i32_le(bytes: &[u8], offset: usize) -> Option<i32> { + Some(i32::from_le_bytes( + bytes.get(offset..offset + 4)?.try_into().ok()?, + )) +} + +fn read_u32_be(bytes: &[u8], offset: usize) -> Option<u32> { + Some(u32::from_be_bytes( + bytes.get(offset..offset + 4)?.try_into().ok()?, + )) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn info_header(width: i32, height: i32, bit_count: u16, compression: u32) -> Vec<u8> { + let mut header = vec![0_u8; 40]; + header[0..4].copy_from_slice(&40_u32.to_le_bytes()); + header[4..8].copy_from_slice(&width.to_le_bytes()); + header[8..12].copy_from_slice(&height.to_le_bytes()); + header[12..14].copy_from_slice(&1_u16.to_le_bytes()); + header[14..16].copy_from_slice(&bit_count.to_le_bytes()); + header[16..20].copy_from_slice(&compression.to_le_bytes()); + header + } + + fn decode_rgba(bytes: &[u8]) -> Vec<u8> { + let mut decoder = png::Decoder::new(Cursor::new(bytes)); + decoder.set_transformations(png::Transformations::EXPAND | png::Transformations::STRIP_16); + let mut reader = decoder.read_info().unwrap(); + let mut output = vec![0_u8; reader.output_buffer_size()]; + let info = reader.next_frame(&mut output).unwrap(); + output.truncate(info.buffer_size()); + output + } + + #[test] + fn converts_bottom_up_24_bit_dib_to_png() { + let mut dib = info_header(2, 2, 24, BI_RGB); + dib.extend_from_slice(&[ + 255, 0, 0, 255, 255, 255, 0, 0, // bottom: blue, white, padding + 0, 0, 255, 0, 255, 0, 0, 0, // top: red, green, padding + ]); + + let png = dib_to_png(&dib).unwrap(); + assert_eq!( + decode_rgba(&png), + vec![255, 0, 0, 255, 0, 255, 0, 255, 0, 0, 255, 255, 255, 255, 255, 255,] + ); + } + + #[test] + fn strips_registered_png_allocator_tail() { + let mut bytes = Vec::new(); + let mut encoder = png::Encoder::new(&mut bytes, 1, 1); + encoder.set_color(png::ColorType::Rgba); + encoder.set_depth(png::BitDepth::Eight); + let mut writer = encoder.write_header().unwrap(); + writer.write_image_data(&[1, 2, 3, 255]).unwrap(); + writer.finish().unwrap(); + let expected = bytes.clone(); + bytes.extend_from_slice(&[0; 32]); + + assert_eq!(validated_png(&bytes), Some(expected)); + } +} |