From f6fa2af7b682531910b37a74b70319dc8ad6c363 Mon Sep 17 00:00:00 2001
From: Edward Paget <edward.paget@chime.com>
Date: Tue, 28 Sep 2021 12:32:40 -0500
Subject: [PATCH] sdk: Allow map of plugins to be passed to SDK

This adds a Plugins field to the SDK Options struct that accepts a
map[string]plugins.Factory similarly to the map created by
runtime.RegisterPlugin this map is passed to the discovery.Factories
plugin similarly, again, to how it works within the Runtime module. This
allows an OPA SDK instance to initialize plugins.

Fixes #3826

Signed-off-by: Edward Paget <edward.paget@chime.com>

Fixed changes requested by @tsandall (squash before merge)

Signed-off-by: Edward Paget <edward.paget@chime.com>
---
 sdk/opa.go      |  4 +++-
 sdk/opa_test.go | 53 +++++++++++++++++++++++++++++++++++++++++++++++++
 sdk/options.go  |  5 +++++
 3 files changed, 61 insertions(+), 1 deletion(-)

diff --git a/sdk/opa.go b/sdk/opa.go
index b2052baa33..64542d527c 100644
--- a/sdk/opa.go
+++ b/sdk/opa.go
@@ -37,6 +37,7 @@ type OPA struct {
 	mtx     sync.Mutex
 	logger  logging.Logger
 	console logging.Logger
+	plugins map[string]plugins.Factory
 	config  []byte
 }
 
@@ -69,6 +70,7 @@ func New(ctx context.Context, opts Options) (*OPA, error) {
 	opa.config = opts.config
 	opa.logger = opts.Logger
 	opa.console = opts.ConsoleLogger
+	opa.plugins = opts.Plugins
 
 	return opa, opa.configure(ctx, opa.config, opts.Ready, opts.block)
 }
@@ -138,7 +140,7 @@ func (opa *OPA) configure(ctx context.Context, bs []byte, ready chan struct{}, b
 		close(ready)
 	})
 
-	d, err := discovery.New(manager)
+	d, err := discovery.New(manager, discovery.Factories(opa.plugins))
 	if err != nil {
 		return err
 	}
diff --git a/sdk/opa_test.go b/sdk/opa_test.go
index d6e0ea91ce..9c1cd64eb3 100644
--- a/sdk/opa_test.go
+++ b/sdk/opa_test.go
@@ -16,10 +16,63 @@ import (
 
 	"github.com/fortytw2/leaktest"
 	loggingtest "github.com/open-policy-agent/opa/logging/test"
+	"github.com/open-policy-agent/opa/plugins"
 	"github.com/open-policy-agent/opa/sdk"
 	sdktest "github.com/open-policy-agent/opa/sdk/test"
 )
 
+// Plugin creates an empty plugin to test plugin initialization
+type plugin struct {
+	manager *plugins.Manager
+}
+
+type factory struct{}
+
+func (p *plugin) Start(ctx context.Context) error {
+	p.manager.UpdatePluginStatus("test_plugin", &plugins.Status{State: plugins.StateOK})
+	return nil
+}
+
+func (p *plugin) Stop(ctx context.Context) {
+}
+
+func (p *plugin) Reconfigure(ctx context.Context, config interface{}) {
+}
+
+func (factory) New(manager *plugins.Manager, config interface{}) plugins.Plugin {
+	return &plugin{
+		manager: manager,
+	}
+}
+
+func (factory) Validate(manager *plugins.Manager, config []byte) (interface{}, error) {
+	return nil, nil
+}
+
+func TestPlugins(t *testing.T) {
+
+	ctx := context.Background()
+
+	config := []byte(`{
+        "plugins": {
+            "test_plugin": {}
+		}
+	}`)
+
+	opa, err := sdk.New(ctx, sdk.Options{
+		Config: bytes.NewReader(config),
+		Plugins: map[string]plugins.Factory{
+			"test_plugin": factory{},
+		},
+	})
+
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	defer opa.Stop(ctx)
+}
+
 func TestDecision(t *testing.T) {
 
 	ctx := context.Background()
diff --git a/sdk/options.go b/sdk/options.go
index 160f94f195..f768fb1f0a 100644
--- a/sdk/options.go
+++ b/sdk/options.go
@@ -11,6 +11,7 @@ import (
 	"github.com/sirupsen/logrus"
 
 	"github.com/open-policy-agent/opa/logging"
+	"github.com/open-policy-agent/opa/plugins"
 )
 
 // Options contains parameters to setup and configure OPA.
@@ -35,6 +36,10 @@ type Options struct {
 	// is closed to signal readiness.
 	Ready chan struct{}
 
+	// Plugins provides a set of plugins.Factory instances that will be
+	// registered with the OPA SDK instance.
+	Plugins map[string]plugins.Factory
+
 	config []byte
 	block  bool
 }