util: Fix UnmarshalJSON parses invalid json in some cases

Fixes: #2331

Signed-off-by: katsew <y.katsew@gmail.com>

topdown/regex_test.go

@@ -39,13 +39,13 @@ func TestRegexFindAllStringSubmatch(t *testing.T) {
 		rules    []string
 		expected interface{}
 	}{
-		{"finds no matches", []string{`p[x] { x = regex.find_all_string_submatch_n("a(x*)b", "-", -1) }`}, `[[]]]`},
-		{"single match without captures", []string{`p[x] { x = regex.find_all_string_submatch_n("a(x*)b", "-ab-", -1) }`}, `[[["ab", ""]]]]`},
-		{"single match with a capture", []string{`p[x] { x = regex.find_all_string_submatch_n("a(x*)b", "-axxb-", -1) }`}, `[[["axxb", "xx"]]]]`},
-		{"multiple matches with captures-1", []string{`p[x] { x = regex.find_all_string_submatch_n("a(x*)b", "-ab-axb-", -1) }`}, `[[["ab", ""], ["axb", "x"]]]]`},
-		{"multiple matches with captures-2", []string{`p[x] { x = regex.find_all_string_submatch_n("a(x*)b", "-axxb-ab-", -1) }`}, `[[["axxb", "xx"], ["ab", ""]]]]`},
-		{"multiple patterns, matches, and captures", []string{`p[x] { x = regex.find_all_string_submatch_n("[^aouiye]([aouiye])([^aouiye])?", "somestri", -1) }`}, `[[["som", "o", "m"], ["ri", "i", ""]]]]`},
-		{"multiple patterns, matches, and captures with specified number of matches", []string{`p[x] { x = regex.find_all_string_submatch_n("[^aouiye]([aouiye])([^aouiye])?", "somestri", 1) }`}, `[[["som", "o", "m"]]]]`},
+		{"finds no matches", []string{`p[x] { x = regex.find_all_string_submatch_n("a(x*)b", "-", -1) }`}, `[[]]`},
+		{"single match without captures", []string{`p[x] { x = regex.find_all_string_submatch_n("a(x*)b", "-ab-", -1) }`}, `[[["ab", ""]]]`},
+		{"single match with a capture", []string{`p[x] { x = regex.find_all_string_submatch_n("a(x*)b", "-axxb-", -1) }`}, `[[["axxb", "xx"]]]`},
+		{"multiple matches with captures-1", []string{`p[x] { x = regex.find_all_string_submatch_n("a(x*)b", "-ab-axb-", -1) }`}, `[[["ab", ""], ["axb", "x"]]]`},
+		{"multiple matches with captures-2", []string{`p[x] { x = regex.find_all_string_submatch_n("a(x*)b", "-axxb-ab-", -1) }`}, `[[["axxb", "xx"], ["ab", ""]]]`},
+		{"multiple patterns, matches, and captures", []string{`p[x] { x = regex.find_all_string_submatch_n("[^aouiye]([aouiye])([^aouiye])?", "somestri", -1) }`}, `[[["som", "o", "m"], ["ri", "i", ""]]]`},
+		{"multiple patterns, matches, and captures with specified number of matches", []string{`p[x] { x = regex.find_all_string_submatch_n("[^aouiye]([aouiye])([^aouiye])?", "somestri", 1) }`}, `[[["som", "o", "m"]]]`},
 	}
 
 	for _, tc := range tests {

topdown/tokens_test.go

@@ -369,6 +369,7 @@ func TestTopDownJWTEncodeSignHeaderErrors(t *testing.T) {
 		runTopDownTestCase(t, data, tc.note, tc.rules, tc.expected)
 	}
 }
+
 func TestTopDownJWTEncodeSignRaw(t *testing.T) {
 
 	const examplePayload = `{"iss":"joe",` + "\r\n" + ` "exp":1300819380,` + "\r\n" + ` "http://example.com/is_root":true}`
@@ -414,7 +415,7 @@ func TestTopDownJWTEncodeSignRaw(t *testing.T) {
 			"`" + "" + "`",
 			"`" + symmetricKey + "`",
 
-			`"eyJ0eXAiOiJ0ZXh0L3BsYWluIiwNCiAiYWxnIjoiSFMyNTYifQ..sXoGQMWwM-SmX495-htA7kndgbkwz1PnqsDeY275gnI""`,
+			`"eyJ0eXAiOiJ0ZXh0L3BsYWluIiwNCiAiYWxnIjoiSFMyNTYifQ..sXoGQMWwM-SmX495-htA7kndgbkwz1PnqsDeY275gnI"`,
 			"",
 		},
 		{
@@ -423,7 +424,7 @@ func TestTopDownJWTEncodeSignRaw(t *testing.T) {
 			"`" + "e" + "`",
 			"`" + symmetricKey + "`",
 
-			`"eyJ0eXAiOiJ0ZXh0L3BsYWluIiwNCiAiYWxnIjoiSFMyNTYifQ.ZQ.oO8Vnc4Jv7-J231a1bEcQrgXfKbNW-kEvVY7BP1v5rM""`,
+			`"eyJ0eXAiOiJ0ZXh0L3BsYWluIiwNCiAiYWxnIjoiSFMyNTYifQ.ZQ.oO8Vnc4Jv7-J231a1bEcQrgXfKbNW-kEvVY7BP1v5rM"`,
 			"",
 		},
 		{

topdown/topdown_test.go

@@ -1414,7 +1414,7 @@ func TestTopDownJSONBuiltins(t *testing.T) {
 		expected interface{}
 	}{
 		{"marshal", []string{`p = x { json.marshal([{"foo": {1,2,3}}], x) }`}, `"[{\"foo\":[1,2,3]}]"`},
-		{"unmarshal", []string{`p = x { json.unmarshal("[{\"foo\":[1,2,3]}]", x) }`}, `[{"foo": [1,2,3]}]"`},
+		{"unmarshal", []string{`p = x { json.unmarshal("[{\"foo\":[1,2,3]}]", x) }`}, `[{"foo": [1,2,3]}]`},
 		{"unmarshal-non-string", []string{`p = x { json.unmarshal(data.a[0], x) }`}, &Error{Code: TypeErr, Message: "operand 1 must be string but got number"}},
 		{"yaml round-trip", []string{`p = y { yaml.marshal([{"foo": {1,2,3}}], x); yaml.unmarshal(x, y) }`}, `[{"foo": [1,2,3]}]`},
 		{"yaml unmarshal error", []string{`p { yaml.unmarshal("[1,2,3", _) } `}, &Error{Code: BuiltinErr, Message: "yaml: line 1: did not find"}},

util/json.go

@@ -7,6 +7,7 @@ package util
 import (
 	"bytes"
 	"encoding/json"
+	"fmt"
 	"io"
 	"reflect"
 
@@ -21,7 +22,20 @@ import (
 func UnmarshalJSON(bs []byte, x interface{}) (err error) {
 	buf := bytes.NewBuffer(bs)
 	decoder := NewJSONDecoder(buf)
-	return decoder.Decode(x)
+	if err := decoder.Decode(x); err != nil {
+		return err
+	}
+
+	// Since decoder.Decode validates only the first json structure in bytes,
+	// check if decoder has more bytes to consume to validate whole input bytes.
+	tok, err := decoder.Token()
+	if tok != nil {
+		return fmt.Errorf("error: invalid character '%s' after top-level value", tok)
+	}
+	if err != nil && err != io.EOF {
+		return err
+	}
+	return nil
 }
 
 // NewJSONDecoder returns a new decoder that reads from r.

util/json_test.go

@@ -13,6 +13,20 @@ import (
 	"github.com/open-policy-agent/opa/util"
 )
 
+func TestInvalidJSONInput(t *testing.T) {
+	cases := [][]byte{
+		[]byte("{ \"k\": 1 }\n{}}"),
+		[]byte("{ \"k\": 1 }\n!!!}"),
+	}
+	for _, tc := range cases {
+		var x interface{}
+		err := util.UnmarshalJSON(tc, &x)
+		if err == nil {
+			t.Errorf("should be an error")
+		}
+	}
+}
+
 func TestRoundTrip(t *testing.T) {
 	cases := []interface{}{
 		nil,