internal/lsp/lsprpc: forward the go environment in initialize requests

The gopls workspace environment defaults to the process environment in
which gopls was started. This means that when switching environments,
gopls can potentially get a different environment when connecting as an
editor sidecar from when forwarding requests to the daemon.

To (hopefully mostly) mitigate this pain point, inject the Go
environment when forwarding the 'initialize' request, which contains
InitializationOptions containing the 'env' configuration. We could go
further and send the entire os.Environ(), but that seems problematic
both in its unbounded nature, and because in many cases the user may not
actually want to send their process env over the wire. Gopls behavior
should *mostly* be parameterized by gopls binary and Go env, and after
this change these should match for forwarder and daemon.

For go1.15, Explicitly set the GOMODCACHE environment variable in the
regtest sandbox. Without this, regtests were failing in the forwarded
environment because they implicitly shared a module cache.

Fixes golang/go#37830

Change-Id: Ic1b335506f8b481505eac9f74c0df6293dc07158
Reviewed-on: https://go-review.googlesource.com/c/tools/+/234109
Run-TryBot: Robert Findley <rfindley@google.com>
Reviewed-by: Ian Cottrell <iancottrell@google.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>

internal/lsp/fake/sandbox.go

@@ -13,6 +13,7 @@ import (
 	"strings"
 
 	"golang.org/x/tools/internal/gocommand"
+	"golang.org/x/tools/internal/testenv"
 	"golang.org/x/tools/txtar"
 )
 
@@ -101,12 +102,16 @@ func (sb *Sandbox) GOPATH() string {
 // GoEnv returns the default environment variables that can be used for
 // invoking Go commands in the sandbox.
 func (sb *Sandbox) GoEnv() []string {
-	return []string{
+	vars := []string{
 		"GOPATH=" + sb.GOPATH(),
 		"GOPROXY=" + sb.Proxy.GOPROXY(),
 		"GO111MODULE=",
 		"GOSUMDB=off",
 	}
+	if testenv.Go1Point() >= 5 {
+		vars = append(vars, "GOMODCACHE=")
+	}
+	return vars
 }
 
 // RunGoCommand executes a go command in the sandbox.

internal/lsp/lsprpc/lsprpc.go

@@ -18,6 +18,7 @@ import (
 	"time"
 
 	"golang.org/x/tools/internal/event"
+	"golang.org/x/tools/internal/gocommand"
 	"golang.org/x/tools/internal/jsonrpc2"
 	"golang.org/x/tools/internal/lsp"
 	"golang.org/x/tools/internal/lsp/cache"
@@ -225,8 +226,9 @@ func (f *Forwarder) ServeStream(ctx context.Context, clientConn jsonrpc2.Conn) e
 	)
 	clientConn.Go(ctx,
 		protocol.Handlers(
-			protocol.ServerHandler(server,
+			forwarderHandler(
-				jsonrpc2.MethodNotFound)))
+				protocol.ServerHandler(server,
+					jsonrpc2.MethodNotFound))))
 
 	select {
 	case <-serverConn.Done():
@@ -321,6 +323,93 @@ func (f *Forwarder) connectToRemote(ctx context.Context) (net.Conn, error) {
 	return nil, fmt.Errorf("dialing remote: %w", err)
 }
 
+// forwarderHandler intercepts 'exit' messages to prevent the shared gopls
+// instance from exiting. In the future it may also intercept 'shutdown' to
+// provide more graceful shutdown of the client connection.
+func forwarderHandler(handler jsonrpc2.Handler) jsonrpc2.Handler {
+	return func(ctx context.Context, reply jsonrpc2.Replier, r jsonrpc2.Request) error {
+		// The gopls workspace environment defaults to the process environment in
+		// which gopls daemon was started. To avoid discrepancies in Go environment
+		// between the editor and daemon, inject any unset variables in `go env`
+		// into the options sent by initialize.
+		//
+		// See also golang.org/issue/37830.
+		if r.Method() == "initialize" {
+			if newr, err := addGoEnvToInitializeRequest(ctx, r); err == nil {
+				r = newr
+			} else {
+				log.Printf("unable to add local env to initialize request: %v", err)
+			}
+		}
+		return handler(ctx, reply, r)
+	}
+}
+
+// addGoEnvToInitializeRequest builds a new initialize request in which we set
+// any environment variables output by `go env` and not already present in the
+// request.
+//
+// It returns an error if r is not an initialize requst, or is otherwise
+// malformed.
+func addGoEnvToInitializeRequest(ctx context.Context, r jsonrpc2.Request) (jsonrpc2.Request, error) {
+	var params protocol.ParamInitialize
+	if err := json.Unmarshal(r.Params(), &params); err != nil {
+		return nil, err
+	}
+	var opts map[string]interface{}
+	switch v := params.InitializationOptions.(type) {
+	case nil:
+		opts = make(map[string]interface{})
+	case map[string]interface{}:
+		opts = v
+	default:
+		return nil, fmt.Errorf("unexpected type for InitializationOptions: %T", v)
+	}
+	envOpt, ok := opts["env"]
+	if !ok {
+		envOpt = make(map[string]interface{})
+	}
+	env, ok := envOpt.(map[string]interface{})
+	if !ok {
+		return nil, fmt.Errorf(`env option is %T, expected a map`, envOpt)
+	}
+	goenv, err := getGoEnv(ctx, env)
+	if err != nil {
+		return nil, err
+	}
+	for govar, value := range goenv {
+		env[govar] = value
+	}
+	opts["env"] = env
+	params.InitializationOptions = opts
+	call, ok := r.(*jsonrpc2.Call)
+	if !ok {
+		return nil, fmt.Errorf("%T is not a *jsonrpc2.Call", r)
+	}
+	return jsonrpc2.NewCall(call.ID(), "initialize", params)
+}
+
+func getGoEnv(ctx context.Context, env map[string]interface{}) (map[string]string, error) {
+	var runEnv []string
+	for k, v := range env {
+		runEnv = append(runEnv, fmt.Sprintf("%s=%s", k, v))
+	}
+	runner := gocommand.Runner{}
+	output, err := runner.Run(ctx, gocommand.Invocation{
+		Verb: "env",
+		Args: []string{"-json"},
+		Env:  runEnv,
+	})
+	if err != nil {
+		return nil, err
+	}
+	envmap := make(map[string]string)
+	if err := json.Unmarshal(output.Bytes(), &envmap); err != nil {
+		return nil, err
+	}
+	return envmap, nil
+}
+
 // A handshakeRequest identifies a client to the LSP server.
 type handshakeRequest struct {
 	// ServerID is the ID of the server on the client. This should usually be 0.

internal/lsp/lsprpc/lsprpc_test.go

@@ -31,14 +31,20 @@ func (c fakeClient) LogMessage(ctx context.Context, params *protocol.LogMessageP
 	return nil
 }
 
-type pingServer struct{ protocol.Server }
+// fakeServer is intended to be embedded in the test fakes below, to trivially
+// implement Shutdown.
+type fakeServer struct {
+	protocol.Server
+}
 
-func (s pingServer) DidOpen(ctx context.Context, params *protocol.DidOpenTextDocumentParams) error {
+func (fakeServer) Shutdown(ctx context.Context) error {
-	event.Log(ctx, "ping")
 	return nil
 }
 
-func (s pingServer) Shutdown(ctx context.Context) error {
+type pingServer struct{ fakeServer }
+
+func (s pingServer) DidOpen(ctx context.Context, params *protocol.DidOpenTextDocumentParams) error {
+	event.Log(ctx, "ping")
 	return nil
 }
 
@@ -78,7 +84,7 @@ func TestClientLogging(t *testing.T) {
 // The requests chosen are arbitrary: we simply needed one that blocks, and
 // another that doesn't.
 type waitableServer struct {
-	protocol.Server
+	fakeServer
 
 	started chan struct{}
 }
@@ -97,10 +103,6 @@ func (s waitableServer) Resolve(_ context.Context, item *protocol.CompletionItem
 	return item, nil
 }
 
-func (s waitableServer) Shutdown(ctx context.Context) error {
-	return nil
-}
-
 func checkClose(t *testing.T, closer func() error) {
 	t.Helper()
 	if err := closer(); err != nil {
@@ -108,22 +110,29 @@ func checkClose(t *testing.T, closer func() error) {
 	}
 }
 
+func setupForwarding(ctx context.Context, t *testing.T, s protocol.Server) (direct, forwarded servertest.Connector, cleanup func()) {
+	t.Helper()
+	serveCtx := debug.WithInstance(ctx, "", "")
+	ss := NewStreamServer(cache.New(serveCtx, nil))
+	ss.serverForTest = s
+	tsDirect := servertest.NewTCPServer(serveCtx, ss, nil)
+
+	forwarderCtx := debug.WithInstance(ctx, "", "")
+	forwarder := NewForwarder("tcp", tsDirect.Addr)
+	tsForwarded := servertest.NewPipeServer(forwarderCtx, forwarder, nil)
+	return tsDirect, tsForwarded, func() {
+		checkClose(t, tsDirect.Close)
+		checkClose(t, tsForwarded.Close)
+	}
+}
+
 func TestRequestCancellation(t *testing.T) {
+	ctx := context.Background()
 	server := waitableServer{
 		started: make(chan struct{}),
 	}
-	baseCtx := context.Background()
+	tsDirect, tsForwarded, cleanup := setupForwarding(ctx, t, server)
-	serveCtx := debug.WithInstance(baseCtx, "", "")
+	defer cleanup()
-	ss := NewStreamServer(cache.New(serveCtx, nil))
-	ss.serverForTest = server
-	tsDirect := servertest.NewTCPServer(serveCtx, ss, nil)
-	defer checkClose(t, tsDirect.Close)
-
-	forwarderCtx := debug.WithInstance(baseCtx, "", "")
-	forwarder := NewForwarder("tcp", tsDirect.Addr)
-	tsForwarded := servertest.NewPipeServer(forwarderCtx, forwarder, nil)
-	defer checkClose(t, tsForwarded.Close)
-
 	tests := []struct {
 		serverType string
 		ts         servertest.Connector
@@ -134,9 +143,9 @@ func TestRequestCancellation(t *testing.T) {
 
 	for _, test := range tests {
 		t.Run(test.serverType, func(t *testing.T) {
-			cc := test.ts.Connect(baseCtx)
+			cc := test.ts.Connect(ctx)
 			sd := protocol.ServerDispatcher(cc)
-			cc.Go(baseCtx,
+			cc.Go(ctx,
 				protocol.Handlers(
 					jsonrpc2.MethodNotFound))
 
@@ -266,3 +275,49 @@ func TestDebugInfoLifecycle(t *testing.T) {
 		t.Errorf("len(server:Sessions()) = %d, want %d", got, want)
 	}
 }
+
+type initServer struct {
+	fakeServer
+
+	params *protocol.ParamInitialize
+}
+
+func (s *initServer) Initialize(ctx context.Context, params *protocol.ParamInitialize) (*protocol.InitializeResult, error) {
+	s.params = params
+	return &protocol.InitializeResult{}, nil
+}
+
+func TestEnvForwarding(t *testing.T) {
+	server := &initServer{}
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	_, tsForwarded, cleanup := setupForwarding(ctx, t, server)
+	defer cleanup()
+
+	conn := tsForwarded.Connect(ctx)
+	conn.Go(ctx, jsonrpc2.MethodNotFound)
+	dispatch := protocol.ServerDispatcher(conn)
+	initParams := &protocol.ParamInitialize{}
+	initParams.InitializationOptions = map[string]interface{}{
+		"env": map[string]interface{}{
+			"GONOPROXY": "example.com",
+		},
+	}
+	_, err := dispatch.Initialize(ctx, initParams)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if server.params == nil {
+		t.Fatalf("initialize params are unset")
+	}
+	env := server.params.InitializationOptions.(map[string]interface{})["env"].(map[string]interface{})
+
+	// Check for an arbitrary Go variable. It should be set.
+	if _, ok := env["GOPRIVATE"]; !ok {
+		t.Errorf("Go environment variable GOPRIVATE unset in initialization options")
+	}
+	// Check that the variable present in our user config was not overwritten.
+	if v := env["GONOPROXY"]; v != "example.com" {
+		t.Errorf("GONOPROXY environment variable was overwritten")
+	}
+}

internal/lsp/regtest/env.go

@@ -549,6 +549,7 @@ func (e *Env) AnyDiagnosticAtCurrentVersion(name string) DiagnosticExpectation {
 // position matching the regexp search string re in the buffer specified by
 // name. Note that this currently ignores the end position.
 func (e *Env) DiagnosticAtRegexp(name, re string) DiagnosticExpectation {
+	e.T.Helper()
 	pos := e.RegexpSearch(name, re)
 	expectation := DiagnosticAt(name, pos.Line, pos.Column)
 	expectation.description += fmt.Sprintf(" (location of %q)", re)

internal/lsp/regtest/runner.go

@@ -225,7 +225,9 @@ func (s *loggingFramer) printBuffers(testname string, w io.Writer) {
 
 	for i, buf := range s.buffers {
 		fmt.Fprintf(os.Stderr, "#### Start Gopls Test Logs %d of %d for %q\n", i+1, len(s.buffers), testname)
-		io.Copy(w, buf)
+		// Re-buffer buf to avoid a data rate (io.Copy mutates src).
+		writeBuf := bytes.NewBuffer(buf.Bytes())
+		io.Copy(w, writeBuf)
 		fmt.Fprintf(os.Stderr, "#### End Gopls Test Logs %d of %d for %q\n", i+1, len(s.buffers), testname)
 	}
 }