From 96b5a5404f303f074e6117d832a9873c439508f0 Mon Sep 17 00:00:00 2001
From: Steve Gilbert <stevegilbert23@gmail.com>
Date: Fri, 15 Dec 2017 08:15:17 -0500
Subject: [PATCH] imports: support files with longer lines

This change set fixes the issue by specifying a max buffer size of ~1 megabyte for Scanner.
Previously the max was not set resulting in the default max which is ~64k.  This increased max should not increase
memory use for smaller, normal files because it is a max and the code expands the buffer as needed for large tokens.

The change set includes an additional change to return an error from the Scanner which was ignored previously.

Fixes golang/go#18201

Change-Id: I11be39af74d5eb3b353ad81ba1cb5404207aa65d
Reviewed-on: https://go-review.googlesource.com/83800
Run-TryBot: Brad Fitzpatrick <bradfitz@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Brad Fitzpatrick <bradfitz@golang.org>
---
 imports/fix_test.go | 93 +++++++++++++++++++++++++++++++++++++++++++++
 imports/imports.go  | 14 +++++--
 2 files changed, 104 insertions(+), 3 deletions(-)

diff --git a/imports/fix_test.go b/imports/fix_test.go
index 2197bafd..d39f28ff 100644
--- a/imports/fix_test.go
+++ b/imports/fix_test.go
@@ -5,6 +5,7 @@
 package imports
 
 import (
+	"bufio"
 	"bytes"
 	"flag"
 	"go/build"
@@ -1957,3 +1958,95 @@ var _ = &bytes.Buffer{}
 		}
 	})
 }
+
+// A happy path test for Process
+func TestProcess(t *testing.T) {
+	in := `package testimports
+
+	var s = fmt.Sprintf("%s", "value")
+`
+	out, err := Process("foo", []byte(in), nil)
+
+	if err != nil {
+		t.Errorf("Process returned error.\n got:\n%v\nwant:\nnil", err)
+	}
+
+	want := `package testimports
+
+import "fmt"
+
+var s = fmt.Sprintf("%s", "value")
+`
+	if got := string(out); got != want {
+		t.Errorf("Process returned unexpected result.\ngot:\n%v\nwant:\n%v", got, want)
+	}
+}
+
+// Ensures a token as large as 500000 bytes can be handled
+// https://golang.org/issues/18201
+func TestProcessLargeToken(t *testing.T) {
+	largeString := strings.Repeat("x", 500000)
+
+	in := `package testimports
+
+import (
+	"fmt"
+	"mydomain.mystuff/mypkg"
+)
+
+const s = fmt.Sprintf("%s", "` + largeString + `")
+const x = mypkg.Sprintf("%s", "my package")
+
+// end
+`
+
+	out, err := Process("foo", []byte(in), nil)
+
+	if err != nil {
+		t.Errorf("Process returned error.\n got:\n%v\nwant:\nnil", err)
+	}
+
+	want := `package testimports
+
+import (
+	"fmt"
+
+	"mydomain.mystuff/mypkg"
+)
+
+const s = fmt.Sprintf("%s", "` + largeString + `")
+const x = mypkg.Sprintf("%s", "my package")
+
+// end
+`
+
+	if got := string(out); got != want {
+		t.Errorf("Process returned unexpected result.\ngot:\n%v\nwant:\n%v", got, want)
+	}
+}
+
+// Ensures a token that is larger that
+// https://golang.org/issues/18201
+func TestProcessTokenTooLarge(t *testing.T) {
+	const largeSize = maxScanTokenSize + 1
+	largeString := strings.Repeat("x", largeSize)
+
+	in := `package testimports
+
+import (
+	"fmt"
+	"mydomain.mystuff/mypkg"
+)
+
+const s = fmt.Sprintf("%s", "` + largeString + `")
+const x = mypkg.Sprintf("%s", "my package")
+
+// end
+`
+
+	_, err := Process("foo", []byte(in), nil)
+
+	if err != bufio.ErrTooLong {
+		t.Errorf("Process did not returned expected error.\n got:\n%v\nwant:\n%v", err, bufio.ErrTooLong)
+	}
+}
diff --git a/imports/imports.go b/imports/imports.go
index 67573f49..d789bb5b 100644
--- a/imports/imports.go
+++ b/imports/imports.go
@@ -98,7 +98,10 @@ func Process(filename string, src []byte, opt *Options) ([]byte, error) {
 		out = adjust(src, out)
 	}
 	if len(spacesBefore) > 0 {
-		out = addImportSpaces(bytes.NewReader(out), spacesBefore)
+		out, err = addImportSpaces(bytes.NewReader(out), spacesBefore)
+		if err != nil {
+			return nil, err
+		}
 	}
 
 	out, err = format.Source(out)
@@ -256,9 +259,14 @@ func matchSpace(orig []byte, src []byte) []byte {
 
 var impLine = regexp.MustCompile(`^\s+(?:[\w\.]+\s+)?"(.+)"`)
 
-func addImportSpaces(r io.Reader, breaks []string) []byte {
+// Used to set Scanner buffer size so that large tokens can be handled.
+// see https://github.com/golang/go/issues/18201
+const maxScanTokenSize = bufio.MaxScanTokenSize * 16
+
+func addImportSpaces(r io.Reader, breaks []string) ([]byte, error) {
 	var out bytes.Buffer
 	sc := bufio.NewScanner(r)
+	sc.Buffer(nil, maxScanTokenSize)
 	inImports := false
 	done := false
 	for sc.Scan() {
@@ -285,5 +293,5 @@ func addImportSpaces(r io.Reader, breaks []string) []byte {
 
 		fmt.Fprintln(&out, s)
 	}
-	return out.Bytes()
+	return out.Bytes(), sc.Err()
 }