Merge pull request #27348 from fredizzimo/fsundvik/fix-ext-hlstate

fix: crashes with large msgpack messages

4 files changed, 335 insertions, 50 deletions

diff --git a/src/mpack/lmpack.c b/src/mpack/lmpack.c
index ff21e29789..4ce4b5f3e5 100644
--- a/src/mpack/lmpack.c
+++ b/src/mpack/lmpack.c
@@ -882,7 +882,9 @@ static int lmpack_session_receive(lua_State *L)
   luaL_argcheck(L, (size_t)startpos <= len, 3,
       "start position must be less than or equal to the input string length");
 
-  str += (size_t)startpos - 1;
+  size_t offset = (size_t)startpos - 1 ;
+  str += offset;
+  len -= offset;
 
   if (session->unpacker != LUA_REFNIL) {
     lmpack_geti(L, session->reg, session->unpacker);
diff --git a/src/nvim/api/ui.c b/src/nvim/api/ui.c
index 5c8ebfb861..c7280253c2 100644
--- a/src/nvim/api/ui.c
+++ b/src/nvim/api/ui.c
@@ -584,7 +584,6 @@ static inline int write_cb(void *vdata, const char *buf, size_t len)
 
   data->pack_totlen += len;
   if (!data->temp_buf && UI_BUF_SIZE - BUF_POS(data) < len) {
-    data->buf_overflow = true;
     return 0;
   }
 
@@ -594,14 +593,42 @@ static inline int write_cb(void *vdata, const char *buf, size_t len)
   return 0;
 }
 
-static bool prepare_call(UI *ui, const char *name)
+static inline int size_cb(void *vdata, const char *buf, size_t len)
+{
+  UIData *data = (UIData *)vdata;
+  if (!buf) {
+    return 0;
+  }
+
+  data->pack_totlen += len;
+  return 0;
+}
+
+static void prepare_call(UI *ui, const char *name, size_t size_needed)
 {
   UIData *data = ui->data;
+  size_t name_len = strlen(name);
+  const size_t overhead = name_len + 20;
+  bool oversized_message = size_needed + overhead > UI_BUF_SIZE;
 
-  if (BUF_POS(data) > UI_BUF_SIZE - EVENT_BUF_SIZE) {
+  if (oversized_message || BUF_POS(data) > UI_BUF_SIZE - size_needed - overhead) {
     remote_ui_flush_buf(ui);
   }
 
+  if (oversized_message) {
+    // TODO(bfredl): manually testable by setting UI_BUF_SIZE to 1024 (mode_info_set)
+    data->temp_buf = xmalloc(20 + name_len + size_needed);
+    data->buf_wptr = data->temp_buf;
+    char **buf = &data->buf_wptr;
+    mpack_array(buf, 3);
+    mpack_uint(buf, 2);
+    mpack_str(buf, S_LEN("redraw"));
+    mpack_array(buf, 1);
+    mpack_array(buf, 2);
+    mpack_str(buf, name, name_len);
+    return;
+  }
+
   // To optimize data transfer(especially for "grid_line"), we bundle adjacent
   // calls to same method together, so only add a new call entry if the last
   // method call is different from "name"
@@ -614,64 +641,42 @@ static bool prepare_call(UI *ui, const char *name)
     mpack_str(buf, name, strlen(name));
     data->nevents++;
     data->ncalls = 1;
-    return true;
+    return;
   }
+}
 
-  return false;
+static void send_oversized_message(UIData *data)
+{
+  if (data->temp_buf) {
+    size_t size = (size_t)(data->buf_wptr - data->temp_buf);
+    WBuffer *buf = wstream_new_buffer(data->temp_buf, size, 1, xfree);
+    rpc_write_raw(data->channel_id, buf);
+    data->temp_buf = NULL;
+    data->buf_wptr = data->buf;
+    data->nevents_pos = NULL;
+  }
 }
 
 /// Pushes data into UI.UIData, to be consumed later by remote_ui_flush().
 static void push_call(UI *ui, const char *name, Array args)
 {
   UIData *data = ui->data;
-  bool pending = data->nevents_pos;
-  char *buf_pos_save = data->buf_wptr;
-
-  bool new_event = prepare_call(ui, name);
 
   msgpack_packer pac;
   data->pack_totlen = 0;
-  data->buf_overflow = false;
+  // First determine the needed size
+  msgpack_packer_init(&pac, data, size_cb);
+  msgpack_rpc_from_array(args, &pac);
+  // Then send the actual message
+  prepare_call(ui, name, data->pack_totlen);
   msgpack_packer_init(&pac, data, write_cb);
   msgpack_rpc_from_array(args, &pac);
-  if (data->buf_overflow) {
-    data->buf_wptr = buf_pos_save;
-    if (new_event) {
-      data->cur_event = NULL;
-      data->nevents--;
-    }
-    if (pending) {
-      remote_ui_flush_buf(ui);
-    }
 
-    size_t name_len = strlen(name);
-    if (data->pack_totlen > UI_BUF_SIZE - name_len - 20) {
-      // TODO(bfredl): manually testable by setting UI_BUF_SIZE to 1024 (mode_info_set)
-      data->temp_buf = xmalloc(20 + name_len + data->pack_totlen);
-      data->buf_wptr = data->temp_buf;
-      char **buf = &data->buf_wptr;
-      mpack_array(buf, 3);
-      mpack_uint(buf, 2);
-      mpack_str(buf, S_LEN("redraw"));
-      mpack_array(buf, 1);
-      mpack_array(buf, 2);
-      mpack_str(buf, name, name_len);
-    } else {
-      prepare_call(ui, name);
-    }
-    data->pack_totlen = 0;
-    data->buf_overflow = false;
-    msgpack_rpc_from_array(args, &pac);
-
-    if (data->temp_buf) {
-      size_t size = (size_t)(data->buf_wptr - data->temp_buf);
-      WBuffer *buf = wstream_new_buffer(data->temp_buf, size, 1, xfree);
-      rpc_write_raw(data->channel_id, buf);
-      data->temp_buf = NULL;
-      data->buf_wptr = data->buf;
-      data->nevents_pos = NULL;
-    }
+  // Oversized messages need to be sent immediately
+  if (data->temp_buf) {
+    send_oversized_message(data);
   }
+
   data->ncalls++;
 }
 
@@ -866,7 +871,7 @@ void remote_ui_raw_line(UI *ui, Integer grid, Integer row, Integer startcol, Int
 {
   UIData *data = ui->data;
   if (ui->ui_ext[kUILinegrid]) {
-    prepare_call(ui, "grid_line");
+    prepare_call(ui, "grid_line", EVENT_BUF_SIZE);
     data->ncalls++;
 
     char **buf = &data->buf_wptr;
@@ -895,7 +900,7 @@ void remote_ui_raw_line(UI *ui, Integer grid, Integer row, Integer startcol, Int
           mpack_bool(buf, false);
           remote_ui_flush_buf(ui);
 
-          prepare_call(ui, "grid_line");
+          prepare_call(ui, "grid_line", EVENT_BUF_SIZE);
           data->ncalls++;
           mpack_array(buf, 5);
           mpack_uint(buf, (uint32_t)grid);
diff --git a/src/nvim/ui_defs.h b/src/nvim/ui_defs.h
index 2245575306..a2071782b6 100644
--- a/src/nvim/ui_defs.h
+++ b/src/nvim/ui_defs.h
@@ -46,7 +46,6 @@ typedef struct {
   // state for write_cb, while packing a single arglist to msgpack. This
   // might fail due to buffer overflow.
   size_t pack_totlen;
-  bool buf_overflow;
   char *temp_buf;
 
   // We start packing the two outermost msgpack arrays before knowing the total
diff --git a/test/functional/ui/hlstate_spec.lua b/test/functional/ui/hlstate_spec.lua
index 278e6e5272..8b36ad5431 100644
--- a/test/functional/ui/hlstate_spec.lua
+++ b/test/functional/ui/hlstate_spec.lua
@@ -309,4 +309,283 @@ describe('ext_hlstate detailed highlights', function()
       {3:                                        }|
     ]])
   end)
+
+  it('combines deleted extmark highlights', function()
+    insert([[
+      line1
+        line2
+        line3
+        line4
+        line5
+      line6]])
+
+    screen:expect {
+      grid = [[
+      line1                                   |
+        line2                                 |
+        line3                                 |
+        line4                                 |
+        line5                                 |
+      line^6                                   |
+      {1:~                                       }|
+      {2:                                        }|
+    ]],
+      attr_ids = {
+        [1] = {
+          { foreground = Screen.colors.Blue, bold = true },
+          { { ui_name = 'EndOfBuffer', hi_name = 'NonText', kind = 'ui' } },
+        },
+        [2] = { {}, { { ui_name = 'MsgArea', hi_name = 'MsgArea', kind = 'ui' } } },
+      },
+    }
+
+    local ns = api.nvim_create_namespace('test')
+
+    local add_indicator = function(line, col)
+      api.nvim_buf_set_extmark(0, ns, line, col, {
+        hl_mode = 'combine',
+        priority = 2,
+        right_gravity = false,
+        virt_text = { { '|', 'Delimiter' } },
+        virt_text_win_col = 0,
+        virt_text_pos = 'overlay',
+      })
+    end
+
+    add_indicator(1, 0)
+    add_indicator(2, 0)
+    add_indicator(3, 0)
+    add_indicator(4, 0)
+
+    screen:expect {
+      grid = [[
+      line1                                   |
+      {1:|} line2                                 |
+      {1:|} line3                                 |
+      {1:|} line4                                 |
+      {1:|} line5                                 |
+      line^6                                   |
+      {2:~                                       }|
+      {3:                                        }|
+    ]],
+      attr_ids = {
+        [1] = {
+          { foreground = Screen.colors.SlateBlue },
+          { { hi_name = 'Special', kind = 'syntax' } },
+        },
+        [2] = {
+          { bold = true, foreground = Screen.colors.Blue },
+          { { ui_name = 'EndOfBuffer', kind = 'ui', hi_name = 'NonText' } },
+        },
+        [3] = { {}, { { ui_name = 'MsgArea', kind = 'ui', hi_name = 'MsgArea' } } },
+      },
+    }
+
+    helpers.feed('3ggV2jd')
+    --screen:redraw_debug()
+    screen:expect {
+      grid = [[
+      line1                                   |
+      {1:|} line2                                 |
+      {2:^|}ine6                                   |
+      {3:~                                       }|*4
+      {4:3 fewer lines                           }|
+    ]],
+      attr_ids = {
+        [1] = {
+          { foreground = Screen.colors.SlateBlue },
+          { { kind = 'syntax', hi_name = 'Special' } },
+        },
+        [2] = { { foreground = Screen.colors.SlateBlue }, { 1, 1, 1 } },
+        [3] = {
+          { bold = true, foreground = Screen.colors.Blue },
+          { { kind = 'ui', ui_name = 'EndOfBuffer', hi_name = 'NonText' } },
+        },
+        [4] = { {}, { { kind = 'ui', ui_name = 'MsgArea', hi_name = 'MsgArea' } } },
+      },
+    }
+  end)
+
+  it('removes deleted extmark highlights with invalidate', function()
+    insert([[
+      line1
+        line2
+        line3
+        line4
+        line5
+      line6]])
+
+    screen:expect {
+      grid = [[
+      line1                                   |
+        line2                                 |
+        line3                                 |
+        line4                                 |
+        line5                                 |
+      line^6                                   |
+      {1:~                                       }|
+      {2:                                        }|
+    ]],
+      attr_ids = {
+        [1] = {
+          { foreground = Screen.colors.Blue, bold = true },
+          { { ui_name = 'EndOfBuffer', hi_name = 'NonText', kind = 'ui' } },
+        },
+        [2] = { {}, { { ui_name = 'MsgArea', hi_name = 'MsgArea', kind = 'ui' } } },
+      },
+    }
+
+    local ns = api.nvim_create_namespace('test')
+
+    local add_indicator = function(line, col)
+      api.nvim_buf_set_extmark(0, ns, line, col, {
+        hl_mode = 'combine',
+        priority = 2,
+        right_gravity = false,
+        virt_text = { { '|', 'Delimiter' } },
+        virt_text_win_col = 0,
+        virt_text_pos = 'overlay',
+        invalidate = true,
+      })
+    end
+
+    add_indicator(1, 0)
+    add_indicator(2, 0)
+    add_indicator(3, 0)
+    add_indicator(4, 0)
+
+    screen:expect {
+      grid = [[
+      line1                                   |
+      {1:|} line2                                 |
+      {1:|} line3                                 |
+      {1:|} line4                                 |
+      {1:|} line5                                 |
+      line^6                                   |
+      {2:~                                       }|
+      {3:                                        }|
+    ]],
+      attr_ids = {
+        [1] = {
+          { foreground = Screen.colors.SlateBlue },
+          { { hi_name = 'Special', kind = 'syntax' } },
+        },
+        [2] = {
+          { bold = true, foreground = Screen.colors.Blue },
+          { { ui_name = 'EndOfBuffer', kind = 'ui', hi_name = 'NonText' } },
+        },
+        [3] = { {}, { { ui_name = 'MsgArea', kind = 'ui', hi_name = 'MsgArea' } } },
+      },
+    }
+
+    helpers.feed('3ggV2jd')
+    --screen:redraw_debug()
+    screen:expect {
+      grid = [[
+      line1                                   |
+      {1:|} line2                                 |
+      ^line6                                   |
+      {2:~                                       }|*4
+      {3:3 fewer lines                           }|
+    ]],
+      attr_ids = {
+        [1] = {
+          { foreground = Screen.colors.SlateBlue },
+          { { kind = 'syntax', hi_name = 'Special' } },
+        },
+        [2] = {
+          { foreground = Screen.colors.Blue, bold = true },
+          { { kind = 'ui', ui_name = 'EndOfBuffer', hi_name = 'NonText' } },
+        },
+        [3] = { {}, { { kind = 'ui', ui_name = 'MsgArea', hi_name = 'MsgArea' } } },
+      },
+    }
+  end)
+
+  it('does not hang when combining too many highlights', function()
+    local num_lines = 500
+    insert('first line\n')
+    for _ = 1, num_lines do
+      insert([[
+        line
+      ]])
+    end
+    insert('last line')
+
+    helpers.feed('gg')
+    screen:expect {
+      grid = [[
+      ^first line                              |
+        line                                  |*6
+      {1:                                        }|
+    ]],
+      attr_ids = {
+        [1] = { {}, { { kind = 'ui', hi_name = 'MsgArea', ui_name = 'MsgArea' } } },
+      },
+    }
+    local ns = api.nvim_create_namespace('test')
+
+    local add_indicator = function(line, col)
+      api.nvim_buf_set_extmark(0, ns, line, col, {
+        hl_mode = 'combine',
+        priority = 2,
+        right_gravity = false,
+        virt_text = { { '|', 'Delimiter' } },
+        virt_text_win_col = 0,
+        virt_text_pos = 'overlay',
+      })
+    end
+
+    for i = 1, num_lines do
+      add_indicator(i, 0)
+    end
+
+    screen:expect {
+      grid = [[
+      ^first line                              |
+      {1:|} line                                  |*6
+      {2:                                        }|
+    ]],
+      attr_ids = {
+        [1] = {
+          { foreground = Screen.colors.SlateBlue },
+          { { kind = 'syntax', hi_name = 'Special' } },
+        },
+        [2] = { {}, { { kind = 'ui', ui_name = 'MsgArea', hi_name = 'MsgArea' } } },
+      },
+    }
+
+    helpers.feed(string.format('3ggV%ijd', num_lines - 2))
+    --screen:redraw_debug(nil, nil, 100000)
+
+    local expected_ids = {}
+    for i = 1, num_lines - 1 do
+      expected_ids[i] = 1
+    end
+    screen:expect {
+      grid = string.format(
+        [[
+        first line                              |
+        {1:|} line                                  |
+        {2:^|}ast line                               |
+        {3:~                                       }|*4
+        {4:%-40s}|
+    ]],
+        tostring(num_lines - 1) .. ' fewer lines'
+      ),
+      attr_ids = {
+        [1] = {
+          { foreground = Screen.colors.SlateBlue },
+          { { kind = 'syntax', hi_name = 'Special' } },
+        },
+        [2] = { { foreground = Screen.colors.SlateBlue }, expected_ids },
+        [3] = {
+          { foreground = Screen.colors.Blue, bold = true },
+          { { kind = 'ui', hi_name = 'NonText', ui_name = 'EndOfBuffer' } },
+        },
+        [4] = { {}, { { kind = 'ui', hi_name = 'MsgArea', ui_name = 'MsgArea' } } },
+      },
+      timeout = 100000,
+    }
+  end)
 end)