From 89ab9fdf7045d812a8cbdb57836a5952fbca2a58 Mon Sep 17 00:00:00 2001
From: Techatrix <19954306+Techatrix@users.noreply.github.com>
Date: Mon, 27 Feb 2023 22:53:16 +0000
Subject: [PATCH] make `diff.edits` memory safe (#1026)

---
 src/diff.zig           | 32 ++++++++++++++++++++++++----
 tests/utility/diff.zig | 47 ++++++++++++++++++++++++------------------
 2 files changed, 55 insertions(+), 24 deletions(-)

diff --git a/src/diff.zig b/src/diff.zig
index be378da..749a07e 100644
--- a/src/diff.zig
+++ b/src/diff.zig
@@ -7,7 +7,7 @@ const dmp = DiffMatchPatch{
     .diff_timeout = 250,
 };
 
-pub const Error = error{ OutOfMemory, InvalidRange, UnknownError };
+pub const Error = error{OutOfMemory};
 
 pub fn edits(
     allocator: std.mem.Allocator,
@@ -15,8 +15,25 @@ pub fn edits(
     after: []const u8,
     encoding: offsets.Encoding,
 ) Error!std.ArrayListUnmanaged(types.TextEdit) {
-    var diffs = try dmp.diff(allocator, before, after, true);
+    var arena = std.heap.ArenaAllocator.init(allocator);
+    defer arena.deinit();
+    var diffs = try dmp.diff(arena.allocator(), before, after, true);
+
+    var edit_count: usize = 0;
+    for (diffs.items) |diff| {
+        switch (diff.operation) {
+            .delete => edit_count += 1,
+            .equal => continue,
+            .insert => edit_count += 1,
+        }
+    }
+
     var eds = std.ArrayListUnmanaged(types.TextEdit){};
+    try eds.ensureTotalCapacity(allocator, edit_count);
+    errdefer {
+        for (eds.items) |edit| allocator.free(edit.newText);
+        eds.deinit(allocator);
+    }
 
     var offset: usize = 0;
     for (diffs.items) |diff| {
@@ -24,13 +41,19 @@ pub fn edits(
         switch (diff.operation) {
             .delete => {
                 offset += diff.text.len;
-                try eds.append(allocator, .{ .range = offsets.locToRange(before, .{ .start = start, .end = offset }, encoding), .newText = "" });
+                eds.appendAssumeCapacity(.{
+                    .range = offsets.locToRange(before, .{ .start = start, .end = offset }, encoding),
+                    .newText = "",
+                });
             },
             .equal => {
                 offset += diff.text.len;
             },
             .insert => {
-                try eds.append(allocator, .{ .range = offsets.locToRange(before, .{ .start = start, .end = start }, encoding), .newText = diff.text });
+                eds.appendAssumeCapacity(.{
+                    .range = offsets.locToRange(before, .{ .start = start, .end = start }, encoding),
+                    .newText = try allocator.dupe(u8, diff.text),
+                });
             },
         }
     }
@@ -95,6 +118,7 @@ pub fn applyTextEdits(
     std.sort.sort(types.TextEdit, text_edits_sortable, {}, textEditLessThan);
 
     var final_text = std.ArrayListUnmanaged(u8){};
+    errdefer final_text.deinit(allocator);
 
     var last: usize = 0;
     for (text_edits_sortable) |te| {
diff --git a/tests/utility/diff.zig b/tests/utility/diff.zig
index 0c4ae54..18cb7b1 100644
--- a/tests/utility/diff.zig
+++ b/tests/utility/diff.zig
@@ -1,30 +1,37 @@
 const std = @import("std");
 const zls = @import("zls");
 
-const allocator = std.testing.allocator;
-
 fn gen(alloc: std.mem.Allocator, rand: std.rand.Random) ![]const u8 {
-    var buffer = try alloc.alloc(u8, rand.intRangeAtMost(usize, 16, 1024));
-    for (buffer) |*b| b.* = rand.intRangeAtMost(u8, 32, 126);
+    var buffer = try alloc.alloc(u8, rand.intRangeAtMost(usize, 0, 256));
+    for (buffer) |*b| b.* = rand.intRangeAtMost(u8, ' ', '~');
     return buffer;
 }
 
 test "diff - random" {
-    var arena = std.heap.ArenaAllocator.init(allocator);
-    defer arena.deinit();
-
-    var rand = std.rand.DefaultPrng.init(0);
-
-    var index: usize = 0;
-
-    while (index < 100) : (index += 1) {
-        defer _ = arena.reset(.retain_capacity);
-
-        const pre = try gen(arena.allocator(), rand.random());
-        const post = try gen(arena.allocator(), rand.random());
-
-        var edits = try zls.diff.edits(arena.allocator(), pre, post, .@"utf-8");
-        const applied = try zls.diff.applyTextEdits(arena.allocator(), pre, edits.items, .@"utf-8");
-        try std.testing.expectEqualStrings(post, applied);
+    const allocator = std.testing.allocator;
+    try std.testing.checkAllAllocationFailures(allocator, testDiff, .{ 0, .@"utf-8" });
+    for (0..30) |i| {
+        try testDiff(allocator, i, .@"utf-8");
+        try testDiff(allocator, i, .@"utf-16");
+        try testDiff(allocator, i, .@"utf-32");
     }
 }
+
+fn testDiff(allocator: std.mem.Allocator, seed: u64, encoding: zls.offsets.Encoding) !void {
+    var rand = std.rand.DefaultPrng.init(seed);
+    const before = try gen(allocator, rand.random());
+    defer allocator.free(before);
+    const after = try gen(allocator, rand.random());
+    defer allocator.free(after);
+
+    var edits = try zls.diff.edits(allocator, before, after, encoding);
+    defer {
+        for (edits.items) |edit| allocator.free(edit.newText);
+        edits.deinit(allocator);
+    }
+
+    const applied = try zls.diff.applyTextEdits(allocator, before, edits.items, encoding);
+    defer allocator.free(applied);
+
+    try std.testing.expectEqualStrings(after, applied);
+}