neovim/runtime/lua/vim/lsp/rpc.lua

local log = require('vim.lsp.log')
local protocol = require('vim.lsp.protocol')
local net_transport = require('vim.net._transport')
local strbuffer = require('vim._core.stringbuffer')
local validate = vim.validate

--- Embeds the given string into a table and correctly computes `Content-Length`.
---
--- @param message string
--- @return string message with `Content-Length` attribute
local function format_message_with_content_length(message)
  return table.concat({
    'Content-Length: ',
    tostring(#message),
    '\r\n\r\n',
    message,
  })
end

--- Extract `content-length` from the header.
---
--- The structure of header fields conforms to [HTTP semantics](https://tools.ietf.org/html/rfc7230#section-3.2),
--- i.e., `header-field = field-name : OWS field-value OWS`. OWS means optional whitespace (space/horizontal tabs).
---
--- We ignore lines ending with `\n` that don't contain `content-length`, since some servers
--- write log to standard output and there's no way to avoid it.
--- See https://github.com/neovim/neovim/pull/35743#pullrequestreview-3379705828
--- @param ptr vim._core.stringbuffer.ptr The ptr to buffer to parse
--- @param start integer The starting index of the buffer to parse, 0-based
--- @param len integer The length of the header to parse
--- @return integer
local function get_content_length(ptr, start, len)
  local state = 'name'
  local i, end_ = start, start + len
  local j, name = 1, 'content-length'
  local buf = strbuffer.new()
  local digit = true
  while i < end_ do
    local c = ptr[i]
    if state == 'name' then
      if c >= 65 and c <= 90 then -- lower case
        c = c + 32
      end
      if c == 10 then -- a blank line or a line with a prefix of header
        j = 1
      elseif (c == 32 or c == 9) and j == 1 then -- luacheck: ignore 542
        -- skip OWS for compatibility only
      elseif c == name:byte(j) then
        j = j + 1
      elseif c == 58 and j == 15 then
        state = 'colon'
      else
        state = 'invalid'
      end
    elseif state == 'colon' then
      if c ~= 32 and c ~= 9 then -- skip OWS normally
        state = 'value'
        i = i - 1
      end
    elseif state == 'value' then
      if c == 13 and ptr[i + 1] == 10 then -- must end with \r\n
        local value = buf:get()
        if digit then
          return vim._assert_integer(value)
        end
        error('value of Content-Length is not number: ' .. value)
      else
        buf:put(string.char(c))
      end
      if c < 48 and c ~= 32 and c ~= 9 or c > 57 then
        digit = false
      end
    elseif state == 'invalid' then
      if c == 10 then -- reset for next line
        state, j = 'name', 1
      end
    end
    i = i + 1
  end
  local header = strbuffer.new()
  for k = start, end_ - 1 do
    header:put(string.char(ptr[k]))
  end
  error('Content-Length not found in header: ' .. header:tostring())
end

local M = {}

--- Mapping of error codes used by the client
--- @enum vim.lsp.rpc.ClientErrors
local client_errors = {
  INVALID_SERVER_MESSAGE = 1,
  INVALID_SERVER_JSON = 2,
  NO_RESULT_CALLBACK_FOUND = 3,
  READ_ERROR = 4,
  NOTIFICATION_HANDLER_ERROR = 5,
  SERVER_REQUEST_HANDLER_ERROR = 6,
  SERVER_RESULT_CALLBACK_ERROR = 7,
}

--- @type table<string,integer> | table<integer,string>
--- @nodoc
M.client_errors = vim.deepcopy(client_errors)
for k, v in pairs(client_errors) do
  M.client_errors[v] = k
end

--- Constructs an error message from an LSP error object.
---
---@param err table The error object
---@return string error_message The formatted error message
function M.format_rpc_error(err)
  validate('err', err, 'table')

  -- There is ErrorCodes in the LSP specification,
  -- but in ResponseError.code it is not used and the actual type is number.
  local code --- @type string
  if protocol.ErrorCodes[err.code] then
    code = string.format('code_name = %s,', protocol.ErrorCodes[err.code])
  else
    code = string.format('code_name = unknown, code = %s,', err.code)
  end

  local message_parts = { 'RPC[Error]', code }
  if err.message then
    table.insert(message_parts, 'message =')
    table.insert(message_parts, string.format('%q', err.message))
  end
  if err.data then
    table.insert(message_parts, 'data =')
    table.insert(message_parts, vim.inspect(err.data))
  end
  return table.concat(message_parts, ' ')
end

--- Creates an RPC response table `error` to be sent to the LSP response.
---
---@param code integer RPC error code defined, see `vim.lsp.protocol.ErrorCodes`
---@param message? string arbitrary message to send to server
---@param data? any arbitrary data to send to server
---
---@see lsp.ErrorCodes See `vim.lsp.protocol.ErrorCodes`
---@return lsp.ResponseError
function M.rpc_response_error(code, message, data)
  -- TODO should this error or just pick a sane error (like InternalError)?
  ---@type string
  local code_name = assert(protocol.ErrorCodes[code], 'Invalid RPC error code')
  return {
    code = code,
    message = message or code_name,
    data = data,
  }
end

--- Dispatchers for LSP message types.
--- @class vim.lsp.rpc.Dispatchers
--- @inlinedoc
--- @field notification fun(method: vim.lsp.protocol.Method.ServerToClient, params: table)
--- @field server_request fun(method: vim.lsp.protocol.Method.ServerToClient, params: table): any?, lsp.ResponseError?
--- @field on_exit fun(code: integer, signal: integer)
--- @field on_error fun(code: integer, err: any)

--- @type vim.lsp.rpc.Dispatchers
local default_dispatchers = {
  --- Default dispatcher for notifications sent to an LSP server.
  ---
  ---@param method vim.lsp.protocol.Method.ServerToClient The invoked LSP method
  ---@param params table Parameters for the invoked LSP method
  notification = function(method, params)
    log.debug('notification', method, params)
  end,

  --- Default dispatcher for requests sent to an LSP server.
  ---
  ---@param method vim.lsp.protocol.Method.ServerToClient The invoked LSP method
  ---@param params table Parameters for the invoked LSP method
  ---@return any result (always nil for the default dispatchers)
  ---@return lsp.ResponseError error `vim.lsp.protocol.ErrorCodes.MethodNotFound`
  server_request = function(method, params)
    log.debug('server_request', method, params)
    return nil, M.rpc_response_error(protocol.ErrorCodes.MethodNotFound)
  end,

  --- Default dispatcher for when a client exits.
  ---
  ---@param code integer Exit code
  ---@param signal integer Number describing the signal used to terminate (if any)
  on_exit = function(code, signal)
    log.info('client_exit', { code = code, signal = signal })
  end,

  --- Default dispatcher for client errors.
  ---
  ---@param code integer Error code
  ---@param err any Details about the error
  on_error = function(code, err)
    log.error('client_error:', M.client_errors[code], err)
  end,
}

--- Parse one `Content-Length` framed message from `strbuf`.
---
--- Returns a body after consuming one full frame, returns nil if more bytes are needed.
--- Raises an error if the buffered data is not a valid frame.
---
---@param strbuf vim._core.stringbuffer
---@return string?
local function message_decoder(strbuf)
  local header_len ---@type integer?
  local ptr, len = strbuf:ref()
  for i = 0, len - 4 do
    -- Find the header boundary "\r\n\r\n"
    -- (compare bytes instead of string.find(), to avoid a string alloc).
    if ptr[i] == 13 and ptr[i + 1] == 10 and ptr[i + 2] == 13 and ptr[i + 3] == 10 then
      header_len = i + 2
      break
    end
  end

  if not header_len then
    return nil
  end

  local content_length = get_content_length(ptr, 0, header_len)
  if strbuffer.len(strbuf) < header_len + 2 + content_length then
    return nil
  end

  strbuf:skip(header_len + 2) -- skip past header boundary
  return strbuf:get(content_length)
end

--- @private
--- @param handle_body fun(body: string)
--- @param on_exit? fun()
--- @param on_error? fun(err: any, errkind: vim.lsp.rpc.ClientErrors)
function M.create_read_loop(handle_body, on_exit, on_error)
  on_exit = on_exit or function() end
  on_error = on_error or function() end
  local message_stream = net_transport.MessageStream.new(
    message_decoder,
    format_message_with_content_length,
    function(err, chunk)
      if err then
        on_error(err, M.client_errors.READ_ERROR)
      elseif chunk then
        handle_body(chunk)
      else
        on_exit()
      end
    end,
    function(err)
      on_error(err, M.client_errors.INVALID_SERVER_MESSAGE)
    end
  )

  return function(err, chunk)
    message_stream:feed(err, chunk)
  end
end

--- Client RPC object
--- @class vim.lsp.rpc.Client
--- @field private message_index integer
--- @field private message_callbacks table<integer, function> dict of message_id to callback
--- @field private notify_reply_callbacks table<integer, function> dict of message_id to callback
--- @field private transport vim.net.Transport
--- @field private message_stream vim.net.MessageStream
--- @field private dispatchers vim.lsp.rpc.Dispatchers
---
--- Sends a request to the LSP server.
--- @field request fun(method: vim.lsp.protocol.Method.ClientToServer.Request, params: table?, callback: fun(err?: lsp.ResponseError, result: any, request_id: integer), notify_reply_callback?: fun(message_id: integer)):boolean,integer?
---
--- Sends a notification to the LSP server.
--- @field notify fun(method: vim.lsp.protocol.Method.ClientToServer.Notification, params: any): boolean
---
--- Indicates if the RPC is closing.
--- @field is_closing fun(): boolean
---
--- Terminates the RPC client.
--- @field terminate fun()
local Client = {}

---@package
---@param dispatchers vim.lsp.rpc.Dispatchers
---@param transport vim.net.Transport
---@param decode fun(buf: vim._core.stringbuffer): string?
---@param format fun(msg: string): string
---@return vim.lsp.rpc.Client
function Client.new(dispatchers, transport, decode, format)
  local result = {
    message_index = 0,
    message_callbacks = {},
    notify_reply_callbacks = {},
    transport = transport,
    dispatchers = dispatchers,
  }

  ---@private
  function result.is_closing()
    return result.transport:is_closing()
  end

  ---@private
  function result.terminate()
    result.transport:terminate()
  end

  ---@private
  ---@param method (vim.lsp.protocol.Method.ClientToServer.Request) The invoked LSP method
  ---@param params (table?) Parameters for the invoked LSP method
  ---@param callback fun(err: lsp.ResponseError?, result: any) Callback to invoke
  ---@param notify_reply_callback? fun(message_id: integer) Callback to invoke as soon as a request is no longer pending
  ---@return boolean success `true` if request could be sent, `false` if not
  ---@return integer? message_id if request could be sent, `nil` if not
  function result.request(method, params, callback, notify_reply_callback)
    return result:_request(method, params, callback, notify_reply_callback)
  end

  ---@private
  ---@param method (vim.lsp.protocol.Method.ClientToServer.Notification) The invoked LSP method
  ---@param params (table?) Parameters for the invoked LSP method
  ---@return boolean `true` if notification could be sent, `false` if not
  function result.notify(method, params)
    return result:_notify(method, params)
  end

  ---@cast result vim.lsp.rpc.Client
  local self = setmetatable(result, { __index = Client })

  self.message_stream = net_transport.MessageStream.new(decode, format, function(err, data)
    if err then
      self:on_error(M.client_errors.READ_ERROR, err)
    elseif data then
      self:handle_body(data)
    else
      ---@diagnostic disable-next-line: invisible
      self.transport:terminate()
    end
  end, function(err)
    self:on_error(M.client_errors.INVALID_SERVER_MESSAGE, err)
    ---@diagnostic disable-next-line: invisible
    self.transport:terminate()
  end)

  transport:listen(function(err, data)
    ---@diagnostic disable-next-line: invisible
    self.message_stream:feed(err, data)
  end, dispatchers.on_exit)
  return self
end

---@private
function Client:encode_and_send(payload)
  log.debug('rpc.send', payload)
  if self.transport:is_closing() then
    return false
  end
  local jsonstr = vim.json.encode(payload)

  self.transport:write(self.message_stream.encode(jsonstr))
  return true
end

---@package
--- Sends a notification to the LSP server.
---@param method vim.lsp.protocol.Method.ClientToServer.Notification The invoked LSP method
---@param params any Parameters for the invoked LSP method
---@return boolean `true` if notification could be sent, `false` if not
function Client:_notify(method, params)
  return self:encode_and_send({
    jsonrpc = '2.0',
    method = method,
    params = params,
  })
end

---@private
--- sends an error object to the remote LSP process.
function Client:send_response(request_id, err, result)
  return self:encode_and_send({
    id = request_id,
    jsonrpc = '2.0',
    error = err,
    result = result,
  })
end

---@package
--- Sends a request to the LSP server and runs {callback} upon response.
---
---@param method vim.lsp.protocol.Method.ClientToServer.Request The invoked LSP method
---@param params table? Parameters for the invoked LSP method
---@param callback fun(err?: lsp.ResponseError, result: any, message_id: integer) Callback to invoke
---@param notify_reply_callback? fun(message_id: integer) Callback to invoke as soon as a request is no longer pending
---@return boolean success `true` if request could be sent, `false` if not
---@return integer? message_id if request could be sent, `nil` if not
function Client:_request(method, params, callback, notify_reply_callback)
  validate('callback', callback, 'function')
  validate('notify_reply_callback', notify_reply_callback, 'function', true)
  self.message_index = self.message_index + 1
  local message_id = self.message_index
  local result = self:encode_and_send({
    id = message_id,
    jsonrpc = '2.0',
    method = method,
    params = params,
  })

  if not result then
    return false
  end

  self.message_callbacks[message_id] = vim.schedule_wrap(callback)
  if notify_reply_callback then
    self.notify_reply_callbacks[message_id] = vim.schedule_wrap(notify_reply_callback)
  end
  return result, message_id
end

---@package
---@param errkind vim.lsp.rpc.ClientErrors
---@param err any
function Client:on_error(errkind, err)
  assert(M.client_errors[errkind])
  -- TODO what to do if this fails?
  pcall(self.dispatchers.on_error, errkind, err)
end

-- TODO periodically check message_callbacks for old requests past a certain
-- time and log them. This would require storing the timestamp. I could call
-- them with an error then, perhaps.

--- @package
--- @param body string
function Client:handle_body(body)
  local ok, decoded = pcall(vim.json.decode, body)
  if not ok then
    self:on_error(M.client_errors.INVALID_SERVER_JSON, decoded)
    return
  elseif type(decoded) ~= 'table' then
    self:on_error(M.client_errors.INVALID_SERVER_MESSAGE, decoded)
    return
  end

  log.debug('rpc.receive', decoded)

  if
    -- Received a request.
    type(decoded.method) == 'string' and decoded.id
  then
    if type(decoded.id) ~= 'number' and type(decoded.id) ~= 'string' and decoded.id ~= vim.NIL then
      log.error(
        'Server request id must be a number or string, got ' .. type(decoded.id),
        decoded.method,
        decoded.id
      )
      self:on_error(M.client_errors.INVALID_SERVER_MESSAGE, decoded)
      return
    end

    -- Schedule here so that the users functions don't trigger an error and
    -- we can still use the result.
    vim.schedule(coroutine.wrap(function()
      xpcall(function()
        local result, err = self.dispatchers.server_request(decoded.method, decoded.params)
        log.debug('server_request: callback result', { result = result, err = err })
        if result == nil and err == nil then
          error(
            string.format(
              'method %q: either a result or an error must be sent to the server in response',
              decoded.method
            )
          )
        end
        if err then
          assert(
            type(err) == 'table',
            'err must be a table. Use rpc_response_error to help format errors.'
          )
          assert(
            protocol.ErrorCodes[err.code],
            'Errors must use protocol.ErrorCodes. Use rpc_response_error to help format errors.'
          )
        end
        self:send_response(decoded.id, err, result)
      end, function(err)
        self:on_error(M.client_errors.SERVER_REQUEST_HANDLER_ERROR, err)
        self:send_response(
          decoded.id,
          M.rpc_response_error(protocol.ErrorCodes.InternalError, err),
          nil
        )
      end)
    end))
  elseif
    -- Received a response to a request we sent.
    decoded.id
  then
    -- If there was an error in detecting the id in the Request object
    -- (e.g. Parse error/Invalid Request), it must be Null.
    if decoded.id == vim.NIL then
      log.warn('Server sent response with null id', decoded)
      self:on_error(M.client_errors.INVALID_SERVER_MESSAGE, decoded)
      return
    end
    -- Proceed only if exactly one of 'result' or 'error' is present,
    -- as required by the JSON-RPC spec:
    -- * If 'error' is nil, then 'result' must be present.
    -- * If 'result' is nil, then 'error' must be present (and not vim.NIL).
    if (decoded.error == nil or decoded.error == vim.NIL) and decoded.result == nil then
      log.error('Server respond empty result and error', decoded)
      self:on_error(M.client_errors.INVALID_SERVER_MESSAGE, decoded)
      return
    end

    -- We sent a number, so we expect a number.
    local result_id = vim._assert_integer(decoded.id)

    -- Notify the user that a response was received for the request
    local notify_reply_callback = self.notify_reply_callbacks[result_id]
    if notify_reply_callback then
      validate('notify_reply_callback', notify_reply_callback, 'function')
      notify_reply_callback(result_id)
      self.notify_reply_callbacks[result_id] = nil
    end

    -- Do not surface RequestCancelled to users, it is RPC-internal.
    if decoded.error then
      assert(type(decoded.error) == 'table')
      if decoded.error.code == protocol.ErrorCodes.RequestCancelled then
        log.debug('Received cancellation ack', decoded)
        -- Clear any callback since this is cancelled now.
        -- This is safe to do assuming that these conditions hold:
        -- - The server will not send a result callback after this cancellation.
        -- - If the server sent this cancellation ACK after sending the result, the user of this RPC
        -- client will ignore the result themselves.
        if result_id then
          self.message_callbacks[result_id] = nil
        end
        return
      end
    end

    local callback = self.message_callbacks[result_id]
    if callback then
      self.message_callbacks[result_id] = nil
      validate('callback', callback, 'function')
      xpcall(function()
        callback(decoded.error, decoded.result ~= vim.NIL and decoded.result or nil, result_id)
      end, function(err)
        self:on_error(M.client_errors.SERVER_RESULT_CALLBACK_ERROR, err)
      end)
    else
      self:on_error(M.client_errors.NO_RESULT_CALLBACK_FOUND, decoded)
      log.error('No callback found for server response id ' .. result_id)
    end
  elseif
    -- Received a notification.
    type(decoded.method) == 'string'
  then
    xpcall(function()
      assert(
        self.dispatchers.notification(decoded.method, decoded.params) == nil,
        'notification handlers should not return a value'
      )
    end, function(err)
      self:on_error(M.client_errors.NOTIFICATION_HANDLER_ERROR, err)
    end)
  else
    -- Invalid server message
    self:on_error(M.client_errors.INVALID_SERVER_MESSAGE, decoded)
  end
end

---@param dispatchers vim.lsp.rpc.Dispatchers?
---@return vim.lsp.rpc.Dispatchers
local function merge_dispatchers(dispatchers)
  if not dispatchers then
    return default_dispatchers
  end
  ---@diagnostic disable-next-line: no-unknown
  for name, fn in pairs(dispatchers) do
    if type(fn) ~= 'function' then
      error(string.format('dispatcher.%s must be a function', name))
    end
  end
  ---@type vim.lsp.rpc.Dispatchers
  local merged = {
    notification = (
      dispatchers.notification and vim.schedule_wrap(dispatchers.notification)
      or default_dispatchers.notification
    ),
    on_error = (
      dispatchers.on_error and vim.schedule_wrap(dispatchers.on_error)
      or default_dispatchers.on_error
    ),
    on_exit = dispatchers.on_exit or default_dispatchers.on_exit,
    server_request = dispatchers.server_request or default_dispatchers.server_request,
  }
  return merged
end

--- Create a LSP RPC client factory that connects to either:
---
---  - a named pipe (windows)
---  - a domain socket (unix)
---  - a host and port via TCP
---
--- Return a function that can be passed to the `cmd` field for
--- |vim.lsp.start()|.
---
---@param host_or_path string host to connect to or path to a pipe/domain socket
---@param port integer? TCP port to connect to. If absent the first argument must be a pipe
---@return fun(dispatchers: vim.lsp.rpc.Dispatchers): vim.lsp.rpc.Client
function M.connect(host_or_path, port)
  log.info('Connecting RPC client', { host_or_path = host_or_path, port = port })

  validate('host_or_path', host_or_path, 'string')
  validate('port', port, 'number', true)

  return function(dispatchers)
    validate('dispatchers', dispatchers, 'table', true)

    dispatchers = merge_dispatchers(dispatchers)

    local transport = net_transport.TransportConnect.new(host_or_path, port, vim.lsp.log._self)
    return Client.new(dispatchers, transport, message_decoder, format_message_with_content_length)
  end
end

--- Additional context for the LSP server process.
--- @class vim.net.transport.ExtraSpawnParams
--- @inlinedoc
--- @field cwd? string Working directory for the LSP server process
--- @field detached? boolean Detach the LSP server process from the current process
--- @field env? table<string,string> Additional environment variables for LSP server process. See |vim.system()|

--- Starts an LSP server process and create an LSP RPC client object to
--- interact with it. Communication with the spawned process happens via stdio. For
--- communication via TCP, spawn a process manually and use |vim.lsp.rpc.connect()|
---
--- @param cmd string[] Command to start the LSP server.
--- @param dispatchers? vim.lsp.rpc.Dispatchers
--- @param extra_spawn_params? vim.net.transport.ExtraSpawnParams
--- @return vim.lsp.rpc.Client
function M.start(cmd, dispatchers, extra_spawn_params)
  log.info('Starting RPC client', { cmd = cmd, extra = extra_spawn_params })

  validate('cmd', cmd, 'table')
  validate('dispatchers', dispatchers, 'table', true)

  dispatchers = merge_dispatchers(dispatchers)

  local transport = net_transport.TransportRun.new(cmd, extra_spawn_params, vim.lsp.log._self)
  return Client.new(dispatchers, transport, message_decoder, format_message_with_content_length)
end

return M