From 40a57dac8a9763ccf6ce30e572c18acd448649e2 Mon Sep 17 00:00:00 2001
From: Connor Peet <connor@peet.io>
Date: Sun, 30 Oct 2016 21:20:44 -0700
Subject: [PATCH] Assert that no extraneous data remains after top-level tokens

---
 gen/decoder.go       | 11 +++++++++++
 jlexer/lexer.go      | 33 +++++++++++++++++++++++++++++++--
 jlexer/lexer_test.go | 24 ++++++++++++++++++++++++
 3 files changed, 66 insertions(+), 2 deletions(-)

diff --git a/gen/decoder.go b/gen/decoder.go
index 681584f..bafa162 100644
--- a/gen/decoder.go
+++ b/gen/decoder.go
@@ -340,10 +340,14 @@ func (g *Generator) genSliceArrayDecoder(t reflect.Type) error {
 	typ := g.getType(t)
 
 	fmt.Fprintln(g.out, "func "+fname+"(in *jlexer.Lexer, out *"+typ+") {")
+	fmt.Fprintln(g.out, " isTopLevel := in.IsStart()")
 	err := g.genTypeDecoderNoCheck(t, "*out", fieldTags{}, 1)
 	if err != nil {
 		return err
 	}
+	fmt.Fprintln(g.out, "  if isTopLevel {")
+	fmt.Fprintln(g.out, "    in.Consumed()")
+	fmt.Fprintln(g.out, "  }")
 	fmt.Fprintln(g.out, "}")
 
 	return nil
@@ -358,7 +362,11 @@ func (g *Generator) genStructDecoder(t reflect.Type) error {
 	typ := g.getType(t)
 
 	fmt.Fprintln(g.out, "func "+fname+"(in *jlexer.Lexer, out *"+typ+") {")
+	fmt.Fprintln(g.out, "  isTopLevel := in.IsStart()")
 	fmt.Fprintln(g.out, "  if in.IsNull() {")
+	fmt.Fprintln(g.out, "    if isTopLevel {")
+	fmt.Fprintln(g.out, "      in.Consumed()")
+	fmt.Fprintln(g.out, "    }")
 	fmt.Fprintln(g.out, "    in.Skip()")
 	fmt.Fprintln(g.out, "    return")
 	fmt.Fprintln(g.out, "  }")
@@ -404,6 +412,9 @@ func (g *Generator) genStructDecoder(t reflect.Type) error {
 	fmt.Fprintln(g.out, "    in.WantComma()")
 	fmt.Fprintln(g.out, "  }")
 	fmt.Fprintln(g.out, "  in.Delim('}')")
+	fmt.Fprintln(g.out, "  if isTopLevel {")
+	fmt.Fprintln(g.out, "    in.Consumed()")
+	fmt.Fprintln(g.out, "  }")
 
 	for _, f := range fs {
 		g.genRequiredFieldCheck(t, f)
diff --git a/jlexer/lexer.go b/jlexer/lexer.go
index c3d4ed2..eac6cf5 100644
--- a/jlexer/lexer.go
+++ b/jlexer/lexer.go
@@ -519,8 +519,9 @@ func (r *Lexer) SkipRecursive() {
 	r.err = &LexerError{
 		Reason: "EOF reached while skipping array/object or token",
 		Offset: r.pos,
-		Data: string(r.Data[r.pos:]),
-	}}
+		Data:   string(r.Data[r.pos:]),
+	}
+}
 
 // Raw fetches the next item recursively as a data slice
 func (r *Lexer) Raw() []byte {
@@ -531,6 +532,34 @@ func (r *Lexer) Raw() []byte {
 	return r.Data[r.start:r.pos]
 }
 
+// IsStart returns whether the lexer is positioned at the start
+// of an input string.
+func (r *Lexer) IsStart() bool {
+	return r.pos == 0
+}
+
+// Consumed reads all remaining bytes from the input, publishing an error if
+// there is anything but whitespace remaining.
+func (r *Lexer) Consumed() {
+	if r.pos > len(r.Data) {
+		return
+	}
+
+	for _, c := range r.Data[r.pos:] {
+		if c != ' ' && c != '\t' && c != '\r' && c != '\n' {
+			r.err = &LexerError{
+				Reason: "invalid character '" + string(c) + "' after top-level value",
+				Offset: r.pos,
+				Data:   string(r.Data[r.pos:]),
+			}
+			return
+		}
+
+		r.pos++
+		r.start++
+	}
+}
+
 // UnsafeString returns the string value if the token is a string literal.
 //
 // Warning: returned string may point to the input buffer, so the string should not outlive
diff --git a/jlexer/lexer_test.go b/jlexer/lexer_test.go
index 58f5fe9..e3add8c 100644
--- a/jlexer/lexer_test.go
+++ b/jlexer/lexer_test.go
@@ -222,3 +222,27 @@ func TestInterface(t *testing.T) {
 		}
 	}
 }
+
+func TestConsumed(t *testing.T) {
+	for i, test := range []struct {
+		toParse   string
+		wantError bool
+	}{
+		{toParse: "", wantError: false},
+		{toParse: "   ", wantError: false},
+		{toParse: "\r\n", wantError: false},
+		{toParse: "\t\t", wantError: false},
+
+		{toParse: "{", wantError: true},
+	} {
+		l := Lexer{Data: []byte(test.toParse)}
+		l.Consumed()
+
+		err := l.Error()
+		if err != nil && !test.wantError {
+			t.Errorf("[%d, %q] Consumed() error: %v", i, test.toParse, err)
+		} else if err == nil && test.wantError {
+			t.Errorf("[%d, %q] Consumed() ok; want error", i, test.toParse)
+		}
+	}
+}