aboutsummaryrefslogtreecommitdiff
path: root/internal/lsp/lsprpc/middleware.go
blob: 2ee83a20301e5885421808d6be030f838d7be083 (plain)
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
// Copyright 2021 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package lsprpc

import (
	"context"
	"encoding/json"
	"sync"

	"golang.org/x/tools/internal/event"
	jsonrpc2_v2 "golang.org/x/tools/internal/jsonrpc2_v2"
	"golang.org/x/xerrors"
)

// Metadata holds arbitrary data transferred between jsonrpc2 peers.
type Metadata map[string]interface{}

// PeerInfo holds information about a peering between jsonrpc2 servers.
type PeerInfo struct {
	// RemoteID is the identity of the current server on its peer.
	RemoteID int64

	// LocalID is the identity of the peer on the server.
	LocalID int64

	// IsClient reports whether the peer is a client. If false, the peer is a
	// server.
	IsClient bool

	// Metadata holds arbitrary information provided by the peer.
	Metadata Metadata
}

// Handshaker handles both server and client handshaking over jsonrpc2. To
// instrument server-side handshaking, use Handshaker.Middleware. To instrument
// client-side handshaking, call Handshaker.ClientHandshake for any new
// client-side connections.
type Handshaker struct {
	// Metadata will be shared with peers via handshaking.
	Metadata Metadata

	mu     sync.Mutex
	prevID int64
	peers  map[int64]PeerInfo
}

// Peers returns the peer info this handshaker knows about by way of either the
// server-side handshake middleware, or client-side handshakes.
func (h *Handshaker) Peers() []PeerInfo {
	h.mu.Lock()
	defer h.mu.Unlock()

	var c []PeerInfo
	for _, v := range h.peers {
		c = append(c, v)
	}
	return c
}

// Middleware is a jsonrpc2 middleware function to augment connection binding
// to handle the handshake method, and record disconnections.
func (h *Handshaker) Middleware(inner jsonrpc2_v2.Binder) jsonrpc2_v2.Binder {
	return BinderFunc(func(ctx context.Context, conn *jsonrpc2_v2.Connection) (jsonrpc2_v2.ConnectionOptions, error) {
		opts, err := inner.Bind(ctx, conn)
		if err != nil {
			return opts, err
		}

		localID := h.nextID()
		info := &PeerInfo{
			RemoteID: localID,
			Metadata: h.Metadata,
		}

		// Wrap the delegated handler to accept the handshake.
		delegate := opts.Handler
		opts.Handler = jsonrpc2_v2.HandlerFunc(func(ctx context.Context, req *jsonrpc2_v2.Request) (interface{}, error) {
			if req.Method == handshakeMethod {
				var peerInfo PeerInfo
				if err := json.Unmarshal(req.Params, &peerInfo); err != nil {
					return nil, xerrors.Errorf("%w: unmarshaling client info: %v", jsonrpc2_v2.ErrInvalidParams, err)
				}
				peerInfo.LocalID = localID
				peerInfo.IsClient = true
				h.recordPeer(peerInfo)
				return info, nil
			}
			return delegate.Handle(ctx, req)
		})

		// Record the dropped client.
		go h.cleanupAtDisconnect(conn, localID)

		return opts, nil
	})
}

// ClientHandshake performs a client-side handshake with the server at the
// other end of conn, recording the server's peer info and watching for conn's
// disconnection.
func (h *Handshaker) ClientHandshake(ctx context.Context, conn *jsonrpc2_v2.Connection) {
	localID := h.nextID()
	info := &PeerInfo{
		RemoteID: localID,
		Metadata: h.Metadata,
	}

	call := conn.Call(ctx, handshakeMethod, info)
	var serverInfo PeerInfo
	if err := call.Await(ctx, &serverInfo); err != nil {
		event.Error(ctx, "performing handshake", err)
		return
	}
	serverInfo.LocalID = localID
	h.recordPeer(serverInfo)

	go h.cleanupAtDisconnect(conn, localID)
}

func (h *Handshaker) nextID() int64 {
	h.mu.Lock()
	defer h.mu.Unlock()

	h.prevID++
	return h.prevID
}

func (h *Handshaker) cleanupAtDisconnect(conn *jsonrpc2_v2.Connection, peerID int64) {
	conn.Wait()

	h.mu.Lock()
	defer h.mu.Unlock()
	delete(h.peers, peerID)
}

func (h *Handshaker) recordPeer(info PeerInfo) {
	h.mu.Lock()
	defer h.mu.Unlock()
	if h.peers == nil {
		h.peers = make(map[int64]PeerInfo)
	}
	h.peers[info.LocalID] = info
}