1// Copyright 2021 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 lsprpc
6
7import (
8	"context"
9	"encoding/json"
10	"fmt"
11
12	"golang.org/x/tools/internal/event"
13	"golang.org/x/tools/internal/gocommand"
14	jsonrpc2_v2 "golang.org/x/tools/internal/jsonrpc2_v2"
15	"golang.org/x/tools/internal/lsp/protocol"
16)
17
18func GoEnvMiddleware() (Middleware, error) {
19	return BindHandler(func(delegate jsonrpc2_v2.Handler) jsonrpc2_v2.Handler {
20		return jsonrpc2_v2.HandlerFunc(func(ctx context.Context, req *jsonrpc2_v2.Request) (interface{}, error) {
21			if req.Method == "initialize" {
22				if err := addGoEnvToInitializeRequestV2(ctx, req); err != nil {
23					event.Error(ctx, "adding go env to initialize", err)
24				}
25			}
26			return delegate.Handle(ctx, req)
27		})
28	}), nil
29}
30
31func addGoEnvToInitializeRequestV2(ctx context.Context, req *jsonrpc2_v2.Request) error {
32	var params protocol.ParamInitialize
33	if err := json.Unmarshal(req.Params, &params); err != nil {
34		return err
35	}
36	var opts map[string]interface{}
37	switch v := params.InitializationOptions.(type) {
38	case nil:
39		opts = make(map[string]interface{})
40	case map[string]interface{}:
41		opts = v
42	default:
43		return fmt.Errorf("unexpected type for InitializationOptions: %T", v)
44	}
45	envOpt, ok := opts["env"]
46	if !ok {
47		envOpt = make(map[string]interface{})
48	}
49	env, ok := envOpt.(map[string]interface{})
50	if !ok {
51		return fmt.Errorf("env option is %T, expected a map", envOpt)
52	}
53	goenv, err := getGoEnv(ctx, env)
54	if err != nil {
55		return err
56	}
57	for govar, value := range goenv {
58		env[govar] = value
59	}
60	opts["env"] = env
61	params.InitializationOptions = opts
62	raw, err := json.Marshal(params)
63	if err != nil {
64		return fmt.Errorf("marshaling updated options: %v", err)
65	}
66	req.Params = json.RawMessage(raw)
67	return nil
68}
69
70func getGoEnv(ctx context.Context, env map[string]interface{}) (map[string]string, error) {
71	var runEnv []string
72	for k, v := range env {
73		runEnv = append(runEnv, fmt.Sprintf("%s=%s", k, v))
74	}
75	runner := gocommand.Runner{}
76	output, err := runner.Run(ctx, gocommand.Invocation{
77		Verb: "env",
78		Args: []string{"-json"},
79		Env:  runEnv,
80	})
81	if err != nil {
82		return nil, err
83	}
84	envmap := make(map[string]string)
85	if err := json.Unmarshal(output.Bytes(), &envmap); err != nil {
86		return nil, err
87	}
88	return envmap, nil
89}
90