Add support to enable ND builtin cache via discovery (#5468)

This commit adds support for enabling/disabling the ND builtin
cache via the OPA `discovery` plugin.

Fixes: #5457

Signed-off-by: Ashutosh Narkar <anarkar4387@gmail.com>
Signed-off-by: Philip Conrad <philipaconrad@gmail.com>
Co-authored-by: Philip Conrad <philipaconrad@gmail.com>

plugins/discovery/discovery_test.go

@@ -1537,6 +1537,70 @@ func TestInterQueryBuiltinCacheConfigUpdate(t *testing.T) {
 	}
 }
 
+func TestNDBuiltinCacheConfigUpdate(t *testing.T) {
+	type exampleConfig struct {
+		v bool
+	}
+	var config1 *exampleConfig
+	var config2 *exampleConfig
+	manager, err := plugins.New([]byte(`{
+		"discovery": {"name": "config"},
+		"services": {
+			"localhost": {
+				"url": "http://localhost:9999"
+			}
+		},
+  }`), "test-id", inmem.New())
+	manager.RegisterNDCacheTrigger(func(x bool) {
+		if config1 == nil {
+			config1 = &exampleConfig{v: x}
+		} else if config2 == nil {
+			config2 = &exampleConfig{v: x}
+		} else {
+			t.Fatal("Expected cache trigger to only be called twice")
+		}
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	testPlugin := &reconfigureTestPlugin{counts: map[string]int{}}
+	testFactory := testFactory{p: testPlugin}
+
+	disco, err := New(manager, Factories(map[string]plugins.Factory{"test_plugin": testFactory}))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	ctx := context.Background()
+
+	initialBundle := makeDataBundle(1, `{
+		"config": {
+			"nd_builtin_cache": true
+		}
+	}`)
+
+	disco.oneShot(ctx, download.Update{Bundle: initialBundle})
+
+	// Verify NDBuiltinCache is triggered with initial config
+	if config1 == nil || config1.v != true {
+		t.Fatalf("Expected ND builtin cache to be enabled after initial discovery, got: %v", config1.v)
+	}
+
+	// Verify NDBuiltinCache is reconfigured
+	updatedBundle := makeDataBundle(2, `{
+		"config": {
+			"nd_builtin_cache": false
+		}
+	}`)
+
+	disco.oneShot(ctx, download.Update{Bundle: updatedBundle})
+
+	if config2 == nil || config2.v != false {
+		t.Fatalf("Expected ND builtin cache to be disabled after discovery reconfigure, got: %v", config2.v)
+	}
+}
+
 func TestPluginManualTriggerLifecycle(t *testing.T) {
 	ctx := context.Background()
 	m := metrics.New()

plugins/plugins.go

@@ -192,6 +192,7 @@ type Manager struct {
 	router                       *mux.Router
 	prometheusRegister           prometheus.Registerer
 	tracerProvider               *trace.TracerProvider
+	registeredNDCacheTriggers    []func(bool)
 }
 
 type managerContextKey string
@@ -682,6 +683,10 @@ func (m *Manager) Reconfigure(config *config.Config) error {
 		trigger(interQueryBuiltinCacheConfig)
 	}
 
+	for _, trigger := range m.registeredNDCacheTriggers {
+		trigger(config.NDBuiltinCache)
+	}
+
 	return nil
 }
 
@@ -923,3 +928,9 @@ func (m *Manager) PrometheusRegister() prometheus.Registerer {
 func (m *Manager) TracerProvider() *trace.TracerProvider {
 	return m.tracerProvider
 }
+
+func (m *Manager) RegisterNDCacheTrigger(trigger func(bool)) {
+	m.mtx.Lock()
+	defer m.mtx.Unlock()
+	m.registeredNDCacheTriggers = append(m.registeredNDCacheTriggers, trigger)
+}

plugins/plugins_test.go

@@ -59,6 +59,44 @@ func TestManagerCacheTriggers(t *testing.T) {
 	}
 }
 
+func TestManagerNDCacheTriggers(t *testing.T) {
+	m, err := New([]byte{}, "test", inmem.New())
+	if err != nil {
+		t.Fatalf("Unexpected error: %s", err)
+	}
+
+	l1Called := false
+	m.RegisterNDCacheTrigger(func(bool) {
+		l1Called = true
+	})
+
+	if m.registeredNDCacheTriggers[0] == nil {
+		t.Fatal("First listener failed to register")
+	}
+
+	l2Called := false
+	m.RegisterNDCacheTrigger(func(bool) {
+		l2Called = true
+	})
+
+	if m.registeredNDCacheTriggers[0] == nil || m.registeredNDCacheTriggers[1] == nil {
+		t.Fatal("Second listener failed to register")
+	}
+
+	if l1Called == true || l2Called == true {
+		t.Fatal("Listeners should not be called yet")
+	}
+
+	err = m.Reconfigure(m.Config)
+	if err != nil {
+		t.Fatalf("Unexpected error: %s", err)
+	}
+
+	if l1Called == false || l2Called == false {
+		t.Fatal("Listeners should hav been called")
+	}
+}
+
 func TestManagerPluginStatusListener(t *testing.T) {
 	m, err := New([]byte{}, "test", inmem.New())
 	if err != nil {
@@ -257,6 +295,24 @@ func TestManagerWithCachingConfig(t *testing.T) {
 	}
 }
 
+func TestManagerWithNDCachingConfig(t *testing.T) {
+	m, err := New([]byte(`{"nd_builtin_cache": true}`), "test", inmem.New())
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	expected := true
+	if !m.Config.NDBuiltinCache == expected {
+		t.Fatalf("want %+v got %+v", expected, m.Config.NDBuiltinCache)
+	}
+
+	// config error
+	_, err = New([]byte(`{"nd_builtin_cache": "x"}`), "test", inmem.New())
+	if err == nil {
+		t.Fatal("expected error but got nil")
+	}
+}
+
 type mockForInitStartOrdering struct {
 	Manager *Manager
 	Started bool

server/server.go

@@ -181,6 +181,7 @@ func (s *Server) Init(ctx context.Context) (*Server, error) {
 	s.defaultDecisionPath = s.generateDefaultDecisionPath()
 	s.interQueryBuiltinCache = iCache.NewInterQueryCache(s.manager.InterQueryBuiltinCacheConfig())
 	s.manager.RegisterCacheTrigger(s.updateCacheConfig)
+	s.manager.RegisterNDCacheTrigger(s.updateNDCache)
 
 	// authorizer, if configured, needs the iCache to be set up already
 	s.Handler = s.initHandlerAuth(s.Handler)
@@ -2615,6 +2616,12 @@ func (s *Server) updateCacheConfig(cacheConfig *iCache.Config) {
 	s.interQueryBuiltinCache.UpdateConfig(cacheConfig)
 }
 
+func (s *Server) updateNDCache(enabled bool) {
+	s.mtx.Lock()
+	defer s.mtx.Unlock()
+	s.ndbCacheEnabled = enabled
+}
+
 func stringPathToDataRef(s string) (r ast.Ref) {
 	result := ast.Ref{ast.DefaultRootDocument}
 	result = append(result, stringPathToRef(s)...)