From 95ece921ff98c72e28d959591cd97f2ec3b3840c Mon Sep 17 00:00:00 2001
From: Heschi Kreinick <heschi@google.com>
Date: Thu, 16 Jan 2020 19:45:26 -0500
Subject: [PATCH] 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 <heschi@google.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Rebecca Stambler <rstambler@golang.org>
---
 internal/lsp/source/format.go      |  9 +++++++--
 internal/lsp/source/format_test.go | 25 +++++++++++++++++++++++++
 2 files changed, 32 insertions(+), 2 deletions(-)
 create mode 100644 internal/lsp/source/format_test.go

diff --git a/internal/lsp/source/format.go b/internal/lsp/source/format.go
index 6d6dfbf512..2149fea434 100644
--- a/internal/lsp/source/format.go
+++ b/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.
diff --git a/internal/lsp/source/format_test.go b/internal/lsp/source/format_test.go
new file mode 100644
index 0000000000..678a6727a9
--- /dev/null
+++ b/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))
+}