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

fix: crashes with large msgpack messages

1 files changed, 53 insertions, 48 deletions

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);