LSP-server-example

1
package server
2

3
import (
4
	"bufio"
5
	"encoding/json"
6
	"fmt"
7
	"io"
8
	"log/slog"
9
	"net"
10
	"net/textproto"
11
	"strconv"
12
	"sync"
13

14
	"lsp-server/internal/server/handlers"
15
)
16

17
type Mux struct {
18
	concurrencyLimit     int64
19
	requestHandlers      map[string]handlers.Request
20
	notificationHandlers map[string]handlers.Notification
21
	writeLock            *sync.Mutex
22
	conn                 *net.Conn
23
}
24

25
const concurrencyLimit = 4
26

27
func NewMux() *Mux {
28
	return &Mux{
29
		concurrencyLimit:     concurrencyLimit,
30
		requestHandlers:      map[string]handlers.Request{},
31
		notificationHandlers: map[string]handlers.Notification{},
32
		writeLock:            &sync.Mutex{},
33
	}
34
}
35

36
func (m *Mux) Process() error {
37
	for {
38
		req, err := Read(bufio.NewReader(*m.conn))
39
		if err != nil {
40
			return err
41
		}
42
		if req.IsNotification() {
43
			if req.Method != "exit" {
44
				// Drop notifications sent before initialization.
45
				slog.Warn("dropping notification sent before initialization", slog.Any("req", req))
46
				continue
47
			}
48
			m.handleMessage(req)
49
			continue
50
		}
51
		if req.Method != "initialize" {
52
			// Return an error if methods used before initialization.
53
			slog.Warn("the client sent a method before initialization", slog.Any("req", req))
54
			if err = m.write(NewResponseError(req.ID, ErrServerNotInitialized)); err != nil {
55
				return err
56
			}
57
			continue
58
		}
59
		m.handleMessage(req)
60
		break
61
	}
62
	slog.Info("initialization complete")
63

64
	// Handle standard flow.
65
	sem := make(chan struct{}, m.concurrencyLimit)
66
	for {
67
		sem <- struct{}{}
68
		req, err := Read(bufio.NewReader(*m.conn))
69
		if err != nil {
70
			return err
71
		}
72
		go func(req Request) {
73
			m.handleMessage(req)
74
			<-sem
75
		}(req)
76
	}
77
}
78

79
func (m *Mux) handleMessage(req Request) {
80
	if req.IsNotification() {
81
		m.handleNotification(req)
82
		return
83
	}
84
	m.handleRequest(req)
85
}
86

87
func (m *Mux) handleNotification(req Request) {
88
	log := slog.With(slog.String("method", req.Method))
89
	nh, ok := m.notificationHandlers[req.Method]
90
	if !ok {
91
		log.Warn("notification not handled")
92
		return
93
	}
94
	// We don't need to notify clients if the notification results in an error.
95
	if err := nh.Call(req.Params); err != nil {
96
		log.Error("failed to handle notification", slog.Any("error", err))
97
	}
98
}
99

100
func (m *Mux) handleRequest(req Request) {
101
	log := slog.With(slog.Any("id", req.ID), slog.String("method", req.Method))
102
	mh, ok := m.requestHandlers[req.Method]
103
	if !ok {
104
		log.Error("method not found")
105
		if err := m.write(NewResponseError(req.ID, ErrMethodNotFound)); err != nil {
106
			log.Error("failed to respond", slog.Any("error", err))
107
		}
108
		return
109
	}
110
	var res Response
111
	result, err := mh.Call(req.Params)
112
	if err != nil {
113
		log.Error("failed to handle", slog.Any("error", err))
114
		res = NewResponseError(req.ID, err)
115
	} else {
116
		res = NewResponse(req.ID, result)
117
	}
118
	if err = m.write(res); err != nil {
119
		log.Error("failed to respond", slog.Any("error", err))
120
	}
121
}
122

123
func (m *Mux) write(msg Message) error {
124
	m.writeLock.Lock()
125
	defer m.writeLock.Unlock()
126
	return Write(bufio.NewWriter(*m.conn), msg)
127
}
128

129
func Read(r *bufio.Reader) (Request, error) {
130
	// Read header.
131
	header, err := textproto.NewReader(r).ReadMIMEHeader()
132
	if err != nil {
133
		return Request{}, err
134
	}
135
	var req Request
136
	contentLength, err := strconv.ParseInt(header.Get("Content-Length"), 10, 64)
137
	if err != nil {
138
		return Request{}, ErrInvalidContentLengthHeader
139
	}
140
	// Read body.
141
	err = json.NewDecoder(io.LimitReader(r, contentLength)).Decode(&req)
142
	if err != nil {
143
		return Request{}, err
144
	}
145
	if !req.IsJSONRPC() {
146
		return req, ErrInvalidRequest
147
	}
148
	return req, nil
149
}
150

151
func Write(w *bufio.Writer, msg Message) error {
152
	// Calculate body size.
153
	body, err := json.Marshal(msg)
154
	if err != nil {
155
		return err
156
	}
157
	// Write the header.
158
	_, err = w.WriteString(fmt.Sprintf("Content-Length: %d\r\n\r\n", len(body)))
159
	if err != nil {
160
		return err
161
	}
162
	// Write the body.
163
	_, err = w.Write(body)
164
	if err != nil {
165
		return err
166
	}
167
	// Flush.
168
	err = w.Flush()
169
	return err
170
}
171
func (m *Mux) Notify(method string, params any) error {
172
	n := Notification{
173
		ProtocolVersion: protocolVersion,
174
		Method:          method,
175
		Params:          params,
176
	}
177
	return m.write(n)
178
}
179
func (m *Mux) HandleRequest(name string, request handlers.Request) {
180
	m.requestHandlers[name] = request
181
}
182

183
func (m *Mux) HandleNotification(name string, notification handlers.Notification) {
184
	m.notificationHandlers[name] = notification
185
}
186

187
func (m *Mux) NewConn(conn *net.Conn) {
188
	m.conn = conn
189
}
190