From 1b7fbd00d1fcc45708f8427deb240d5c3cb4dd5c Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Sun, 20 Aug 2023 15:02:29 -0700
Subject: [PATCH] terminal/kitty-gfx: add some validation from Kitty

---
 src/terminal/kitty/graphics_exec.zig  | 24 +++++++++++++-----
 src/terminal/kitty/graphics_image.zig | 35 +++++++++++++++++++++++++++
 2 files changed, 53 insertions(+), 6 deletions(-)

diff --git a/src/terminal/kitty/graphics_exec.zig b/src/terminal/kitty/graphics_exec.zig
index 0cb77a9a0..f21e82c18 100644
--- a/src/terminal/kitty/graphics_exec.zig
+++ b/src/terminal/kitty/graphics_exec.zig
@@ -50,6 +50,13 @@ pub fn execute(
 fn query(alloc: Allocator, cmd: *Command) Response {
     const t = cmd.control.query;
 
+    // Query requires image ID. We can't actually send a response without
+    // an image ID either but we return an error and this will be logged
+    // downstream.
+    if (t.image_id == 0) {
+        return .{ .message = "EINVAL: image ID required" };
+    }
+
     // Build a partial response to start
     var result: Response = .{
         .id = t.image_id,
@@ -61,14 +68,19 @@ fn query(alloc: Allocator, cmd: *Command) Response {
     if (Image.load(alloc, t, cmd.data)) |img| {
         // Tell the command we've consumed the data.
         _ = cmd.toOwnedData();
+        defer {
+            // We need a mutable reference to deinit the image.
+            var img_c = img;
+            img_c.deinit(alloc);
+        }
 
-        // We need a mutable reference to deinit the image.
-        var img_c = img;
-        img_c.deinit(alloc);
+        // If the image is greater than a predetermined max size, then we
+        // error. The max size here is taken directly from Kitty.
     } else |err| switch (err) {
-        error.InvalidData => result.message = "ERROR: invalid data",
-        error.UnsupportedFormat => result.message = "ERROR: unsupported format",
-        error.DimensionsRequired => result.message = "ERROR: dimensions required",
+        error.InvalidData => result.message = "EINVAL: invalid data",
+        error.UnsupportedFormat => result.message = "EINVAL: unsupported format",
+        error.DimensionsRequired => result.message = "EINVAL: dimensions required",
+        error.DimensionsTooLarge => result.message = "EINVAL: dimensions too large",
     }
 
     return result;
diff --git a/src/terminal/kitty/graphics_image.zig b/src/terminal/kitty/graphics_image.zig
index b327d1d68..de564511c 100644
--- a/src/terminal/kitty/graphics_image.zig
+++ b/src/terminal/kitty/graphics_image.zig
@@ -5,6 +5,9 @@ const ArenaAllocator = std.heap.ArenaAllocator;
 
 const command = @import("graphics_command.zig");
 
+/// Maximum width or height of an image. Taken directly from Kitty.
+const max_dimension = 10000;
+
 pub const Image = struct {
     id: u32 = 0,
     data: []const u8,
@@ -12,6 +15,7 @@ pub const Image = struct {
     pub const Error = error{
         InvalidData,
         DimensionsRequired,
+        DimensionsTooLarge,
         UnsupportedFormat,
     };
 
@@ -33,6 +37,7 @@ pub const Image = struct {
         data: []const u8,
     ) !Image {
         if (t.width == 0 or t.height == 0) return error.DimensionsRequired;
+        if (t.width > max_dimension or t.height > max_dimension) return error.DimensionsTooLarge;
 
         // Data length must be what we expect
         // NOTE: we use a "<" check here because Kitty itself doesn't validate
@@ -71,3 +76,33 @@ test "image load with invalid RGB data" {
     }, data);
     defer img.deinit(alloc);
 }
+
+test "image load with image too wide" {
+    const testing = std.testing;
+    const alloc = testing.allocator;
+
+    var data = try alloc.dupe(u8, "AAAA");
+    defer alloc.free(data);
+
+    try testing.expectError(error.DimensionsTooLarge, Image.load(alloc, .{
+        .format = .rgb,
+        .width = max_dimension + 1,
+        .height = 1,
+        .image_id = 31,
+    }, data));
+}
+
+test "image load with image too tall" {
+    const testing = std.testing;
+    const alloc = testing.allocator;
+
+    var data = try alloc.dupe(u8, "AAAA");
+    defer alloc.free(data);
+
+    try testing.expectError(error.DimensionsTooLarge, Image.load(alloc, .{
+        .format = .rgb,
+        .height = max_dimension + 1,
+        .width = 1,
+        .image_id = 31,
+    }, data));
+}