1package lsp
2
3import (
4 "bufio"
5 "bytes"
6 "errors"
7 "fmt"
8 "io"
9 "strconv"
10 "strings"
11)
12
13const (
14 maxMessageBytes = 32 * 1024 * 1024
15 maxHeaderBytes = 8 * 1024
16 maxHeaderLines = 64
17)
18
19// WriteMessage writes one LSP Content-Length framed JSON-RPC message.
20func WriteMessage(w io.Writer, payload []byte) error {
21 if len(payload) > maxMessageBytes {
22 return fmt.Errorf("LSP message is too large: %d bytes", len(payload))
23 }
24 header := []byte(fmt.Sprintf("Content-Length: %d\r\n\r\n", len(payload)))
25 if n, err := w.Write(header); err != nil {
26 return err
27 } else if n != len(header) {
28 return io.ErrShortWrite
29 }
30 n, err := w.Write(payload)
31 if err == nil && n != len(payload) {
32 return io.ErrShortWrite
33 }
34 return err
35}
36
37// ReadMessage reads one Content-Length framed LSP message. It accepts other
38// headers (including Content-Type) and either CRLF or LF header endings.
39func ReadMessage(r *bufio.Reader) ([]byte, error) {
40 length := -1
41 headerBytes := 0
42 headerLines := 0
43 for {
44 line, err := r.ReadSlice('\n')
45 if err != nil {
46 if err == bufio.ErrBufferFull {
47 return nil, errors.New("LSP headers are too large")
48 }
49 return nil, err
50 }
51 headerBytes += len(line)
52 headerLines++
53 if headerBytes > maxHeaderBytes || headerLines > maxHeaderLines {
54 return nil, errors.New("LSP headers are too large")
55 }
56 lineText := strings.TrimRight(string(line), "\r\n")
57 if lineText == "" {
58 break
59 }
60 key, value, ok := strings.Cut(lineText, ":")
61 if !ok {
62 return nil, fmt.Errorf("malformed LSP header %q", lineText)
63 }
64 if strings.EqualFold(strings.TrimSpace(key), "content-length") {
65 n, err := strconv.Atoi(strings.TrimSpace(value))
66 if err != nil || n < 0 {
67 return nil, fmt.Errorf("invalid Content-Length %q", strings.TrimSpace(value))
68 }
69 length = n
70 }
71 }
72 if length < 0 {
73 return nil, errors.New("LSP message has no Content-Length")
74 }
75 if length > maxMessageBytes {
76 return nil, fmt.Errorf("LSP message is too large: %d bytes", length)
77 }
78 payload := make([]byte, length)
79 if _, err := io.ReadFull(r, payload); err != nil {
80 return nil, err
81 }
82 return payload, nil
83}
84
85// Frame is useful in tests and for callers which already have a complete
86// message. It returns the exact wire representation.
87func Frame(payload []byte) []byte {
88 var b bytes.Buffer
89 if err := WriteMessage(&b, payload); err != nil {
90 panic(err)
91 }
92 return b.Bytes()
93}