Fix use-after-free in font.Atlas.grow (#8249)

Grow needs to allocate and might fail midway. It tries to handle this
using "undo" pattern, and restoring old state on error. But this is
exactly what steps into UAF, as, on error, both errdefer and defer are
run, and the old data is freed.

Instead, use a more robust "reservation" pattern, where we first
fallibly resrve all the resources we need, without applying any changes,
and than do the actual change once we are sure that cannot fail.

src/font/Atlas.zig

@@ -86,6 +86,11 @@ pub const Region = extern struct {
     height: u32,
 };
 
+/// Number of nodes to preallocate in the list on init.
+///
+/// TODO: figure out optimal prealloc based on real world usage
+const node_prealloc: usize = 64;
+
 pub fn init(alloc: Allocator, size: u32, format: Format) Allocator.Error!Atlas {
     var result = Atlas{
         .data = try alloc.alloc(u8, size * size * format.depth()),
@@ -95,8 +100,8 @@ pub fn init(alloc: Allocator, size: u32, format: Format) Allocator.Error!Atlas {
     };
     errdefer result.deinit(alloc);
 
-    // TODO: figure out optimal prealloc based on real world usage
-    try result.nodes.ensureUnusedCapacity(alloc, 64);
+    // Prealloc some nodes.
+    result.nodes = try .initCapacity(alloc, node_prealloc);
 
     // This sets up our initial state
     result.clear();
@@ -287,30 +292,30 @@ pub fn grow(self: *Atlas, alloc: Allocator, size_new: u32) Allocator.Error!void
     assert(size_new >= self.size);
     if (size_new == self.size) return;
 
-    // Preserve our old values so we can copy the old data
+    // We reserve space ahead of time for the new node, so that we
+    // won't have to handle any errors after allocating our new data.
+    try self.nodes.ensureUnusedCapacity(alloc, 1);
+
+    const data_new = try alloc.alloc(
+        u8,
+        size_new * size_new * self.format.depth(),
+    );
+
+    // Function is infallible from this point.
+    errdefer comptime unreachable;
+
+    // Keep track of our old data so that we can copy it.
     const data_old = self.data;
     const size_old = self.size;
 
-    // Allocate our new data
-    self.data = try alloc.alloc(u8, size_new * size_new * self.format.depth());
-    defer alloc.free(data_old);
-    errdefer {
-        alloc.free(self.data);
-        self.data = data_old;
-    }
-
-    // Add our new rectangle for our added righthand space. We do this
-    // right away since its the only operation that can fail and we want
-    // to make error cleanup easier.
-    try self.nodes.append(alloc, .{
-        .x = size_old - 1,
-        .y = 1,
-        .width = size_new - size_old,
-    });
-
-    // If our allocation and rectangle add succeeded, we can go ahead
-    // and persist our new size and copy over the old data.
+    // Update our data and size to our new ones.
+    self.data = data_new;
     self.size = size_new;
+
+    // Free the old data once we're done with it.
+    defer alloc.free(data_old);
+
+    // Zero the new data out and copy the old data over.
     @memset(self.data, 0);
     self.set(.{
         .x = 0, // don't bother skipping border so we can avoid strides
@@ -319,6 +324,13 @@ pub fn grow(self: *Atlas, alloc: Allocator, size_new: u32) Allocator.Error!void
         .height = size_old - 2, // skip the last border row
     }, data_old[size_old * self.format.depth() ..]);
 
+    // Add the new rectangle for our added righthand space.
+    self.nodes.appendAssumeCapacity(.{
+        .x = size_old - 1,
+        .y = 1,
+        .width = size_new - size_old,
+    });
+
     // We are both modified and resized
     _ = self.modified.fetchAdd(1, .monotonic);
     _ = self.resized.fetchAdd(1, .monotonic);
@@ -737,3 +749,49 @@ test "grow BGR" {
     _ = try atlas.reserve(alloc, 2, 1);
     try testing.expectError(Error.AtlasFull, atlas.reserve(alloc, 1, 1));
 }
+
+test "grow OOM" {
+    // We use a fixed buffer allocator so that we can consistently hit OOM.
+    //
+    // We calculate the size to exactly fit the 4x4 pixels and node list.
+    var buf: [
+        4 * 4 * 1 // 4x4 pixels, each 1 byte.
+        + node_prealloc * @sizeOf(Node) // preallocated nodes.
+    ]u8 = undefined;
+    var fba: std.heap.FixedBufferAllocator = .init(&buf);
+    const alloc = fba.allocator();
+
+    var atlas = try init(alloc, 4, .grayscale); // +2 for 1px border
+    defer atlas.deinit(alloc);
+
+    const reg = try atlas.reserve(alloc, 2, 2);
+    try testing.expectError(
+        Error.AtlasFull,
+        atlas.reserve(alloc, 1, 1),
+    );
+
+    // Write some data so we can verify that attempted growing doesn't mess it up.
+    atlas.set(reg, &[_]u8{ 1, 2, 3, 4 });
+    try testing.expectEqual(@as(u8, 1), atlas.data[5]);
+    try testing.expectEqual(@as(u8, 2), atlas.data[6]);
+    try testing.expectEqual(@as(u8, 3), atlas.data[9]);
+    try testing.expectEqual(@as(u8, 4), atlas.data[10]);
+
+    // Expand by 1, should give OOM, modified and resized should be unchanged.
+    const old_modified = atlas.modified.load(.monotonic);
+    const old_resized = atlas.resized.load(.monotonic);
+    try testing.expectError(
+        Allocator.Error.OutOfMemory,
+        atlas.grow(alloc, atlas.size + 1),
+    );
+    const new_modified = atlas.modified.load(.monotonic);
+    const new_resized = atlas.resized.load(.monotonic);
+    try testing.expectEqual(old_modified, new_modified);
+    try testing.expectEqual(old_resized, new_resized);
+
+    // Ensure our data is still set.
+    try testing.expectEqual(@as(u8, 1), atlas.data[5]);
+    try testing.expectEqual(@as(u8, 2), atlas.data[6]);
+    try testing.expectEqual(@as(u8, 3), atlas.data[9]);
+    try testing.expectEqual(@as(u8, 4), atlas.data[10]);
+}