playground: format go.mod files; fix paths in formatting errors

Previously, handleFmt applied gofmt/goimports formatting to .go files
only. This change makes it apply formatting to go.mod files as well.
The cmd/go/internal/modfile package (well, a non-internal copy thereof)
is used to perform the go.mod file formatting.

Add test cases for error messages, and fix some cases where the paths
weren't accurate.

Detect when the error was returned by format.Source and needs to be
prefixed by using the fixImports variable instead of checking for the
presence of the prefix. This makes the code simpler and more readable.

Replace old fs.m[f] usage with fs.Data(f) in fmt.go and txtar_test.go.

Updates golang/go#32040
Updates golang/go#31944

Change-Id: Iefef7337f962914817558bcf0c622a952160ac44
Reviewed-on: https://go-review.googlesource.com/c/playground/+/177421
Reviewed-by: Brad Fitzpatrick <bradfitz@golang.org>
diff --git a/Dockerfile b/Dockerfile
index 27adf06..9b3c852 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -51,6 +51,7 @@
 RUN go install github.com/bradfitz/gomemcache/memcache
 RUN go install golang.org/x/tools/godoc/static
 RUN go install golang.org/x/tools/imports
+RUN go install github.com/rogpeppe/go-internal/modfile
 RUN go install github.com/rogpeppe/go-internal/txtar
 
 # Add and compile playground daemon
diff --git a/fmt.go b/fmt.go
index c5aa943..002bd34 100644
--- a/fmt.go
+++ b/fmt.go
@@ -9,8 +9,9 @@
 	"fmt"
 	"go/format"
 	"net/http"
-	"strings"
+	"path"
 
+	"github.com/rogpeppe/go-internal/modfile"
 	"golang.org/x/tools/imports"
 )
 
@@ -30,30 +31,46 @@
 
 	fixImports := r.FormValue("imports") != ""
 	for _, f := range fs.files {
-		if !strings.HasSuffix(f, ".go") {
-			continue
-		}
-		var out []byte
-		var err error
-		in := fs.m[f]
-		if fixImports {
-			// TODO: pass options to imports.Process so it
-			// can find symbols in sibling files.
-			out, err = imports.Process(progName, in, nil)
-		} else {
-			out, err = format.Source(in)
-		}
-		if err != nil {
-			errMsg := err.Error()
-			// Prefix the error returned by format.Source.
-			if !strings.HasPrefix(errMsg, f) {
-				errMsg = fmt.Sprintf("%v:%v", f, errMsg)
+		switch {
+		case path.Ext(f) == ".go":
+			var out []byte
+			var err error
+			in := fs.Data(f)
+			if fixImports {
+				// TODO: pass options to imports.Process so it
+				// can find symbols in sibling files.
+				out, err = imports.Process(f, in, nil)
+			} else {
+				out, err = format.Source(in)
 			}
-			json.NewEncoder(w).Encode(fmtResponse{Error: errMsg})
-			return
+			if err != nil {
+				errMsg := err.Error()
+				if !fixImports {
+					// Unlike imports.Process, format.Source does not prefix
+					// the error with the file path. So, do it ourselves here.
+					errMsg = fmt.Sprintf("%v:%v", f, errMsg)
+				}
+				json.NewEncoder(w).Encode(fmtResponse{Error: errMsg})
+				return
+			}
+			fs.AddFile(f, out)
+		case path.Base(f) == "go.mod":
+			out, err := formatGoMod(f, fs.Data(f))
+			if err != nil {
+				json.NewEncoder(w).Encode(fmtResponse{Error: err.Error()})
+				return
+			}
+			fs.AddFile(f, out)
 		}
-		fs.AddFile(f, out)
 	}
 
 	json.NewEncoder(w).Encode(fmtResponse{Body: string(fs.Format())})
 }
+
+func formatGoMod(file string, data []byte) ([]byte, error) {
+	f, err := modfile.Parse(file, data, nil)
+	if err != nil {
+		return nil, err
+	}
+	return f.Format()
+}
diff --git a/fmt_test.go b/fmt_test.go
index b8a1b9a..e601dc3 100644
--- a/fmt_test.go
+++ b/fmt_test.go
@@ -47,9 +47,53 @@
 			want: "package main\n-- two.go --\npackage main\n\nvar X = 5\n",
 		},
 		{
-			name: "only_format_go",
-			body: "    package main\n\n\n-- go.mod --\n   module foo\n",
-			want: "package main\n-- go.mod --\n   module foo\n",
+			name: "single_go.mod_with_header",
+			body: "-- go.mod --\n   module   \"foo\"   ",
+			want: "-- go.mod --\nmodule foo\n",
+		},
+		{
+			name: "multi_go.mod_with_header",
+			body: "-- a/go.mod --\n  module foo\n\n\n-- b/go.mod --\n   module  \"bar\"",
+			want: "-- a/go.mod --\nmodule foo\n-- b/go.mod --\nmodule bar\n",
+		},
+		{
+			name: "only_format_go_and_go.mod",
+			body: "    package   main   \n\n\n" +
+				"-- go.mod --\n   module   foo   \n\n\n" +
+				"-- plain.txt --\n   plain   text   \n\n\n",
+			want: "package main\n-- go.mod --\nmodule foo\n-- plain.txt --\n   plain   text   \n\n\n",
+		},
+		{
+			name:    "error_gofmt",
+			body:    "package 123\n",
+			wantErr: "prog.go:1:9: expected 'IDENT', found 123",
+		},
+		{
+			name:    "error_gofmt_with_header",
+			body:    "-- dir/one.go --\npackage 123\n",
+			wantErr: "dir/one.go:1:9: expected 'IDENT', found 123",
+		},
+		{
+			name:    "error_goimports",
+			body:    "package 123\n",
+			imports: true,
+			wantErr: "prog.go:1:9: expected 'IDENT', found 123",
+		},
+		{
+			name:    "error_goimports_with_header",
+			body:    "-- dir/one.go --\npackage 123\n",
+			imports: true,
+			wantErr: "dir/one.go:1:9: expected 'IDENT', found 123",
+		},
+		{
+			name:    "error_go.mod",
+			body:    "-- go.mod --\n123\n",
+			wantErr: "go.mod:1: unknown directive: 123",
+		},
+		{
+			name:    "error_go.mod_with_header",
+			body:    "-- dir/go.mod --\n123\n",
+			wantErr: "dir/go.mod:1: unknown directive: 123",
 		},
 	} {
 		t.Run(tt.name, func(t *testing.T) {
diff --git a/txtar_test.go b/txtar_test.go
index ae1ef96..493d140 100644
--- a/txtar_test.go
+++ b/txtar_test.go
@@ -147,7 +147,7 @@
 		if i == 0 && f == progName && fs.noHeader {
 			implicit = " (implicit)"
 		}
-		fmt.Fprintf(&sb, "[file %q%s]: %q\n", f, implicit, fs.m[f])
+		fmt.Fprintf(&sb, "[file %q%s]: %q\n", f, implicit, fs.Data(f))
 	}
 	return sb.String()
 }