internal/lsp/source: trim file very carefully

Apparently the AST will sometimes give us offsets past the end of the
file. Don't crash when it does.

Fixes golang/go#36610.

Change-Id: I3cfbf8645bfcea94a5d87bca5bef4236d657b2c0
Reviewed-on: https://go-review.googlesource.com/c/tools/+/215119
Run-TryBot: Heschi Kreinick <[email protected]>
TryBot-Result: Gobot Gobot <[email protected]>
Reviewed-by: Rebecca Stambler <[email protected]>

Author: heschi

internal/lsp/source/format.go

@@ -245,8 +245,13 @@ func trimToImports(fset *token.FileSet, f *ast.File, src []byte) ([]byte, int) {
 	tok := fset.File(f.Pos())
 	start := firstImport.Pos()
 	end := lastImport.End()
-	if tok.LineCount() > fset.Position(end).Line {
-		end = fset.File(f.Pos()).LineStart(fset.Position(lastImport.End()).Line + 1)
+	// The parser will happily feed us nonsense. See golang/go#36610.
+	tokStart, tokEnd := token.Pos(tok.Base()), token.Pos(tok.Base()+tok.Size())
+	if start < tokStart || start > tokEnd || end < tokStart || end > tokEnd {
+		return nil, 0
+	}
+	if nextLine := fset.Position(end).Line + 1; tok.LineCount() >= nextLine {
+		end = fset.File(f.Pos()).LineStart(nextLine)
 	}
 
 	startLineOffset := fset.Position(start).Line - 1 // lines are 1-indexed.

internal/lsp/source/format_test.go

@@ -0,0 +1,25 @@
+package source
+
+import (
+	"go/parser"
+	"go/token"
+	"testing"
+)
+
+func TestTrimToImports(t *testing.T) {
+	const input = `package source
+
+import (
+	m
+	"fmt"
+)
+
+func foo() {
+	fmt.Println("hi")
+}
+`
+
+	fs := token.NewFileSet()
+	f, _ := parser.ParseFile(fs, "foo.go", input, parser.ImportsOnly)
+	trimToImports(fs, f, []byte(input))
+}