channels: implement sockopen() to connect to socket

Helped-By: oni-link <knil.ino@gmail.com>

4 files changed, 142 insertions, 0 deletions

diff --git a/src/nvim/eval.c b/src/nvim/eval.c
index 7bfd638e86..4e303414a3 100644
--- a/src/nvim/eval.c
+++ b/src/nvim/eval.c
@@ -15058,6 +15058,54 @@ static void f_simplify(typval_T *argvars, typval_T *rettv, FunPtr fptr)
   rettv->v_type = VAR_STRING;
 }
 
+/// "sockconnect()" function
+static void f_sockconnect(typval_T *argvars, typval_T *rettv, FunPtr fptr)
+{
+  if (argvars[0].v_type != VAR_STRING || argvars[1].v_type != VAR_STRING) {
+    EMSG(_(e_invarg));
+    return;
+  }
+  if (argvars[2].v_type != VAR_DICT && argvars[2].v_type != VAR_UNKNOWN) {
+    // Wrong argument types
+    EMSG2(_(e_invarg2), "expected dictionary");
+    return;
+  }
+
+  const char *mode = tv_get_string(&argvars[0]);
+  const char *address = tv_get_string(&argvars[1]);
+
+  bool tcp;
+  if (strcmp(mode, "tcp") == 0) {
+    tcp = true;
+  } else if (strcmp(mode, "pipe") == 0) {
+    tcp = false;
+  } else {
+    EMSG2(_(e_invarg2), "invalid mode");
+    return;
+  }
+
+  bool rpc = false;
+  if (argvars[2].v_type == VAR_DICT) {
+    dict_T *opts = argvars[2].vval.v_dict;
+    rpc = tv_dict_get_number(opts, "rpc") != 0;
+  }
+
+  if (!rpc) {
+    EMSG2(_(e_invarg2), "rpc option must be true");
+    return;
+  }
+
+  const char *error = NULL;
+  uint64_t id = channel_connect(tcp, address, 50, &error);
+
+  if (error) {
+    EMSG2(_("connection failed: %s"), error);
+  }
+
+  rettv->vval.v_number = (varnumber_T)id;
+  rettv->v_type = VAR_NUMBER;
+}
+
 /// struct used in the array that's given to qsort()
 typedef struct {
   listitem_T *item;
diff --git a/src/nvim/eval.lua b/src/nvim/eval.lua
index 533403b2b0..334e10eb6c 100644
--- a/src/nvim/eval.lua
+++ b/src/nvim/eval.lua
@@ -268,6 +268,7 @@ return {
     simplify={args=1},
     sin={args=1, func="float_op_wrapper", data="&sin"},
     sinh={args=1, func="float_op_wrapper", data="&sinh"},
+    sockconnect={args={2,3}},
     sort={args={1, 3}},
     soundfold={args=1},
     spellbadword={args={0, 1}},
diff --git a/src/nvim/event/socket.c b/src/nvim/event/socket.c
index bc5a4ec75e..30a71a5586 100644
--- a/src/nvim/event/socket.c
+++ b/src/nvim/event/socket.c
@@ -15,6 +15,7 @@
 #include "nvim/vim.h"
 #include "nvim/strings.h"
 #include "nvim/path.h"
+#include "nvim/main.h"
 #include "nvim/memory.h"
 #include "nvim/macros.h"
 #include "nvim/charset.h"
@@ -189,3 +190,76 @@ static void close_cb(uv_handle_t *handle)
     watcher->close_cb(watcher, watcher->data);
   }
 }
+
+static void connect_cb(uv_connect_t *req, int status)
+{
+  int *ret_status = req->data;
+  *ret_status = status;
+  if (status != 0) {
+    uv_close((uv_handle_t *)req->handle, NULL);
+  }
+}
+
+bool socket_connect(Loop *loop, Stream *stream,
+                    bool is_tcp, const char *address,
+                    int timeout, const char **error)
+{
+  bool success = false;
+  int status;
+  uv_connect_t req;
+  req.data = &status;
+  uv_stream_t *uv_stream;
+
+  uv_tcp_t *tcp = &stream->uv.tcp;
+  uv_getaddrinfo_t addr_req;
+  addr_req.addrinfo = NULL;
+  const struct addrinfo *addrinfo = NULL;
+  char *addr = NULL;
+  if (is_tcp) {
+    addr = xstrdup(address);
+    char *host_end = strrchr(addr, ':');
+    if (!host_end) {
+      *error = _("tcp address must be host:port");
+      goto cleanup;
+    }
+    *host_end = NUL;
+
+    const struct addrinfo hints = { .ai_family = AF_UNSPEC,
+                                    .ai_socktype = SOCK_STREAM,
+                                    .ai_flags  = AI_NUMERICSERV };
+    int retval = uv_getaddrinfo(&loop->uv, &addr_req, NULL,
+                                addr, host_end+1, &hints);
+    if (retval != 0) {
+      *error = _("failed to lookup host or port");
+      goto cleanup;
+    }
+    addrinfo = addr_req.addrinfo;
+
+tcp_retry:
+    uv_tcp_init(&loop->uv, tcp);
+    uv_tcp_connect(&req,  tcp, addrinfo->ai_addr, connect_cb);
+    uv_stream = (uv_stream_t *)tcp;
+
+  } else {
+    uv_pipe_t *pipe = &stream->uv.pipe;
+    uv_pipe_init(&loop->uv, pipe, 0);
+    uv_pipe_connect(&req,  pipe, address, connect_cb);
+    uv_stream = (uv_stream_t *)pipe;
+  }
+  status = 1;
+  LOOP_PROCESS_EVENTS_UNTIL(&main_loop, NULL, timeout, status != 1);
+  if (status == 0) {
+    stream_init(NULL, stream, -1, uv_stream);
+    success = true;
+  } else if (is_tcp && addrinfo->ai_next) {
+    addrinfo = addrinfo->ai_next;
+    goto tcp_retry;
+  } else {
+    *error = _("connection refused");
+  }
+
+cleanup:
+  xfree(addr);
+  uv_freeaddrinfo(addr_req.addrinfo);
+  return success;
+}
diff --git a/src/nvim/msgpack_rpc/channel.c b/src/nvim/msgpack_rpc/channel.c
index cd64e14976..0ff749649b 100644
--- a/src/nvim/msgpack_rpc/channel.c
+++ b/src/nvim/msgpack_rpc/channel.c
@@ -145,6 +145,25 @@ void channel_from_connection(SocketWatcher *watcher)
   rstream_start(&channel->data.stream, parse_msgpack, channel);
 }
 
+uint64_t channel_connect(bool tcp, const char *address,
+                         int timeout, const char **error)
+{
+  Channel *channel = register_channel(kChannelTypeSocket, 0, NULL);
+  if (!socket_connect(&main_loop, &channel->data.stream,
+                      tcp, address, timeout, error)) {
+    decref(channel);
+    return 0;
+  }
+
+  incref(channel);  // close channel only after the stream is closed
+  channel->data.stream.internal_close_cb = close_cb;
+  channel->data.stream.internal_data = channel;
+  wstream_init(&channel->data.stream, 0);
+  rstream_init(&channel->data.stream, CHANNEL_BUFFER_SIZE);
+  rstream_start(&channel->data.stream, parse_msgpack, channel);
+  return channel->id;
+}
+
 /// Sends event/arguments to channel
 ///
 /// @param id The channel id. If 0, the event will be sent to all