fix(inccommand): restrict cmdpreview undo calls (#24289)

Problem:
The cmdpreview saved undo nodes on cmdpreview_prepare() from ex_getln.c may
become invalid (free) if the preview function makes undo operations, causing
heap-use-after-free errors.

Solution:
Save the buffer undo list on cmdpreview_prepare)_ and start a new empty one. On
cmdpreview_restore_state(), undo all the entries in the new undo list and
restore the original one. With this approach, the preview function will be
allowed to undo only its own changes.

Fix #20036
Fix #20248

2 files changed, 99 insertions, 26 deletions

diff --git a/src/nvim/ex_getln.c b/src/nvim/ex_getln.c
index fe72860628..9dcfa99a37 100644
--- a/src/nvim/ex_getln.c
+++ b/src/nvim/ex_getln.c
@@ -130,25 +130,38 @@ typedef struct command_line_state {
   buf_T *b_im_ptr_buf;  ///< buffer where b_im_ptr is valid
 } CommandLineState;
 
-typedef struct cmdpreview_win_info {
-  win_T *win;
-  pos_T save_w_cursor;
-  viewstate_T save_viewstate;
-  int save_w_p_cul;
-  int save_w_p_cuc;
-} CpWinInfo;
+typedef struct cmdpreview_undo_info {
+  u_header_T *save_b_u_oldhead;
+  u_header_T *save_b_u_newhead;
+  u_header_T *save_b_u_curhead;
+  int save_b_u_numhead;
+  bool save_b_u_synced;
+  long save_b_u_seq_last;
+  long save_b_u_save_nr_last;
+  long save_b_u_seq_cur;
+  time_t save_b_u_time_cur;
+  long save_b_u_save_nr_cur;
+  char *save_b_u_line_ptr;
+  linenr_T save_b_u_line_lnum;
+  colnr_T save_b_u_line_colnr;
+} CpUndoInfo;
 
 typedef struct cmdpreview_buf_info {
   buf_T *buf;
-  bool save_b_u_synced;
-  time_t save_b_u_time_cur;
-  long save_b_u_seq_cur;
-  u_header_T *save_b_u_newhead;
   long save_b_p_ul;
   int save_b_changed;
   varnumber_T save_changedtick;
+  CpUndoInfo undo_info;
 } CpBufInfo;
 
+typedef struct cmdpreview_win_info {
+  win_T *win;
+  pos_T save_w_cursor;
+  viewstate_T save_viewstate;
+  int save_w_p_cul;
+  int save_w_p_cuc;
+} CpWinInfo;
+
 typedef struct cmdpreview_info {
   kvec_t(CpWinInfo) win_info;
   kvec_t(CpBufInfo) buf_info;
@@ -2277,8 +2290,48 @@ static void cmdpreview_close_win(void)
   }
 }
 
+/// Save the undo state of a buffer for command preview.
+static void cmdpreview_save_undo(CpUndoInfo *cp_undoinfo, buf_T *buf)
+  FUNC_ATTR_NONNULL_ALL
+{
+  cp_undoinfo->save_b_u_synced = buf->b_u_synced;
+  cp_undoinfo->save_b_u_oldhead = buf->b_u_oldhead;
+  cp_undoinfo->save_b_u_newhead = buf->b_u_newhead;
+  cp_undoinfo->save_b_u_curhead = buf->b_u_curhead;
+  cp_undoinfo->save_b_u_numhead = buf->b_u_numhead;
+  cp_undoinfo->save_b_u_seq_last = buf->b_u_seq_last;
+  cp_undoinfo->save_b_u_save_nr_last = buf->b_u_save_nr_last;
+  cp_undoinfo->save_b_u_seq_cur = buf->b_u_seq_cur;
+  cp_undoinfo->save_b_u_time_cur = buf->b_u_time_cur;
+  cp_undoinfo->save_b_u_save_nr_cur = buf->b_u_save_nr_cur;
+  cp_undoinfo->save_b_u_line_ptr = buf->b_u_line_ptr;
+  cp_undoinfo->save_b_u_line_lnum = buf->b_u_line_lnum;
+  cp_undoinfo->save_b_u_line_colnr = buf->b_u_line_colnr;
+}
+
+/// Restore the undo state of a buffer for command preview.
+static void cmdpreview_restore_undo(const CpUndoInfo *cp_undoinfo, buf_T *buf)
+{
+  buf->b_u_oldhead = cp_undoinfo->save_b_u_oldhead;
+  buf->b_u_newhead = cp_undoinfo->save_b_u_newhead;
+  buf->b_u_curhead = cp_undoinfo->save_b_u_curhead;
+  buf->b_u_numhead = cp_undoinfo->save_b_u_numhead;
+  buf->b_u_seq_last = cp_undoinfo->save_b_u_seq_last;
+  buf->b_u_save_nr_last = cp_undoinfo->save_b_u_save_nr_last;
+  buf->b_u_seq_cur = cp_undoinfo->save_b_u_seq_cur;
+  buf->b_u_time_cur = cp_undoinfo->save_b_u_time_cur;
+  buf->b_u_save_nr_cur = cp_undoinfo->save_b_u_save_nr_cur;
+  buf->b_u_line_ptr = cp_undoinfo->save_b_u_line_ptr;
+  buf->b_u_line_lnum = cp_undoinfo->save_b_u_line_lnum;
+  buf->b_u_line_colnr = cp_undoinfo->save_b_u_line_colnr;
+  if (buf->b_u_curhead == NULL) {
+    buf->b_u_synced = cp_undoinfo->save_b_u_synced;
+  }
+}
+
 /// Save current state and prepare windows and buffers for command preview.
 static void cmdpreview_prepare(CpInfo *cpinfo)
+  FUNC_ATTR_NONNULL_ALL
 {
   kv_init(cpinfo->buf_info);
   kv_init(cpinfo->win_info);
@@ -2294,14 +2347,13 @@ static void cmdpreview_prepare(CpInfo *cpinfo)
     CpBufInfo cp_bufinfo;
     cp_bufinfo.buf = buf;
 
-    cp_bufinfo.save_b_u_synced = buf->b_u_synced;
-    cp_bufinfo.save_b_u_time_cur = buf->b_u_time_cur;
-    cp_bufinfo.save_b_u_seq_cur = buf->b_u_seq_cur;
-    cp_bufinfo.save_b_u_newhead = buf->b_u_newhead;
     cp_bufinfo.save_b_p_ul = buf->b_p_ul;
     cp_bufinfo.save_b_changed = buf->b_changed;
     cp_bufinfo.save_changedtick = buf_get_changedtick(buf);
 
+    cmdpreview_save_undo(&cp_bufinfo.undo_info, buf);
+    u_clearall(buf);
+
     kv_push(cpinfo->buf_info, cp_bufinfo);
 
     buf->b_p_ul = LONG_MAX;     // Make sure we can undo all changes
@@ -2336,8 +2388,9 @@ static void cmdpreview_prepare(CpInfo *cpinfo)
   u_sync(true);
 }
 
-// Restore the state of buffers and windows before command preview.
+/// Restore the state of buffers and windows for command preview.
 static void cmdpreview_restore_state(CpInfo *cpinfo)
+  FUNC_ATTR_NONNULL_ALL
 {
   for (size_t i = 0; i < cpinfo->buf_info.size; i++) {
     CpBufInfo cp_bufinfo = cpinfo->buf_info.items[i];
@@ -2345,31 +2398,29 @@ static void cmdpreview_restore_state(CpInfo *cpinfo)
 
     buf->b_changed = cp_bufinfo.save_b_changed;
 
-    if (buf->b_u_seq_cur != cp_bufinfo.save_b_u_seq_cur) {
+    if (buf->b_u_seq_cur != cp_bufinfo.undo_info.save_b_u_seq_cur) {
       int count = 0;
 
       // Calculate how many undo steps are necessary to restore earlier state.
       for (u_header_T *uhp = buf->b_u_curhead ? buf->b_u_curhead : buf->b_u_newhead;
-           uhp != NULL && uhp->uh_seq > cp_bufinfo.save_b_u_seq_cur;
+           uhp != NULL;
            uhp = uhp->uh_next.ptr, ++count) {}
 
       aco_save_T aco;
       aucmd_prepbuf(&aco, buf);
+      // Ensure all the entries will be undone
+      if (curbuf->b_u_synced == false) {
+        u_sync(true);
+      }
       // Undo invisibly. This also moves the cursor!
       if (!u_undo_and_forget(count, false)) {
         abort();
       }
       aucmd_restbuf(&aco);
-
-      // Restore newhead. It is meaningless when curhead is valid, but we must
-      // restore it so that undotree() is identical before/after the preview.
-      buf->b_u_newhead = cp_bufinfo.save_b_u_newhead;
-      buf->b_u_time_cur = cp_bufinfo.save_b_u_time_cur;
     }
 
-    if (buf->b_u_curhead == NULL) {
-      buf->b_u_synced = cp_bufinfo.save_b_u_synced;
-    }
+    u_blockfree(buf);
+    cmdpreview_restore_undo(&cp_bufinfo.undo_info, buf);
 
     if (cp_bufinfo.save_changedtick != buf_get_changedtick(buf)) {
       buf_set_changedtick(buf, cp_bufinfo.save_changedtick);
diff --git a/test/functional/ui/inccommand_user_spec.lua b/test/functional/ui/inccommand_user_spec.lua
index 9cc6e095c5..8a1030fa25 100644
--- a/test/functional/ui/inccommand_user_spec.lua
+++ b/test/functional/ui/inccommand_user_spec.lua
@@ -435,6 +435,28 @@ describe("'inccommand' for user commands", function()
     ]])
     assert_alive()
   end)
+
+  it("no crash if preview callback executes undo", function()
+    command('set inccommand=nosplit')
+    exec_lua([[
+      vim.api.nvim_create_user_command('Foo', function() end, {
+        nargs = '?',
+        preview = function(_, _, _)
+          vim.cmd.undo()
+        end,
+      })
+    ]])
+
+    -- Clear undo history
+    command('set undolevels=-1')
+    feed('ggyyp')
+    command('set undolevels=1000')
+
+    feed('yypp:Fo')
+    assert_alive()
+    feed('<Esc>:Fo')
+    assert_alive()
+  end)
 end)
 
 describe("'inccommand' with multiple buffers", function()