terminal/kitty: prevent auto-assigned image ID collisions

Fixes #2197

Image transmissions without an explicit ID (i=) were assigned IDs from
a wrapping counter starting with no collision check. The protocol allows
clients to choose IDs anywhere in the u32 range, so an auto-assigned ID could
collide.

Number-based transmissions (I= without i=) now receive the smallest ID
not currently in use. This probes the image map in an `O(N)` fashion but
performance issues here require a pathological client and this implementation
matches Kitty's performance as well.

src/terminal/kitty/graphics_exec.zig

@@ -436,12 +436,12 @@ fn loadAndAddImage(
 
     // If the image has no ID, we assign one
     if (loading.image.id == 0) {
-        loading.image.id = storage.next_image_id;
-        storage.next_image_id +%= 1;
-
-        // If the image also has no number then its auto-ID is "implicit".
-        // See the doc comment on Image.metadata.implicit_id for more detail.
-        if (loading.image.number == 0) loading.image.metadata.implicit_id = true;
+        if (loading.image.number > 0) {
+            loading.image.id = storage.nextImageId(.explicit);
+        } else {
+            loading.image.id = storage.nextImageId(.implicit);
+            loading.image.metadata.implicit_id = true;
+        }
     }
 
     // If this is chunked, this is the beginning of a new chunked transmission.
@@ -700,7 +700,7 @@ test "kittygfx chunked success response uses initial identifiers" {
         const resp = execute(io, alloc, &t, &cmd).?;
 
         try testing.expect(resp.ok());
-        try testing.expectEqual(@as(u32, 2147483647), resp.id);
+        try testing.expectEqual(@as(u32, 1), resp.id);
         try testing.expectEqual(@as(u32, 93), resp.image_number);
         try testing.expectEqual(@as(u32, 7), resp.placement_id);
 
@@ -708,7 +708,7 @@ test "kittygfx chunked success response uses initial identifiers" {
         var writer: std.Io.Writer = .fixed(&buf);
         try resp.encode(&writer);
         try testing.expectEqualStrings(
-            "\x1b_Gi=2147483647,I=93,p=7;OK\x1b\\",
+            "\x1b_Gi=1,I=93,p=7;OK\x1b\\",
             writer.buffered(),
         );
     }
@@ -1448,3 +1448,163 @@ test "kittygfx out of range display keys are tolerated" {
     const entry = it.next().?;
     try testing.expect(entry.value_ptr.location == .virtual);
 }
+
+test "kittygfx number-based transmission assigns smallest free id" {
+    const testing = std.testing;
+    const alloc = testing.allocator;
+    const io = testing.io;
+
+    var t = try Terminal.init(io, alloc, .{ .rows = 5, .cols = 5 });
+    defer t.deinit(alloc);
+    const storage = &t.screens.active.kitty_images;
+
+    // Empty storage: the first free ID is 1, as in Kitty.
+    {
+        const cmd = try command.Parser.parseString(
+            alloc,
+            "a=t,f=24,s=1,v=1,I=42;/wAA",
+        );
+        defer cmd.deinit(alloc);
+        const resp = execute(io, alloc, &t, &cmd).?;
+        try testing.expect(resp.ok());
+        try testing.expectEqual(@as(u32, 1), resp.id);
+        try testing.expectEqual(@as(u32, 42), resp.image_number);
+    }
+
+    // Occupy ID 2 explicitly; the next number gets 3.
+    {
+        const cmd = try command.Parser.parseString(
+            alloc,
+            "a=t,f=24,s=1,v=1,i=2;/wAA",
+        );
+        defer cmd.deinit(alloc);
+        try testing.expect(execute(io, alloc, &t, &cmd).?.ok());
+    }
+    {
+        const cmd = try command.Parser.parseString(
+            alloc,
+            "a=t,f=24,s=1,v=1,I=43;/wAA",
+        );
+        defer cmd.deinit(alloc);
+        const resp = execute(io, alloc, &t, &cmd).?;
+        try testing.expect(resp.ok());
+        try testing.expectEqual(@as(u32, 3), resp.id);
+    }
+
+    // Deleting ID 1 opens a gap that the next number fills.
+    {
+        const cmd = try command.Parser.parseString(alloc, "a=d,d=I,i=1");
+        defer cmd.deinit(alloc);
+        try testing.expect(execute(io, alloc, &t, &cmd) == null);
+    }
+    {
+        const cmd = try command.Parser.parseString(
+            alloc,
+            "a=t,f=24,s=1,v=1,I=44;/wAA",
+        );
+        defer cmd.deinit(alloc);
+        const resp = execute(io, alloc, &t, &cmd).?;
+        try testing.expect(resp.ok());
+        try testing.expectEqual(@as(u32, 1), resp.id);
+    }
+
+    try testing.expectEqual(@as(usize, 3), storage.images.count());
+    try testing.expectEqual(@as(u32, 44), storage.imageById(1).?.number);
+    try testing.expectEqual(@as(u32, 0), storage.imageById(2).?.number);
+    try testing.expectEqual(@as(u32, 43), storage.imageById(3).?.number);
+}
+
+test "kittygfx number-based id assignment does not replace client image" {
+    const testing = std.testing;
+    const alloc = testing.allocator;
+    const io = testing.io;
+
+    var t = try Terminal.init(io, alloc, .{ .rows = 5, .cols = 5 });
+    defer t.deinit(alloc);
+    const storage = &t.screens.active.kitty_images;
+
+    // A client stores and places a red pixel on the ID the old wrapping
+    // counter would assign first.
+    {
+        const cmd = try command.Parser.parseString(
+            alloc,
+            "a=t,f=24,s=1,v=1,i=2147483647;/wAA",
+        );
+        defer cmd.deinit(alloc);
+        try testing.expect(execute(io, alloc, &t, &cmd).?.ok());
+    }
+    {
+        const cmd = try command.Parser.parseString(alloc, "a=p,i=2147483647,C=1");
+        defer cmd.deinit(alloc);
+        try testing.expect(execute(io, alloc, &t, &cmd).?.ok());
+    }
+
+    // A number-based transmission must not collide with it.
+    {
+        const cmd = try command.Parser.parseString(
+            alloc,
+            "a=t,f=24,s=1,v=1,I=42;AAD/",
+        );
+        defer cmd.deinit(alloc);
+        const resp = execute(io, alloc, &t, &cmd).?;
+        try testing.expect(resp.ok());
+        try testing.expectEqual(@as(u32, 1), resp.id);
+    }
+
+    // The client's image and placement are untouched.
+    try testing.expectEqual(@as(usize, 2), storage.images.count());
+    try testing.expectEqual(@as(usize, 1), storage.placements.count());
+    try testing.expectEqualSlices(
+        u8,
+        &.{ 255, 0, 0 },
+        storage.imageById(2147483647).?.data.bytes().?,
+    );
+}
+
+test "kittygfx implicit id assignment does not replace client image" {
+    const testing = std.testing;
+    const alloc = testing.allocator;
+    const io = testing.io;
+
+    var t = try Terminal.init(io, alloc, .{ .rows = 5, .cols = 5 });
+    defer t.deinit(alloc);
+    const storage = &t.screens.active.kitty_images;
+
+    // A client stores and places a red pixel on the ID the implicit
+    // counter assigns first.
+    {
+        const cmd = try command.Parser.parseString(
+            alloc,
+            "a=t,f=24,s=1,v=1,i=2147483647;/wAA",
+        );
+        defer cmd.deinit(alloc);
+        try testing.expect(execute(io, alloc, &t, &cmd).?.ok());
+    }
+    {
+        const cmd = try command.Parser.parseString(alloc, "a=p,i=2147483647,C=1");
+        defer cmd.deinit(alloc);
+        try testing.expect(execute(io, alloc, &t, &cmd).?.ok());
+    }
+
+    // Transmit without an ID or number: no response, and the counter
+    // skips over the in-use ID instead of replacing that image.
+    {
+        const cmd = try command.Parser.parseString(
+            alloc,
+            "a=t,f=24,s=1,v=1;AAD/",
+        );
+        defer cmd.deinit(alloc);
+        try testing.expect(execute(io, alloc, &t, &cmd) == null);
+    }
+
+    try testing.expectEqual(@as(usize, 2), storage.images.count());
+    try testing.expectEqual(@as(usize, 1), storage.placements.count());
+    try testing.expectEqualSlices(
+        u8,
+        &.{ 255, 0, 0 },
+        storage.imageById(2147483647).?.data.bytes().?,
+    );
+    const implicit = storage.imageById(2147483648).?;
+    try testing.expect(implicit.metadata.implicit_id);
+    try testing.expectEqualSlices(u8, &.{ 0, 0, 255 }, implicit.data.bytes().?);
+}

src/terminal/kitty/graphics_image.zig

@@ -635,8 +635,12 @@ pub const Image = struct {
         transient: bool = false,
 
         /// Set this if the image was loaded without an ID or number. Such
-        /// images must not receive responses even though they currently get
-        /// IDs in the public range (which is bad!).
+        /// images must not receive responses. Kitty gives these client ID
+        /// 0 (unaddressable); our storage keys everything by one public
+        /// u32 ID, so they get an ID from the upper half of the range
+        /// that is guaranteed unused at assignment time, but a client
+        /// that explicitly transmits that ID later can still replace
+        /// them.
         implicit_id: bool = false,
 
         /// Number of placements referencing this image.

src/terminal/kitty/graphics_storage.zig

@@ -101,11 +101,10 @@ pub const ImageStorage = struct {
     /// This field must only be written via markMutated.
     generation: u64 = 0,
 
-    /// This is the next automatically assigned image ID. We start mid-way
-    /// through the u32 range to avoid collisions with buggy programs.
-    /// TODO: This isn't good enough, it's perfectly legal for programs
-    ///       to use IDs in the latter half of the range and collisions
-    ///       are not gracefully handled.
+    /// This is the next automatically assigned image ID for images
+    /// transmitted without an ID or number. We start mid-way through
+    /// the u32 range to stay clear of the low IDs client programs
+    /// typically pick. See nextImageId.
     next_image_id: u32 = 2147483647,
 
     /// This is the next automatically assigned placement ID. This is never
@@ -228,6 +227,40 @@ pub const ImageStorage = struct {
         self.total_limit = limit;
     }
 
+    /// Returns the next ID to automatically assign to an image that
+    /// was transmitted without an explicit ID (i=). The result is
+    /// never zero (zero means "no ID") and never an ID currently in
+    /// use, so an automatic ID never replaces an existing image.
+    pub fn nextImageId(
+        self: *ImageStorage,
+        mode: enum { explicit, implicit },
+    ) u32 {
+        // Starting ID depends on mode.
+        var id: u32 = switch (mode) {
+            // Yeah, starting from 1 is wasteful. This matches what Kitty
+            // does. In the future we can probably cache a low-water-mark
+            // based on delete behavior.
+            .explicit => 1,
+            .implicit => self.next_image_id,
+        };
+
+        // Go through all possible images. We have +2 because we can
+        // exceed the total by 1 (new ID).
+        const count: usize = self.images.count();
+        for (0..count + 2) |_| {
+            if (id != 0 and !self.images.contains(id)) break;
+            id +%= 1;
+        }
+
+        // Zero is never allowed.
+        if (id == 0) id = 1;
+
+        // Keep track of our next image id for implicit to speed that up.
+        if (mode == .implicit) self.next_image_id = id +% 1;
+
+        return id;
+    }
+
     /// Add an image to the storage. This will automatically free any existing
     /// image with the same ID. Prefer addPendingImage for pending data so the
     /// caller receives a completion token.
@@ -2802,3 +2835,54 @@ test "storage: pending images share exact eviction ordering" {
     try testing.expect(!s.images.contains(3));
     try testing.expectEqual(@as(usize, 128), s.total_bytes);
 }
+
+test "storage: nextImageId number matches Kitty get_free_client_id" {
+    const testing = std.testing;
+    const io = testing.io;
+    const alloc = testing.allocator;
+    var t = try terminal.Terminal.init(io, alloc, .{ .rows = 3, .cols = 3 });
+    defer t.deinit(alloc);
+
+    var s: ImageStorage = .{};
+    defer s.deinit(alloc, t.screens.active);
+
+    // Empty storage returns 1, as does Kitty.
+    try testing.expectEqual(@as(u32, 1), s.nextImageId(.explicit));
+
+    // Contiguous IDs from 1: the next ID after them.
+    try s.addImage(io, alloc, t.screens.active, .{ .id = 1 });
+    try s.addImage(io, alloc, t.screens.active, .{ .id = 2 });
+    try s.addImage(io, alloc, t.screens.active, .{ .id = 3 });
+    try testing.expectEqual(@as(u32, 4), s.nextImageId(.explicit));
+
+    // The first gap is filled, even when higher IDs exist.
+    try s.addImage(io, alloc, t.screens.active, .{ .id = 100 });
+    s.delete(io, alloc, &t, .{ .id = .{ .image_id = 2, .delete = true } });
+    try testing.expectEqual(@as(u32, 2), s.nextImageId(.explicit));
+
+    // 1 is reused as soon as it is free.
+    s.delete(io, alloc, &t, .{ .id = .{ .image_id = 1, .delete = true } });
+    try testing.expectEqual(@as(u32, 1), s.nextImageId(.explicit));
+}
+
+test "storage: nextImageId implicit skips in-use ids and zero" {
+    const testing = std.testing;
+    const io = testing.io;
+    const alloc = testing.allocator;
+    var t = try terminal.Terminal.init(io, alloc, .{ .rows = 3, .cols = 3 });
+    defer t.deinit(alloc);
+
+    var s: ImageStorage = .{};
+    defer s.deinit(alloc, t.screens.active);
+
+    // A client image already sits on the counter's first two values.
+    try s.addImage(io, alloc, t.screens.active, .{ .id = 2147483647 });
+    try s.addImage(io, alloc, t.screens.active, .{ .id = 2147483648 });
+    try testing.expectEqual(@as(u32, 2147483649), s.nextImageId(.implicit));
+    try testing.expectEqual(@as(u32, 2147483650), s.nextImageId(.implicit));
+
+    // Wrapping skips zero, which means "no ID" protocol-wide.
+    s.next_image_id = std.math.maxInt(u32);
+    try testing.expectEqual(std.math.maxInt(u32), s.nextImageId(.implicit));
+    try testing.expectEqual(@as(u32, 1), s.nextImageId(.implicit));
+}