3 files changed, 45 insertions, 17 deletions

diff --git a/src/nvim/event/socket.c b/src/nvim/event/socket.c
index 7744d3ff26..2205fc08e4 100644
--- a/src/nvim/event/socket.c
+++ b/src/nvim/event/socket.c
@@ -97,10 +97,10 @@ int socket_watcher_init(Loop *loop, SocketWatcher *watcher, const char *endpoint
   return 0;
 }
 
-/// Callback for closing a handle initialized by socket_connect().
-static void connect_close_cb(uv_handle_t *handle)
+/// Callback after closing a Stream initialized by socket_connect().
+static void connect_close_cb(Stream *stream, void *data)
 {
-  bool *closed = handle->data;
+  bool *closed = data;
   *closed = true;
 }
 
@@ -110,7 +110,7 @@ static void connect_close_cb(uv_handle_t *handle)
 /// @return true if socket is alive (connection succeeded), false otherwise
 static bool socket_alive(Loop *loop, const char *addr)
 {
-  RStream stream;
+  RStream stream = { 0 };
   const char *error = NULL;
 
   // Try to connect with a 500ms timeout (fast failure for dead sockets)
@@ -121,13 +121,20 @@ static bool socket_alive(Loop *loop, const char *addr)
 
   // Connection succeeded - socket is alive. Close the probe connection properly.
   bool closed = false;
-  stream.s.uv.pipe.data = &closed;
-  uv_close((uv_handle_t *)&stream.s.uv.pipe, connect_close_cb);
+  stream.s.internal_close_cb = connect_close_cb;
+  stream.s.internal_data = &closed;
+  stream_may_close(&stream.s);
   LOOP_PROCESS_EVENTS_UNTIL(&main_loop, NULL, -1, closed);
 
   return true;
 }
 
+static void early_server_close_cb(uv_handle_t *handle)
+{
+  bool *closed = handle->data;
+  *closed = true;
+}
+
 int socket_watcher_start(SocketWatcher *watcher, int backlog, socket_cb cb)
   FUNC_ATTR_NONNULL_ALL
 {
@@ -181,7 +188,7 @@ int socket_watcher_start(SocketWatcher *watcher, int backlog, socket_cb cb)
           uv_loop_t *uv_loop = watcher->uv.pipe.handle.loop;
           bool closed = false;
           watcher->uv.pipe.handle.data = &closed;
-          uv_close((uv_handle_t *)&watcher->uv.pipe.handle, connect_close_cb);
+          uv_close((uv_handle_t *)&watcher->uv.pipe.handle, early_server_close_cb);
           LOOP_PROCESS_EVENTS_UNTIL(&main_loop, NULL, -1, closed);
 
           uv_pipe_init(uv_loop, &watcher->uv.pipe.handle, 0);
@@ -276,8 +283,8 @@ static void connect_cb(uv_connect_t *req, int status)
   int *ret_status = req->data;
   *ret_status = status;
   uv_handle_t *handle = (uv_handle_t *)req->handle;
-  if (status != 0 && !uv_is_closing(handle)) {
-    uv_close(handle, connect_close_cb);
+  if (status != 0) {
+    stream_may_close(handle->data);
   }
 }
 
@@ -327,18 +334,16 @@ tcp_retry:
     uv_pipe_connect(&req,  pipe, address, connect_cb);
     uv_stream = (uv_stream_t *)pipe;
   }
-  uv_stream->data = &closed;
+  stream_init(NULL, &stream->s, -1, uv_stream);
+  stream->s.internal_close_cb = connect_close_cb;
+  stream->s.internal_data = &closed;
   closed = false;
   status = 1;
   LOOP_PROCESS_EVENTS_UNTIL(&main_loop, NULL, timeout, status != 1);
   if (status == 0) {
-    stream_init(NULL, &stream->s, -1, uv_stream);
-    assert(uv_stream->data != &closed);  // Should have been set by stream_init().
     success = true;
   } else {
-    if (!uv_is_closing((uv_handle_t *)uv_stream)) {
-      uv_close((uv_handle_t *)uv_stream, connect_close_cb);
-    }
+    stream_may_close(&stream->s);
     // Wait for the close callback to arrive before retrying or returning, otherwise
     // it may lead to a hang or stack-use-after-return.
     LOOP_PROCESS_EVENTS_UNTIL(&main_loop, NULL, -1, closed);
@@ -352,6 +357,8 @@ tcp_retry:
   }
 
 cleanup:
+  stream->s.internal_close_cb = NULL;
+  stream->s.internal_data = NULL;
   xfree(addr);
   uv_freeaddrinfo(addr_req.addrinfo);
   return success;
diff --git a/src/nvim/event/stream.c b/src/nvim/event/stream.c
index 24bf8b5fe4..be3120a196 100644
--- a/src/nvim/event/stream.c
+++ b/src/nvim/event/stream.c
@@ -143,8 +143,8 @@ void stream_close_handle(Stream *stream)
 static void close_cb(uv_handle_t *handle)
 {
   Stream *stream = handle->data;
-  // Need to check if handle->data is NULL here as this callback may be called between
-  // the handle's initialization and stream_init() (e.g. in socket_connect()).
+  // Check if handle->data is NULL here, in case this callback is called between
+  // the handle's initialization and stream_init().
   if (stream && stream->close_cb) {
     stream->close_cb(stream, stream->close_cb_data);
   }
diff --git a/test/functional/core/channels_spec.lua b/test/functional/core/channels_spec.lua
index 3329e1c783..6470909fd9 100644
--- a/test/functional/core/channels_spec.lua
+++ b/test/functional/core/channels_spec.lua
@@ -505,6 +505,27 @@ describe('channels', function()
       )
     end)
   end)
+
+  it('no stack-buffer-overflow with Nvim exit during connection #39387', function()
+    local nvim0 = n.get_session()
+    -- Need a valid pipe so that connecting to it doesn't fail immediately.
+    local server = fn.serverstart()
+    finally(function()
+      nvim0:close()
+    end)
+
+    n.set_session(n.new_session(true))
+    exec_lua(function()
+      vim.defer_fn(function()
+        vim.uv.sleep(50) -- Block the uv event loop.
+        vim.fn.sockconnect('pipe', server)
+      end, 10)
+    end)
+    vim.uv.sleep(20)
+    -- The server uv event loop is currently blocked, so the channel close will
+    -- be processed when sockconnect() polls.
+    n.check_close()
+  end)
 end)
 
 describe('loopback', function()