internal/lsp: get file URI from beginFileRequest in SemanticTokens

Unguarded calls to span.URI.Filename() can panic. beginFileRequest
handles this, so use the URI of the returned FileHandle instead.

Fixes golang/vscode-go#1498

Change-Id: Ie48c27854e4a8ed8cca52ff6547ff580eccb5fd5
Reviewed-on: https://go-review.googlesource.com/c/tools/+/319529
Trust: Robert Findley <[email protected]>
Run-TryBot: Robert Findley <[email protected]>
Reviewed-by: Heschi Kreinick <[email protected]>
gopls-CI: kokoro <[email protected]>
TryBot-Result: Go Bot <[email protected]>

Author: findleyr

gopls/internal/regtest/misc/semantictokens_test.go

@@ -0,0 +1,44 @@
+// 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 misc
+
+import (
+	"testing"
+
+	"golang.org/x/tools/internal/lsp/protocol"
+	. "golang.org/x/tools/internal/lsp/regtest"
+)
+
+func TestBadURICrash_VSCodeIssue1498(t *testing.T) {
+	const src = `
+-- go.mod --
+module example.com
+
+go 1.12
+
+-- main.go --
+package main
+
+func main() {}
+
+`
+	WithOptions(
+		Modes(Singleton),
+		EditorConfig{
+			AllExperiments: true,
+		},
+	).Run(t, src, func(t *testing.T, env *Env) {
+		params := &protocol.SemanticTokensParams{}
+		const badURI = "http://foo"
+		params.TextDocument.URI = badURI
+		// This call panicked in the past: golang/vscode-go#1498.
+		if _, err := env.Editor.Server.SemanticTokensFull(env.Ctx, params); err != nil {
+			// Requests to an invalid URI scheme shouldn't result in an error, we
+			// simply don't support this so return empty result. This could be
+			// changed, but for now assert on the current behavior.
+			t.Errorf("SemanticTokensFull(%q): %v", badURI, err)
+		}
+	})
+}

internal/lsp/semantic.go

@@ -49,8 +49,7 @@ func (s *Server) computeSemanticTokens(ctx context.Context, td protocol.TextDocu
 	ans := protocol.SemanticTokens{
 		Data: []uint32{},
 	}
-	kind := source.DetectLanguage("", td.URI.SpanURI().Filename())
-	snapshot, _, ok, release, err := s.beginFileRequest(ctx, td.URI, kind)
+	snapshot, fh, ok, release, err := s.beginFileRequest(ctx, td.URI, source.UnknownKind)
 	defer release()
 	if !ok {
 		return nil, err
@@ -61,7 +60,7 @@ func (s *Server) computeSemanticTokens(ctx context.Context, td protocol.TextDocu
 		// the client won't remember the wrong answer
 		return nil, errors.Errorf("semantictokens are disabled")
 	}
-	if kind == source.Tmpl {
+	if fh.Kind() == source.Tmpl {
 		// this is a little cumbersome to avoid both exporting 'encoded' and its methods
 		// and to avoid import cycles
 		e := &encoded{
@@ -76,14 +75,14 @@ func (s *Server) computeSemanticTokens(ctx context.Context, td protocol.TextDocu
 		data := func() ([]uint32, error) {
 			return e.Data()
 		}
-		return template.SemanticTokens(ctx, snapshot, td.URI.SpanURI(), add, data)
+		return template.SemanticTokens(ctx, snapshot, fh.URI(), add, data)
 	}
-	pkg, err := snapshot.PackageForFile(ctx, td.URI.SpanURI(), source.TypecheckFull, source.WidestPackage)
+	pkg, err := snapshot.PackageForFile(ctx, fh.URI(), source.TypecheckFull, source.WidestPackage)
 	if err != nil {
 		return nil, err
 	}
 	info := pkg.GetTypesInfo()
-	pgf, err := pkg.File(td.URI.SpanURI())
+	pgf, err := pkg.File(fh.URI())
 	if err != nil {
 		return nil, err
 	}
@@ -92,7 +91,7 @@ func (s *Server) computeSemanticTokens(ctx context.Context, td protocol.TextDocu
 	}
 	if rng == nil && len(pgf.Src) > maxFullFileSize {
 		err := fmt.Errorf("semantic tokens: file %s too large for full (%d>%d)",
-			td.URI.SpanURI().Filename(), len(pgf.Src), maxFullFileSize)
+			fh.URI().Filename(), len(pgf.Src), maxFullFileSize)
 		return nil, err
 	}
 	e := &encoded{