1// Copyright 2018 The Go Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style
3// license that can be found in the LICENSE file.
4
5package protocol
6
7import (
8	"context"
9	"encoding/json"
10	"fmt"
11	"io"
12
13	"golang.org/x/tools/internal/event"
14	"golang.org/x/tools/internal/jsonrpc2"
15	jsonrpc2_v2 "golang.org/x/tools/internal/jsonrpc2_v2"
16	"golang.org/x/tools/internal/xcontext"
17	errors "golang.org/x/xerrors"
18)
19
20var (
21	// RequestCancelledError should be used when a request is cancelled early.
22	RequestCancelledError   = jsonrpc2.NewError(-32800, "JSON RPC cancelled")
23	RequestCancelledErrorV2 = jsonrpc2_v2.NewError(-32800, "JSON RPC cancelled")
24)
25
26type ClientCloser interface {
27	Client
28	io.Closer
29}
30
31type connSender interface {
32	io.Closer
33
34	Notify(ctx context.Context, method string, params interface{}) error
35	Call(ctx context.Context, method string, params, result interface{}) error
36}
37
38type clientDispatcher struct {
39	sender connSender
40}
41
42func (c *clientDispatcher) Close() error {
43	return c.sender.Close()
44}
45
46// ClientDispatcher returns a Client that dispatches LSP requests across the
47// given jsonrpc2 connection.
48func ClientDispatcher(conn jsonrpc2.Conn) ClientCloser {
49	return &clientDispatcher{sender: clientConn{conn}}
50}
51
52type clientConn struct {
53	conn jsonrpc2.Conn
54}
55
56func (c clientConn) Close() error {
57	return c.conn.Close()
58}
59
60func (c clientConn) Notify(ctx context.Context, method string, params interface{}) error {
61	return c.conn.Notify(ctx, method, params)
62}
63
64func (c clientConn) Call(ctx context.Context, method string, params interface{}, result interface{}) error {
65	id, err := c.conn.Call(ctx, method, params, result)
66	if ctx.Err() != nil {
67		cancelCall(ctx, c, id)
68	}
69	return err
70}
71
72func ClientDispatcherV2(conn *jsonrpc2_v2.Connection) ClientCloser {
73	return &clientDispatcher{clientConnV2{conn}}
74}
75
76type clientConnV2 struct {
77	conn *jsonrpc2_v2.Connection
78}
79
80func (c clientConnV2) Close() error {
81	return c.conn.Close()
82}
83
84func (c clientConnV2) Notify(ctx context.Context, method string, params interface{}) error {
85	return c.conn.Notify(ctx, method, params)
86}
87
88func (c clientConnV2) Call(ctx context.Context, method string, params interface{}, result interface{}) error {
89	call := c.conn.Call(ctx, method, params)
90	err := call.Await(ctx, result)
91	if ctx.Err() != nil {
92		detached := xcontext.Detach(ctx)
93		c.conn.Notify(detached, "$/cancelRequest", &CancelParams{ID: call.ID().Raw()})
94	}
95	return err
96}
97
98// ServerDispatcher returns a Server that dispatches LSP requests across the
99// given jsonrpc2 connection.
100func ServerDispatcher(conn jsonrpc2.Conn) Server {
101	return &serverDispatcher{sender: clientConn{conn}}
102}
103
104func ServerDispatcherV2(conn *jsonrpc2_v2.Connection) Server {
105	return &serverDispatcher{sender: clientConnV2{conn}}
106}
107
108type serverDispatcher struct {
109	sender connSender
110}
111
112func ClientHandler(client Client, handler jsonrpc2.Handler) jsonrpc2.Handler {
113	return func(ctx context.Context, reply jsonrpc2.Replier, req jsonrpc2.Request) error {
114		if ctx.Err() != nil {
115			ctx := xcontext.Detach(ctx)
116			return reply(ctx, nil, RequestCancelledError)
117		}
118		handled, err := clientDispatch(ctx, client, reply, req)
119		if handled || err != nil {
120			return err
121		}
122		return handler(ctx, reply, req)
123	}
124}
125
126func ClientHandlerV2(client Client) jsonrpc2_v2.Handler {
127	return jsonrpc2_v2.HandlerFunc(func(ctx context.Context, req *jsonrpc2_v2.Request) (interface{}, error) {
128		if ctx.Err() != nil {
129			return nil, RequestCancelledErrorV2
130		}
131		req1 := req2to1(req)
132		var (
133			result interface{}
134			resErr error
135		)
136		replier := func(_ context.Context, res interface{}, err error) error {
137			result, resErr = res, err
138			return nil
139		}
140		_, err := clientDispatch(ctx, client, replier, req1)
141		if err != nil {
142			return nil, err
143		}
144		return result, resErr
145	})
146}
147
148func ServerHandler(server Server, handler jsonrpc2.Handler) jsonrpc2.Handler {
149	return func(ctx context.Context, reply jsonrpc2.Replier, req jsonrpc2.Request) error {
150		if ctx.Err() != nil {
151			ctx := xcontext.Detach(ctx)
152			return reply(ctx, nil, RequestCancelledError)
153		}
154		handled, err := serverDispatch(ctx, server, reply, req)
155		if handled || err != nil {
156			return err
157		}
158		//TODO: This code is wrong, it ignores handler and assumes non standard
159		// request handles everything
160		// non standard request should just be a layered handler.
161		var params interface{}
162		if err := json.Unmarshal(req.Params(), &params); err != nil {
163			return sendParseError(ctx, reply, err)
164		}
165		resp, err := server.NonstandardRequest(ctx, req.Method(), params)
166		return reply(ctx, resp, err)
167
168	}
169}
170
171func ServerHandlerV2(server Server) jsonrpc2_v2.Handler {
172	return jsonrpc2_v2.HandlerFunc(func(ctx context.Context, req *jsonrpc2_v2.Request) (interface{}, error) {
173		if ctx.Err() != nil {
174			return nil, RequestCancelledErrorV2
175		}
176		req1 := req2to1(req)
177		var (
178			result interface{}
179			resErr error
180		)
181		replier := func(_ context.Context, res interface{}, err error) error {
182			result, resErr = res, err
183			return nil
184		}
185		_, err := serverDispatch(ctx, server, replier, req1)
186		if err != nil {
187			return nil, err
188		}
189		return result, resErr
190	})
191}
192
193func req2to1(req2 *jsonrpc2_v2.Request) jsonrpc2.Request {
194	if req2.ID.IsValid() {
195		raw := req2.ID.Raw()
196		var idv1 jsonrpc2.ID
197		switch v := raw.(type) {
198		case int64:
199			idv1 = jsonrpc2.NewIntID(v)
200		case string:
201			idv1 = jsonrpc2.NewStringID(v)
202		default:
203			panic(fmt.Sprintf("unsupported ID type %T", raw))
204		}
205		req1, err := jsonrpc2.NewCall(idv1, req2.Method, req2.Params)
206		if err != nil {
207			panic(err)
208		}
209		return req1
210	}
211	req1, err := jsonrpc2.NewNotification(req2.Method, req2.Params)
212	if err != nil {
213		panic(err)
214	}
215	return req1
216}
217
218func Handlers(handler jsonrpc2.Handler) jsonrpc2.Handler {
219	return CancelHandler(
220		jsonrpc2.AsyncHandler(
221			jsonrpc2.MustReplyHandler(handler)))
222}
223
224func CancelHandler(handler jsonrpc2.Handler) jsonrpc2.Handler {
225	handler, canceller := jsonrpc2.CancelHandler(handler)
226	return func(ctx context.Context, reply jsonrpc2.Replier, req jsonrpc2.Request) error {
227		if req.Method() != "$/cancelRequest" {
228			// TODO(iancottrell): See if we can generate a reply for the request to be cancelled
229			// at the point of cancellation rather than waiting for gopls to naturally reply.
230			// To do that, we need to keep track of whether a reply has been sent already and
231			// be careful about racing between the two paths.
232			// TODO(iancottrell): Add a test that watches the stream and verifies the response
233			// for the cancelled request flows.
234			replyWithDetachedContext := func(ctx context.Context, resp interface{}, err error) error {
235				// https://microsoft.github.io/language-server-protocol/specifications/specification-current/#cancelRequest
236				if ctx.Err() != nil && err == nil {
237					err = RequestCancelledError
238				}
239				ctx = xcontext.Detach(ctx)
240				return reply(ctx, resp, err)
241			}
242			return handler(ctx, replyWithDetachedContext, req)
243		}
244		var params CancelParams
245		if err := json.Unmarshal(req.Params(), &params); err != nil {
246			return sendParseError(ctx, reply, err)
247		}
248		if n, ok := params.ID.(float64); ok {
249			canceller(jsonrpc2.NewIntID(int64(n)))
250		} else if s, ok := params.ID.(string); ok {
251			canceller(jsonrpc2.NewStringID(s))
252		} else {
253			return sendParseError(ctx, reply, fmt.Errorf("request ID %v malformed", params.ID))
254		}
255		return reply(ctx, nil, nil)
256	}
257}
258
259func Call(ctx context.Context, conn jsonrpc2.Conn, method string, params interface{}, result interface{}) error {
260	id, err := conn.Call(ctx, method, params, result)
261	if ctx.Err() != nil {
262		cancelCall(ctx, clientConn{conn}, id)
263	}
264	return err
265}
266
267func cancelCall(ctx context.Context, sender connSender, id jsonrpc2.ID) {
268	ctx = xcontext.Detach(ctx)
269	ctx, done := event.Start(ctx, "protocol.canceller")
270	defer done()
271	// Note that only *jsonrpc2.ID implements json.Marshaler.
272	sender.Notify(ctx, "$/cancelRequest", &CancelParams{ID: &id})
273}
274
275func sendParseError(ctx context.Context, reply jsonrpc2.Replier, err error) error {
276	return reply(ctx, nil, errors.Errorf("%w: %s", jsonrpc2.ErrParse, err))
277}
278