loader: change extension handler type (open-policy-agent#6015)

This 1. changes the extension.Handler type to make it more flexible 2. simplifies the extension usage -- it used to be called in many places, but it could all be handled through util.Unmarshal and util.UnmarshalJSON instead We've previously marked it as "EXPERIMENTAL", so we should have enough leeway to change this now. NOTE: As a consequence of (2.), we're no longer accepting trailing data for json files loaded with OPA. I believe it wasn't intentional to ignore bad data before -- now, it'll be an error. Signed-off-by: Stephan Renatus <[email protected]>
danielbass37 · Jun 16, 2023 · cca8197 · cca8197
1 parent 6abcc29
commit cca8197
Show file tree

Hide file tree

Showing 14 changed files with 36 additions and 54 deletions.
diff --git a/bundle/bundle.go b/bundle/bundle.go
@@ -23,7 +23,6 @@ import (
 	"github.com/open-policy-agent/opa/format"
 	"github.com/open-policy-agent/opa/internal/file/archive"
 	"github.com/open-policy-agent/opa/internal/merge"
-	"github.com/open-policy-agent/opa/loader/extension"
 	"github.com/open-policy-agent/opa/metrics"
 	"github.com/open-policy-agent/opa/util"
 )
@@ -608,15 +607,10 @@ func (r *Reader) Read() (Bundle, error) {
 				continue
 			}
 
-			var err error
 			var value interface{}
 
 			r.metrics.Timer(metrics.RegoDataParse).Start()
-			if handler := extension.FindExtension(".json"); handler != nil {
-				value, err = handler(buf.Bytes())
-			} else {
-				err = util.NewJSONDecoder(&buf).Decode(&value)
-			}
+			err := util.UnmarshalJSON(buf.Bytes(), &value)
 			r.metrics.Timer(metrics.RegoDataParse).Stop()
 
 			if err != nil {

diff --git a/bundle/bundle_test.go b/bundle/bundle_test.go
@@ -127,7 +127,7 @@ func TestReadBundleInLazyMode(t *testing.T) {
 		{"/a/b/d/data.json", "true"},
 		{"/a/b/y/data.yaml", `foo: 1`},
 		{"/example/example.rego", `package example`},
-		{"/data.json", `{"x": {"y": true}, "a": {"b": {"z": true}}}}`},
+		{"/data.json", `{"x": {"y": true}, "a": {"b": {"z": true}}}`},
 		{"/.manifest", `{"revision": "foo", "roots": ["example"]}`}, // data is outside roots but validation skipped in lazy mode
 	}
 
@@ -196,7 +196,7 @@ func testReadBundle(t *testing.T, baseDir string, useMemoryFS bool) {
 		{"/example/example.rego", `package example`},
 		{"/policy.wasm", `legacy-wasm-module`},
 		{wasmResolverPath, `wasm-module`},
-		{"/data.json", `{"x": {"y": true}, "a": {"b": {"z": true}}}}`},
+		{"/data.json", `{"x": {"y": true}, "a": {"b": {"z": true}}}`},
 		{"/.manifest", fmt.Sprintf(`{"wasm":[{"entrypoint": "authz/allow", "module": "%s"}]}`, fullWasmResolverPath)},
 	}
 
@@ -665,7 +665,7 @@ func TestVerifyBundleFileHash(t *testing.T) {
 		{"/a/b/y/data.yaml", `foo: 1`},
 		{"/example/example.rego", `package example`},
 		{"/policy.wasm", `modules-compiled-as-wasm-binary`},
-		{"/data.json", `{"x": {"y": true}, "a": {"b": {"z": true}}}}`},
+		{"/data.json", `{"x": {"y": true}, "a": {"b": {"z": true}}}`},
 	}
 
 	buf := archive.MustWriteTarGz(files)

diff --git a/bundle/file_test.go b/bundle/file_test.go
@@ -83,7 +83,7 @@ func TestIteratorOrder(t *testing.T) {
 	var archFiles = map[string]string{
 		"/a/b/c/data.json":     "[1,2,3]",
 		"/a/b/d/e/data.json":   `e: true`,
-		"/data.json":           `{"x": {"y": true}, "a": {"b": {"z": true}}}}`,
+		"/data.json":           `{"x": {"y": true}, "a": {"b": {"z": true}}}`,
 		"/a/b/y/x/z/data.yaml": `foo: 1`,
 		"/a/b/data.json":       "[4,5,6]",
 		"/a/data.json":         "hello",

diff --git a/bundle/sign_test.go b/bundle/sign_test.go
@@ -27,7 +27,7 @@ func TestGenerateSignedToken(t *testing.T) {
 		{"/a/b/y/data.yaml", `foo: 1`},
 		{"/example/example.rego", `package example`},
 		{"/policy.wasm", `modules-compiled-as-wasm-binary`},
-		{"/data.json", `{"x": {"y": true}, "a": {"b": {"z": true}}}}`},
+		{"/data.json", `{"x": {"y": true}, "a": {"b": {"z": true}}}`},
 	}
 
 	input := []FileInfo{}

diff --git a/bundle/store_test.go b/bundle/store_test.go
@@ -307,7 +307,7 @@ func TestBundleLazyModeLifecycleRaw(t *testing.T) {
 		{"/a/b/y/data.yaml", `foo: 1`},
 		{"/example/example.rego", `package example`},
 		{"/authz/allow/policy.wasm", `wasm-module`},
-		{"/data.json", `{"x": {"y": true}, "a": {"b": {"z": true}}}}`},
+		{"/data.json", `{"x": {"y": true}, "a": {"b": {"z": true}}}`},
 		{"/.manifest", `{"revision": "foo", "roots": ["a", "example", "x", "authz"],"wasm":[{"entrypoint": "authz/allow", "module": "/authz/allow/policy.wasm"}]}`},
 	}
 
@@ -705,7 +705,7 @@ func TestBundleLazyModeLifecycleRawNoBundleRoots(t *testing.T) {
 		{"/a/b/d/data.json", "true"},
 		{"/a/b/y/data.yaml", `foo: 1`},
 		{"/example/example.rego", `package example`},
-		{"/data.json", `{"x": {"y": true}, "a": {"b": {"z": true}}}}`},
+		{"/data.json", `{"x": {"y": true}, "a": {"b": {"z": true}}}`},
 		{"/.manifest", `{"revision": "rev-1"}`},
 	}
 
@@ -903,7 +903,7 @@ func TestBundleLazyModeLifecycleRawNoBundleRootsDiskStorage(t *testing.T) {
 			{"/a/b/d/data.json", "true"},
 			{"/a/b/y/data.yaml", `foo: 1`},
 			{"/example/example.rego", `package example`},
-			{"/data.json", `{"x": {"y": true}, "a": {"b": {"z": true}}}}`},
+			{"/data.json", `{"x": {"y": true}, "a": {"b": {"z": true}}}`},
 			{"/.manifest", `{"revision": "rev-1"}`},
 		}
 

diff --git a/cmd/inspect_test.go b/cmd/inspect_test.go
@@ -23,7 +23,7 @@ import (
 func TestDoInspect(t *testing.T) {
 	files := [][2]string{
 		{"/.manifest", `{"revision": "rev", "roots": ["foo", "bar", "fuz", "baz", "a", "x"]}`},
-		{"/data.json", `{"x": {"y": true}, "a": {"b": {"z": true}}}}`},
+		{"/data.json", `{"x": {"y": true}, "a": {"b": {"z": true}}}`},
 		{"/example/foo.rego", `package foo`},
 	}
 
@@ -79,7 +79,7 @@ func TestDoInspectPretty(t *testing.T) {
 
 	files := [][2]string{
 		{"/.manifest", manifest},
-		{"/data.json", `{"x": {"y": true}, "a": {"b": {"z": true}}}}`},
+		{"/data.json", `{"x": {"y": true}, "a": {"b": {"z": true}}}`},
 		{"/http/example/authz/foo.rego", `package http.example.authz`},
 		{"/http/example/authz/data.json", `{"faz": "baz"}`},
 		{"/example/foo.rego", `package foo`},

diff --git a/cmd/sign_test.go b/cmd/sign_test.go
@@ -78,7 +78,7 @@ func TestBundleSignVerification(t *testing.T) {
 		"/a/b/y/data.yaml":      `foo: 1`,
 		"/example/example.rego": `package example`,
 		"/policy.wasm":          `modules-compiled-as-wasm-binary`,
-		"/data.json":            `{"x": {"y": true}, "a": {"b": {"z": true}}}}`,
+		"/data.json":            `{"x": {"y": true}, "a": {"b": {"z": true}}}`,
 	}
 
 	test.WithTempFS(files, func(rootDir string) {

diff --git a/internal/bundle/inspect/inspect_test.go b/internal/bundle/inspect/inspect_test.go
@@ -172,7 +172,7 @@ func TestGenerateBundleInfoWithBundleTarGz(t *testing.T) {
 		{"/example/example.rego", `package example`},
 		{"/example/policy.wasm", `modules-compiled-as-wasm-binary`},
 		{"/policy.wasm", `modules-compiled-as-wasm-binary`},
-		{"/data.json", `{"x": {"y": true}, "a": {"b": {"z": true}}}}`},
+		{"/data.json", `{"x": {"y": true}, "a": {"b": {"z": true}}}`},
 	}
 
 	buf := archive.MustWriteTarGz(files)

diff --git a/loader/extension/extension.go b/loader/extension/extension.go
@@ -14,7 +14,7 @@ var bundleExtensions map[string]Handler
 // Handler is used to unmarshal a byte slice of a registered extension
 // EXPERIMENTAL: Please don't rely on this functionality, it may go
 // away or change in the future.
-type Handler func([]byte) (any, error)
+type Handler func([]byte, any) error
 
 // RegisterExtension registers a Handler for a certain file extension, including
 // the dot: ".json", not "json".
@@ -36,11 +36,5 @@ func RegisterExtension(name string, handler Handler) {
 func FindExtension(ext string) Handler {
 	pluginMtx.Lock()
 	defer pluginMtx.Unlock()
-
-	for e, handler := range bundleExtensions {
-		if e == ext {
-			return handler
-		}
-	}
-	return nil
+	return bundleExtensions[ext]
 }
diff --git a/loader/extension/extension_test.go b/loader/extension/extension_test.go
@@ -18,8 +18,8 @@ import (
 
 func TestLoaderExtensionUnmarshal(t *testing.T) {
 	sentinelErr := fmt.Errorf("test handler called")
-	extension.RegisterExtension(".json", func([]byte) (any, error) {
-		return nil, sentinelErr
+	extension.RegisterExtension(".json", func([]byte, any) error {
+		return sentinelErr
 	})
 	defer extension.RegisterExtension(".json", nil)
 
@@ -36,8 +36,9 @@ func TestLoaderExtensionUnmarshal(t *testing.T) {
 
 func TestLoaderExtensionBundle(t *testing.T) {
 	data := map[string]any{"foo": "bar"}
-	extension.RegisterExtension(".json", func([]byte) (any, error) {
-		return data, nil
+	extension.RegisterExtension(".json", func(_ []byte, x any) error {
+		*(x.(*any)) = data
+		return nil
 	})
 	defer extension.RegisterExtension(".json", nil)
 

diff --git a/loader/loader.go b/loader/loader.go
@@ -21,7 +21,6 @@ import (
 	"github.com/open-policy-agent/opa/bundle"
 	fileurl "github.com/open-policy-agent/opa/internal/file/url"
 	"github.com/open-policy-agent/opa/internal/merge"
-	"github.com/open-policy-agent/opa/loader/extension"
 	"github.com/open-policy-agent/opa/loader/filter"
 	"github.com/open-policy-agent/opa/metrics"
 	"github.com/open-policy-agent/opa/storage"
@@ -750,14 +749,8 @@ func loadRego(path string, bs []byte, m metrics.Metrics, opts ast.ParserOptions)
 
 func loadJSON(path string, bs []byte, m metrics.Metrics) (interface{}, error) {
 	m.Timer(metrics.RegoDataParse).Start()
-	var err error
 	var x interface{}
-	if handler := extension.FindExtension(".json"); handler != nil {
-		x, err = handler(bs)
-	} else {
-		buf := bytes.NewBuffer(bs)
-		err = util.NewJSONDecoder(buf).Decode(&x)
-	}
+	err := util.UnmarshalJSON(bs, &x)
 	m.Timer(metrics.RegoDataParse).Stop()
 
 	if err != nil {

diff --git a/storage/disk/disk_test.go b/storage/disk/disk_test.go
@@ -198,7 +198,7 @@ func runTruncateTest(t *testing.T, dir string) {
 	var archiveFiles = map[string]string{
 		"/a/b/c/data.json":   "[1,2,3]",
 		"/a/b/d/data.json":   `e: true`,
-		"/data.json":         `{"x": {"y": true}, "a": {"b": {"z": true}}}}`,
+		"/data.json":         `{"x": {"y": true}, "a": {"b": {"z": true}}}`,
 		"/a/b/y/data.yaml":   `foo: 1`,
 		"/policy.rego":       "package foo\n p = 1",
 		"/roles/policy.rego": "package bar\n p = 1",
@@ -334,7 +334,7 @@ func TestTruncateMultipleTxn(t *testing.T) {
 		}
 
 		// additional data file at root
-		archiveFiles["/data.json"] = `{"a": {"b": {"z": true}}}}`
+		archiveFiles["/data.json"] = `{"a": {"b": {"z": true}}}`
 
 		files := make([][2]string, 0, len(archiveFiles))
 		for name, content := range archiveFiles {

diff --git a/storage/inmem/inmem_test.go b/storage/inmem/inmem_test.go
@@ -390,7 +390,7 @@ func TestTruncate(t *testing.T) {
 	var archiveFiles = map[string]string{
 		"/a/b/c/data.json":   "[1,2,3]",
 		"/a/b/d/data.json":   "true",
-		"/data.json":         `{"x": {"y": true}, "a": {"b": {"z": true}}}}`,
+		"/data.json":         `{"x": {"y": true}, "a": {"b": {"z": true}}}`,
 		"/a/b/y/data.yaml":   `foo: 1`,
 		"/policy.rego":       "package foo\n p = 1",
 		"/roles/policy.rego": "package bar\n p = 1",

diff --git a/util/json.go b/util/json.go
@@ -21,10 +21,17 @@ import (
 //
 // This function is intended to be used in place of the standard json.Marshal
 // function when json.Number is required.
-func UnmarshalJSON(bs []byte, x interface{}) (err error) {
+func UnmarshalJSON(bs []byte, x interface{}) error {
+	return unmarshalJSON(bs, x, true)
+}
+
+func unmarshalJSON(bs []byte, x interface{}, ext bool) error {
 	buf := bytes.NewBuffer(bs)
 	decoder := NewJSONDecoder(buf)
 	if err := decoder.Decode(x); err != nil {
+		if handler := extension.FindExtension(".json"); handler != nil && ext {
+			return handler(bs, x)
+		}
 		return err
 	}
 
@@ -108,22 +115,15 @@ func Reference(x interface{}) *interface{} {
 // Unmarshal decodes a YAML, JSON or JSON extension value into the specified type.
 func Unmarshal(bs []byte, v interface{}) error {
 	if json.Valid(bs) {
-		return UnmarshalJSON(bs, v)
+		return unmarshalJSON(bs, v, false)
 	}
 	nbs, err := yaml.YAMLToJSON(bs)
 	if err == nil {
-		return UnmarshalJSON(nbs, v)
+		return unmarshalJSON(nbs, v, false)
 	}
 	// not json or yaml: try extensions
-	if value, ok := v.(*any); ok {
-		if handler := extension.FindExtension(".json"); handler != nil {
-			retval, err := handler(bs)
-			if err != nil {
-				return err
-			}
-			*value = retval
-			return nil
-		}
+	if handler := extension.FindExtension(".json"); handler != nil {
+		return handler(bs, v)
 	}
 	return err
 }