From 7a9a459a402814611d8b47a72050efa2e90eaa38 Mon Sep 17 00:00:00 2001
From: Hexorg <Hexorg@users.noreply.github.com>
Date: Sun, 14 Apr 2024 10:40:10 -0400
Subject: [PATCH] Fixed crash when transcoding one- or two-channel KTX2
 textures (#12629)

# Objective

Fixes a crash when transcoding one- or two-channel KTX2 textures

## Solution

transcoded array has been pre-allocated up to levels.len using a macros.
Rgb8 transcoding already uses that and addresses transcoded array by an
index. R8UnormSrgb and Rg8UnormSrgb were pushing on top of the
transcoded vec, resulting in first levels.len() vectors to stay empty,
and second levels.len() levels actually being transcoded, which then
resulted in out of bounds read when copying levels to gpu
---
 crates/bevy_render/src/texture/ktx2.rs | 61 +++++++++++++++++++-------
 1 file changed, 45 insertions(+), 16 deletions(-)

diff --git a/crates/bevy_render/src/texture/ktx2.rs b/crates/bevy_render/src/texture/ktx2.rs
index 8318a6fad7..9e37c47175 100644
--- a/crates/bevy_render/src/texture/ktx2.rs
+++ b/crates/bevy_render/src/texture/ktx2.rs
@@ -90,14 +90,12 @@ pub fn ktx2_buffer_to_image(
                 TranscodeFormat::R8UnormSrgb => {
                     let (mut original_width, mut original_height) = (width, height);
 
-                    for level_data in &levels {
-                        transcoded.push(
-                            level_data
-                                .iter()
-                                .copied()
-                                .map(|v| (Srgba::gamma_function(v as f32 / 255.) * 255.).floor() as u8)
-                                .collect::<Vec<u8>>(),
-                        );
+                    for (level, level_data) in levels.iter().enumerate() {
+                        transcoded[level] = level_data
+                            .iter()
+                            .copied()
+                            .map(|v| (Srgba::gamma_function(v as f32 / 255.) * 255.).floor() as u8)
+                            .collect::<Vec<u8>>();
 
                         // Next mip dimensions are half the current, minimum 1x1
                         original_width = (original_width / 2).max(1);
@@ -109,14 +107,12 @@ pub fn ktx2_buffer_to_image(
                 TranscodeFormat::Rg8UnormSrgb => {
                     let (mut original_width, mut original_height) = (width, height);
 
-                    for level_data in &levels {
-                        transcoded.push(
-                            level_data
-                                .iter()
-                                .copied()
-                                .map(|v| (Srgba::gamma_function(v as f32 / 255.) * 255.).floor() as u8)
-                                .collect::<Vec<u8>>(),
-                        );
+                    for (level, level_data) in levels.iter().enumerate() {
+                        transcoded[level] = level_data
+                            .iter()
+                            .copied()
+                            .map(|v| (Srgba::gamma_function(v as f32 / 255.) * 255.).floor() as u8)
+                            .collect::<Vec<u8>>();
 
                         // Next mip dimensions are half the current, minimum 1x1
                         original_width = (original_width / 2).max(1);
@@ -1493,3 +1489,36 @@ pub fn ktx2_format_to_texture_format(
         }
     })
 }
+
+#[cfg(test)]
+mod tests {
+    use crate::texture::CompressedImageFormats;
+
+    use super::ktx2_buffer_to_image;
+
+    #[test]
+    fn test_ktx_levels() {
+        // R8UnormSrgb textture with 4x4 pixels data and 3 levels of mipmaps
+        let buffer = vec![
+            0xab, 0x4b, 0x54, 0x58, 0x20, 0x32, 0x30, 0xbb, 0x0d, 10, 0x1a, 10, 0x0f, 0, 0, 0, 1,
+            0, 0, 0, 4, 0, 0, 0, 4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 3, 0, 0, 0, 0, 0,
+            0, 0, 0x98, 0, 0, 0, 0x2c, 0, 0, 0, 0xc4, 0, 0, 0, 0x5c, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+            0, 0, 0, 0, 0, 0, 0, 0, 0, 0x28, 1, 0, 0, 0, 0, 0, 0, 0x10, 0, 0, 0, 0, 0, 0, 0, 0x10,
+            0, 0, 0, 0, 0, 0, 0, 0x24, 1, 0, 0, 0, 0, 0, 0, 4, 0, 0, 0, 0, 0, 0, 0, 4, 0, 0, 0, 0,
+            0, 0, 0, 0x20, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0,
+            0x2c, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0x28, 0, 1, 1, 2, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0,
+            0, 0, 0, 7, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0xff, 0, 0, 0, 0x12, 0, 0, 0, 0x4b, 0x54, 0x58,
+            0x6f, 0x72, 0x69, 0x65, 0x6e, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0, 0x72, 0x64, 0, 0,
+            0, 0x10, 0, 0, 0, 0x4b, 0x54, 0x58, 0x73, 0x77, 0x69, 0x7a, 0x7a, 0x6c, 0x65, 0, 0x72,
+            0x72, 0x72, 0x31, 0, 0x2c, 0, 0, 0, 0x4b, 0x54, 0x58, 0x77, 0x72, 0x69, 0x74, 0x65,
+            0x72, 0, 0x74, 0x6f, 0x6b, 0x74, 0x78, 0x20, 0x76, 0x34, 0x2e, 0x33, 0x2e, 0x30, 0x7e,
+            0x32, 0x38, 0x20, 0x2f, 0x20, 0x6c, 0x69, 0x62, 0x6b, 0x74, 0x78, 0x20, 0x76, 0x34,
+            0x2e, 0x33, 0x2e, 0x30, 0x7e, 0x31, 0, 0x4a, 0, 0, 0, 0x4a, 0x4a, 0x4a, 0x4a, 0x4a,
+            0x4a, 0x4a, 0x4a, 0x4a, 0x4a, 0x4a, 0x4a, 0x4a, 0x4a, 0x4a, 0x4a, 0x4a, 0x4a, 0x4a,
+            0x4a,
+        ];
+        let supported_compressed_formats = CompressedImageFormats::empty();
+        let result = ktx2_buffer_to_image(&buffer, supported_compressed_formats, true);
+        assert!(result.is_ok());
+    }
+}