From b3bda2f0438da44a0a74fd5d64d6eecdac98d691 Mon Sep 17 00:00:00 2001
From: Sean Dewar <seandewar@users.noreply.github.com>
Date: Tue, 6 Feb 2024 01:13:35 +0000
Subject: [PATCH] fix(tui): `space_buf` overflow when clearing screen (#27352)

Problem: `tui->space_buf` may be smaller than the width of the TUI or widest
grid, causing an overflow when calling `tui_grid_clear` if `print_spaces` is
called from `clear_region` (clears the TUI's screen since #23428).

Solution: resize `space_buf` to be wide enough to fit the TUI or widest grid.

Didn't bother shrinking the allocation if the max widths decrease.
---
 src/nvim/tui/tui.c                    | 16 +++++++---
 test/functional/terminal/tui_spec.lua | 42 +++++++++++++++++++++++++++
 2 files changed, 54 insertions(+), 4 deletions(-)

diff --git a/src/nvim/tui/tui.c b/src/nvim/tui/tui.c
index d4bbb8a924..5565bd1055 100644
--- a/src/nvim/tui/tui.c
+++ b/src/nvim/tui/tui.c
@@ -137,6 +137,7 @@ struct TUIData {
     int sync;
   } unibi_ext;
   char *space_buf;
+  size_t space_buf_len;
   bool stopped;
   int seen_error_exit;
   int width;
@@ -1055,10 +1056,7 @@ void tui_grid_resize(TUIData *tui, Integer g, Integer width, Integer height)
 {
   UGrid *grid = &tui->grid;
   ugrid_resize(grid, (int)width, (int)height);
-
-  xfree(tui->space_buf);
-  tui->space_buf = xmalloc((size_t)width * sizeof(*tui->space_buf));
-  memset(tui->space_buf, ' ', (size_t)width);
+  ensure_space_buf_size(tui, (size_t)width);
 
   // resize might not always be followed by a clear before flush
   // so clip the invalid region
@@ -1642,6 +1640,15 @@ static void invalidate(TUIData *tui, int top, int bot, int left, int right)
   }
 }
 
+static void ensure_space_buf_size(TUIData *tui, size_t len)
+{
+  if (len > tui->space_buf_len) {
+    tui->space_buf = xrealloc(tui->space_buf, len * sizeof *tui->space_buf);
+    memset(tui->space_buf + tui->space_buf_len, ' ', len - tui->space_buf_len);
+    tui->space_buf_len = len;
+  }
+}
+
 /// Tries to get the user's wanted dimensions (columns and rows) for the entire
 /// application (i.e., the host terminal).
 void tui_guess_size(TUIData *tui)
@@ -1678,6 +1685,7 @@ void tui_guess_size(TUIData *tui)
 
   tui->width = width;
   tui->height = height;
+  ensure_space_buf_size(tui, (size_t)tui->width);
 
   // Redraw on SIGWINCH event if size didn't change. #23411
   ui_client_set_size(width, height);
diff --git a/test/functional/terminal/tui_spec.lua b/test/functional/terminal/tui_spec.lua
index 121664ae84..a78994f998 100644
--- a/test/functional/terminal/tui_spec.lua
+++ b/test/functional/terminal/tui_spec.lua
@@ -2095,6 +2095,48 @@ describe('TUI', function()
     ]],
     }
   end)
+
+  it('no heap-buffer-overflow when changing &columns', function()
+    -- Set a different bg colour and change $TERM to something dumber so the `print_spaces()`
+    -- codepath in `clear_region()` is hit.
+    local screen = thelpers.setup_child_nvim({
+      '-u',
+      'NONE',
+      '-i',
+      'NONE',
+      '--cmd',
+      'set notermguicolors | highlight Normal ctermbg=red',
+      '--cmd',
+      'call setline(1, ["a"->repeat(&columns)])',
+    }, { env = { TERM = 'ansi' } })
+
+    screen:expect {
+      grid = [[
+      aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa|
+      ~                                                 |*3
+      [No Name] [+]                   1,1            All|
+                                                        |
+      -- TERMINAL --                                    |
+    ]],
+      attr_ids = {},
+    }
+
+    feed_data(':set columns=12\n')
+    screen:expect {
+      grid = [[
+      aaaaaaaaaaaa                                      |*4
+      < [+] 1,1                                         |
+                                                        |
+      -- TERMINAL --                                    |
+    ]],
+      attr_ids = {},
+    }
+
+    -- Wider than TUI, so screen state will look weird.
+    -- Wait for the statusline to redraw to confirm that the TUI lives and ASAN is happy.
+    feed_data(':set columns=99|set stl=redrawn%m\n')
+    screen:expect({ any = 'redrawn%[%+%]' })
+  end)
 end)
 
 describe('TUI UIEnter/UILeave', function()