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() }