gopls/internal/golang: refine crash golang/go#70553

This CL adds defensive assertions to ExtractToNewFile, which
experiences a crash due to out-of-bounds indexing. The logic
looks sound, but clearly something is off.

This CL also fixes a mistake in the logic added to parsego.Parse
to work around golang/go#70162 (missing File.FileStart in go/parser),
but I don't think that bug explains golang/go#70553 because it is
concerned only with deeply broken files.

Updates golang/go#70553
Updates golang/go#70162

Change-Id: Ie4f6fbbde2046023d7a3b865e5810bbed3be3118
Reviewed-on: https://go-review.googlesource.com/c/tools/+/631675
Reviewed-by: Robert Findley <[email protected]>
LUCI-TryBot-Result: Go LUCI <[email protected]>

Author: adonovan

Diff for: gopls/internal/cache/parsego/file.go

@@ -12,6 +12,7 @@ import (
 	"sync"
 
 	"golang.org/x/tools/gopls/internal/protocol"
+	"golang.org/x/tools/gopls/internal/util/bug"
 	"golang.org/x/tools/gopls/internal/util/safetoken"
 )
 
@@ -116,6 +117,22 @@ func (pgf *File) RangePos(r protocol.Range) (token.Pos, token.Pos, error) {
 	return pgf.Tok.Pos(start), pgf.Tok.Pos(end), nil
 }
 
+// CheckNode asserts that the Node's positions are valid w.r.t. pgf.Tok.
+func (pgf *File) CheckNode(node ast.Node) {
+	// Avoid safetoken.Offsets, and put each assertion on its own source line.
+	pgf.CheckPos(node.Pos())
+	pgf.CheckPos(node.End())
+}
+
+// CheckPos asserts that the position is valid w.r.t. pgf.Tok.
+func (pgf *File) CheckPos(pos token.Pos) {
+	if !pos.IsValid() {
+		bug.Report("invalid token.Pos")
+	} else if _, err := safetoken.Offset(pgf.Tok, pos); err != nil {
+		bug.Report("token.Pos out of range")
+	}
+}
+
 // Resolve lazily resolves ast.Ident.Objects in the enclosed syntax tree.
 //
 // Resolve must be called before accessing any of:

Diff for: gopls/internal/cache/parsego/parse.go

@@ -25,6 +25,7 @@ import (
 	"golang.org/x/tools/gopls/internal/label"
 	"golang.org/x/tools/gopls/internal/protocol"
 	"golang.org/x/tools/gopls/internal/util/astutil"
+	"golang.org/x/tools/gopls/internal/util/bug"
 	"golang.org/x/tools/gopls/internal/util/safetoken"
 	"golang.org/x/tools/internal/diff"
 	"golang.org/x/tools/internal/event"
@@ -77,12 +78,21 @@ func Parse(ctx context.Context, fset *token.FileSet, uri protocol.DocumentURI, s
 	tokenFile := func(file *ast.File) *token.File {
 		tok := fset.File(file.FileStart)
 		if tok == nil {
+			// Invalid File.FileStart (also File.{Package,Name.Pos}).
+			if file.Package.IsValid() {
+				bug.Report("ast.File has valid Package but no FileStart")
+			}
+			if file.Name.Pos().IsValid() {
+				bug.Report("ast.File has valid Name.Pos but no FileStart")
+			}
 			tok = fset.AddFile(uri.Path(), -1, len(src))
 			tok.SetLinesForContent(src)
-			if file.FileStart.IsValid() {
-				file.FileStart = token.Pos(tok.Base())
-				file.FileEnd = token.Pos(tok.Base() + tok.Size())
-			}
+			// If the File contained any valid token.Pos values,
+			// they would all be invalid wrt the new token.File,
+			// but we have established that it lacks FileStart,
+			// Package, and Name.Pos.
+			file.FileStart = token.Pos(tok.Base())
+			file.FileEnd = token.Pos(tok.Base() + tok.Size())
 		}
 		return tok
 	}

Diff for: gopls/internal/golang/extracttofile.go

@@ -97,14 +97,19 @@ func ExtractToNewFile(ctx context.Context, snapshot *cache.Snapshot, fh file.Han
 	if !ok {
 		return nil, bug.Errorf("invalid selection")
 	}
+	pgf.CheckPos(start) // #70553
+	// Inv: start is valid wrt pgf.Tok.
 
 	// select trailing empty lines
 	offset, err := safetoken.Offset(pgf.Tok, end)
 	if err != nil {
 		return nil, err
 	}
 	rest := pgf.Src[offset:]
-	end += token.Pos(len(rest) - len(bytes.TrimLeft(rest, " \t\n")))
+	spaces := len(rest) - len(bytes.TrimLeft(rest, " \t\n"))
+	end += token.Pos(spaces)
+	pgf.CheckPos(end) // #70553
+	// Inv: end is valid wrt pgf.Tok.
 
 	replaceRange, err := pgf.PosRange(start, end)
 	if err != nil {
@@ -172,6 +177,7 @@ func ExtractToNewFile(ctx context.Context, snapshot *cache.Snapshot, fh file.Han
 	}
 
 	fileStart := pgf.File.FileStart
+	pgf.CheckPos(fileStart) // #70553
 	buf.Write(pgf.Src[start-fileStart : end-fileStart])
 
 	newFileContent, err := format.Source(buf.Bytes())
@@ -221,31 +227,42 @@ func selectedToplevelDecls(pgf *parsego.File, start, end token.Pos) (token.Pos,
 	firstName := ""
 	for _, decl := range pgf.File.Decls {
 		if posRangeIntersects(start, end, decl.Pos(), decl.End()) {
-			var id *ast.Ident
-			switch v := decl.(type) {
+			var (
+				comment *ast.CommentGroup // (include comment preceding decl)
+				id      *ast.Ident
+			)
+			switch decl := decl.(type) {
 			case *ast.BadDecl:
 				return 0, 0, "", false
+
 			case *ast.FuncDecl:
 				// if only selecting keyword "func" or function name, extend selection to the
 				// whole function
-				if posRangeContains(v.Pos(), v.Name.End(), start, end) {
-					start, end = v.Pos(), v.End()
+				if posRangeContains(decl.Pos(), decl.Name.End(), start, end) {
+					pgf.CheckNode(decl) // #70553
+					start, end = decl.Pos(), decl.End()
+					// Inv: start, end are valid wrt pgf.Tok.
 				}
-				id = v.Name
+				comment = decl.Doc
+				id = decl.Name
+
 			case *ast.GenDecl:
 				// selection cannot intersect an import declaration
-				if v.Tok == token.IMPORT {
+				if decl.Tok == token.IMPORT {
 					return 0, 0, "", false
 				}
 				// if only selecting keyword "type", "const", or "var", extend selection to the
 				// whole declaration
-				if v.Tok == token.TYPE && posRangeContains(v.Pos(), v.Pos()+token.Pos(len("type")), start, end) ||
-					v.Tok == token.CONST && posRangeContains(v.Pos(), v.Pos()+token.Pos(len("const")), start, end) ||
-					v.Tok == token.VAR && posRangeContains(v.Pos(), v.Pos()+token.Pos(len("var")), start, end) {
-					start, end = v.Pos(), v.End()
+				if decl.Tok == token.TYPE && posRangeContains(decl.Pos(), decl.Pos()+token.Pos(len("type")), start, end) ||
+					decl.Tok == token.CONST && posRangeContains(decl.Pos(), decl.Pos()+token.Pos(len("const")), start, end) ||
+					decl.Tok == token.VAR && posRangeContains(decl.Pos(), decl.Pos()+token.Pos(len("var")), start, end) {
+					pgf.CheckNode(decl) // #70553
+					start, end = decl.Pos(), decl.End()
+					// Inv: start, end are valid wrt pgf.Tok.
 				}
-				if len(v.Specs) > 0 {
-					switch spec := v.Specs[0].(type) {
+				comment = decl.Doc
+				if len(decl.Specs) > 0 {
+					switch spec := decl.Specs[0].(type) {
 					case *ast.TypeSpec:
 						id = spec.Name
 					case *ast.ValueSpec:
@@ -261,16 +278,10 @@ func selectedToplevelDecls(pgf *parsego.File, start, end token.Pos) (token.Pos,
 				// may be "_"
 				firstName = id.Name
 			}
-			// extends selection to docs comments
-			var c *ast.CommentGroup
-			switch decl := decl.(type) {
-			case *ast.GenDecl:
-				c = decl.Doc
-			case *ast.FuncDecl:
-				c = decl.Doc
-			}
-			if c != nil && c.Pos() < start {
-				start = c.Pos()
+			if comment != nil && comment.Pos() < start {
+				pgf.CheckNode(comment) // #70553
+				start = comment.Pos()
+				// Inv: start is valid wrt pgf.Tok.
 			}
 		}
 	}