1// Copyright 2020 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
5// Package lsprpc implements a jsonrpc2.StreamServer that may be used to
6// serve the LSP on a jsonrpc2 channel.
7package lsprpc
8
9import (
10	"context"
11	"fmt"
12	"net"
13
14	"golang.org/x/sync/errgroup"
15	"golang.org/x/tools/internal/jsonrpc2"
16	"golang.org/x/tools/internal/lsp"
17	"golang.org/x/tools/internal/lsp/protocol"
18	"golang.org/x/tools/internal/lsp/source"
19)
20
21// The StreamServer type is a jsonrpc2.StreamServer that handles incoming
22// streams as a new LSP session, using a shared cache.
23type StreamServer struct {
24	withTelemetry bool
25
26	// accept is mutable for testing.
27	accept func(protocol.Client) protocol.Server
28}
29
30// NewStreamServer creates a StreamServer using the shared cache. If
31// withTelemetry is true, each session is instrumented with telemetry that
32// records RPC statistics.
33func NewStreamServer(cache source.Cache, withTelemetry bool) *StreamServer {
34	s := &StreamServer{
35		withTelemetry: withTelemetry,
36	}
37	s.accept = func(c protocol.Client) protocol.Server {
38		session := cache.NewSession()
39		return lsp.NewServer(session, c)
40	}
41	return s
42}
43
44// ServeStream implements the jsonrpc2.StreamServer interface, by handling
45// incoming streams using a new lsp server.
46func (s *StreamServer) ServeStream(ctx context.Context, stream jsonrpc2.Stream) error {
47	conn := jsonrpc2.NewConn(stream)
48	client := protocol.ClientDispatcher(conn)
49	server := s.accept(client)
50	conn.AddHandler(protocol.ServerHandler(server))
51	conn.AddHandler(protocol.Canceller{})
52	if s.withTelemetry {
53		conn.AddHandler(telemetryHandler{})
54	}
55	return conn.Run(protocol.WithClient(ctx, client))
56}
57
58// A Forwarder is a jsonrpc2.StreamServer that handles an LSP stream by
59// forwarding it to a remote. This is used when the gopls process started by
60// the editor is in the `-remote` mode, which means it finds and connects to a
61// separate gopls daemon. In these cases, we still want the forwarder gopls to
62// be instrumented with telemetry, and want to be able to in some cases hijack
63// the jsonrpc2 connection with the daemon.
64type Forwarder struct {
65	remote        string
66	withTelemetry bool
67}
68
69// NewForwarder creates a new Forwarder, ready to forward connections to the
70// given remote.
71func NewForwarder(remote string, withTelemetry bool) *Forwarder {
72	return &Forwarder{
73		remote:        remote,
74		withTelemetry: withTelemetry,
75	}
76}
77
78// ServeStream dials the forwarder remote and binds the remote to serve the LSP
79// on the incoming stream.
80func (f *Forwarder) ServeStream(ctx context.Context, stream jsonrpc2.Stream) error {
81	clientConn := jsonrpc2.NewConn(stream)
82	client := protocol.ClientDispatcher(clientConn)
83
84	netConn, err := net.Dial("tcp", f.remote)
85	if err != nil {
86		return fmt.Errorf("forwarder: dialing remote: %v", err)
87	}
88	serverConn := jsonrpc2.NewConn(jsonrpc2.NewHeaderStream(netConn, netConn))
89	server := protocol.ServerDispatcher(serverConn)
90
91	// Forward between connections.
92	serverConn.AddHandler(protocol.ClientHandler(client))
93	serverConn.AddHandler(protocol.Canceller{})
94	clientConn.AddHandler(protocol.ServerHandler(server))
95	clientConn.AddHandler(protocol.Canceller{})
96	if f.withTelemetry {
97		clientConn.AddHandler(telemetryHandler{})
98	}
99
100	g, ctx := errgroup.WithContext(ctx)
101	g.Go(func() error {
102		return serverConn.Run(ctx)
103	})
104	g.Go(func() error {
105		return clientConn.Run(ctx)
106	})
107	return g.Wait()
108}
109