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
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253

package jsonrpc

import (
	"bufio"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"sync"
)

// Conn is a JSON-RPC 2.0 connection over a framed stream.
//
// It is safe for concurrent use: several goroutines may call and notify at
// once, and replies are matched back to their callers by request id.
type Conn struct {
	stream io.ReadWriteCloser
	reader *bufio.Reader
	framer Framer

	onNotify  NotificationFunc
	onRequest RequestFunc

	writeMu sync.Mutex

	mu      sync.Mutex
	nextID  int64
	pending map[int64]chan *message
	closed  bool
}

// NewConn returns a connection over stream, using framer to mark one message
// off from the next. Call Run in a goroutine of its own to start reading.
//
// onNotify and onRequest may both be nil, in which case notifications are
// dropped and the peer's requests are answered with "method not found".
//
//	conn := jsonrpc.NewConn(stream, lsp.Framing{}, client.onNotify, client.onRequest)
//	go conn.Run()
func NewConn(stream io.ReadWriteCloser, framer Framer, onNotify NotificationFunc, onRequest RequestFunc) *Conn {
	return &Conn{
		stream:    stream,
		reader:    bufio.NewReader(stream),
		framer:    framer,
		onNotify:  onNotify,
		onRequest: onRequest,
		pending:   map[int64]chan *message{},
	}
}

// Run reads messages until the stream ends or the connection is closed, then
// fails every call still waiting for an answer.
//
// It returns nil for a clean end of stream and the read error otherwise.
func (c *Conn) Run() error {
	err := c.readLoop()
	c.failPending()
	return err
}

// readLoop dispatches every message that arrives.
func (c *Conn) readLoop() error {
	for {
		body, err := c.framer.ReadMessage(c.reader)
		if err != nil {
			if errors.Is(err, io.EOF) || c.isClosed() {
				return nil
			}
			return err
		}

		var msg message
		if err := json.Unmarshal(body, &msg); err != nil {
			return fmt.Errorf("jsonrpc: unreadable message: %w", err)
		}
		c.dispatch(&msg)
	}
}

// dispatch sends a message where it belongs: to the call that is waiting for
// it, to the notification handler, or to the request handler.
//
// A notification and a request differ only by the id: the peer expects an
// answer to exactly the messages that carry one.
func (c *Conn) dispatch(msg *message) {
	switch {
	case msg.Method != "" && len(msg.ID) == 0:
		if c.onNotify != nil {
			c.onNotify(msg.Method, msg.Params)
		}
	case msg.Method != "":
		c.answer(msg)
	default:
		c.deliver(msg)
	}
}

// deliver hands a response to the call that is waiting for it.
//
// The id is looked up among the ids *this* connection sent. A peer numbering
// its own requests from zero — which the Agent Client Protocol's agents do —
// therefore cannot collide with a call in flight here.
func (c *Conn) deliver(msg *message) {
	var id int64
	if err := json.Unmarshal(msg.ID, &id); err != nil {
		return // a response to an id we never sent; nothing to do with it
	}

	c.mu.Lock()
	waiting, ok := c.pending[id]
	delete(c.pending, id)
	c.mu.Unlock()

	if ok {
		waiting <- msg
	}
}

// answer hands a request the peer made to the handler, which replies now or
// later.
func (c *Conn) answer(msg *message) {
	request := &Request{Method: msg.Method, Params: msg.Params, conn: c, id: msg.ID}

	if c.onRequest == nil {
		request.Reply(nil, &ResponseError{Code: CodeMethodNotFound, Message: msg.Method})
		return
	}
	c.onRequest(request)
}

// Call sends a request and waits for its answer, decoding it into result.
//
// Passing a nil result discards the answer. The call ends early if ctx does,
// which is what stops a slow request from freezing the caller.
//
//	var items CompletionList
//	err := conn.Call(ctx, "textDocument/completion", params, &items)
func (c *Conn) Call(ctx context.Context, method string, params, result any) error {
	id, waiting, err := c.register()
	if err != nil {
		return err
	}
	defer c.forget(id)

	if err := c.sendRequest(id, method, params); err != nil {
		return err
	}

	select {
	case <-ctx.Done():
		return ctx.Err()
	case reply := <-waiting:
		return decodeReply(reply, result)
	}
}

// Notify sends a notification, which expects no answer.
//
//	err := conn.Notify("session/cancel", map[string]string{"sessionId": id})
func (c *Conn) Notify(method string, params any) error {
	encodedParams, err := encodeParams(params)
	if err != nil {
		return err
	}
	return c.write(&message{JSONRPC: Version, Method: method, Params: encodedParams})
}

// sendRequest encodes and writes one request.
func (c *Conn) sendRequest(id int64, method string, params any) error {
	encodedID, err := json.Marshal(id)
	if err != nil {
		return fmt.Errorf("jsonrpc: encoding the request id: %w", err)
	}
	encodedParams, err := encodeParams(params)
	if err != nil {
		return err
	}
	return c.write(&message{JSONRPC: Version, ID: encodedID, Method: method, Params: encodedParams})
}

// register allocates a request id and the channel its answer will arrive on.
func (c *Conn) register() (int64, chan *message, error) {
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.closed {
		return 0, nil, ErrClosed
	}
	c.nextID++
	// The channel is buffered so that the read loop never blocks handing over
	// an answer whose caller has already given up.
	waiting := make(chan *message, 1)
	c.pending[c.nextID] = waiting
	return c.nextID, waiting, nil
}

// forget drops a pending request, whether it was answered or abandoned.
func (c *Conn) forget(id int64) {
	c.mu.Lock()
	delete(c.pending, id)
	c.mu.Unlock()
}

// write frames and sends one message.
func (c *Conn) write(msg *message) error {
	body, err := json.Marshal(msg)
	if err != nil {
		return fmt.Errorf("jsonrpc: encoding a message: %w", err)
	}

	c.writeMu.Lock()
	defer c.writeMu.Unlock()
	if c.isClosed() {
		return ErrClosed
	}
	return c.framer.WriteMessage(c.stream, body)
}

// failPending wakes every waiting call once the connection has ended.
//
// Requests the *peer* sent are not answered here. There is no answer to give,
// and the peer whose connection has just ended is not waiting for one.
func (c *Conn) failPending() {
	c.mu.Lock()
	defer c.mu.Unlock()

	c.closed = true
	for id, waiting := range c.pending {
		close(waiting)
		delete(c.pending, id)
	}
}

// isClosed reports whether the connection has been closed.
func (c *Conn) isClosed() bool {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.closed
}

// Close shuts the connection down. Calls still waiting fail with ErrClosed.
func (c *Conn) Close() error {
	c.mu.Lock()
	if c.closed {
		c.mu.Unlock()
		return nil
	}
	c.closed = true
	c.mu.Unlock()

	return c.stream.Close()
}