1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102

// Package lsp is a small Language Server Protocol client: enough of it to ask
// a language server for completions, hovers, definitions and diagnostics.
//
// The protocol is JSON-RPC 2.0 inside HTTP-style frames. The JSON-RPC half is
// jsonrpc's, shared with every other protocol built on it; what is here is the
// framing, the methods, and the conversation a language server expects.
package lsp

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"net/textproto"
	"strconv"
)

// MaxMessageBytes caps how large a single message may be. A language server
// should never send anything near this; the limit is there so that a corrupt
// or hostile Content-Length cannot make the editor allocate without bound.
const MaxMessageBytes = 64 << 20

// ErrMessageTooLarge is returned when a frame announces more bytes than
// MaxMessageBytes.
var ErrMessageTooLarge = errors.New("lsp: message too large")

// contentLengthHeader is the one header the protocol actually requires.
const contentLengthHeader = "Content-Length"

// Framing is the Language Server Protocol's way of marking one message off
// from the next: a Content-Length header, a blank line, then the JSON.
//
// It is the jsonrpc.Framer a language server connection is built with.
//
//	conn := jsonrpc.NewConn(stream, lsp.Framing{}, onNotify, onRequest)
type Framing struct{}

// ReadMessage reads one framed message and returns its body.
func (Framing) ReadMessage(r *bufio.Reader) ([]byte, error) { return ReadMessage(r) }

// WriteMessage frames body and writes it out.
func (Framing) WriteMessage(w io.Writer, body []byte) error { return WriteMessage(w, body) }

// WriteMessage frames body and writes it out.
//
// The frame is the header, a blank line, then the JSON — the same shape HTTP
// uses, which is what the specification asks for.
func WriteMessage(w io.Writer, body []byte) error {
	if _, err := fmt.Fprintf(w, "%s: %d\r\n\r\n", contentLengthHeader, len(body)); err != nil {
		return fmt.Errorf("lsp: writing the header: %w", err)
	}
	if _, err := w.Write(body); err != nil {
		return fmt.Errorf("lsp: writing the body: %w", err)
	}
	return nil
}

// ReadMessage reads one framed message and returns its body.
//
// It returns io.EOF when the stream ends cleanly between messages, which is
// how a server shutting down is told apart from one that died mid-message.
func ReadMessage(r *bufio.Reader) ([]byte, error) {
	length, err := readContentLength(r)
	if err != nil {
		return nil, err
	}

	body := make([]byte, length)
	if _, err := io.ReadFull(r, body); err != nil {
		return nil, fmt.Errorf("lsp: reading a %d-byte body: %w", length, err)
	}
	return body, nil
}

// readContentLength reads the header block and returns the announced size.
func readContentLength(r *bufio.Reader) (int, error) {
	header, err := textproto.NewReader(r).ReadMIMEHeader()
	if err != nil {
		if errors.Is(err, io.EOF) {
			return 0, io.EOF
		}
		return 0, fmt.Errorf("lsp: reading the header: %w", err)
	}
	return parseContentLength(header.Get(contentLengthHeader))
}

// parseContentLength turns the header's value into a size, rejecting the ones
// that are missing, malformed, or large enough to be an attack.
func parseContentLength(raw string) (int, error) {
	if raw == "" {
		return 0, fmt.Errorf("lsp: the frame has no %s header", contentLengthHeader)
	}

	length, err := strconv.Atoi(raw)
	if err != nil || length < 0 {
		return 0, fmt.Errorf("lsp: %s is %q, which is not a size", contentLengthHeader, raw)
	}
	if length > MaxMessageBytes {
		return 0, fmt.Errorf("%w: %d bytes", ErrMessageTooLarge, length)
	}
	return length, nil
}