fix: resolve concurrent map writes for audit cache

internal/audit/agent.go

@@ -21,14 +21,20 @@ type Labels = map[string]any
 type Agent interface {
 	Trace(context.Context, Audit)
 	Cache() AuditCache
-	SetGlobalLabels(labels Labels)
+}
+
+type AgentPool interface {
+	New() Agent
 }
 
 // noopAgent is a lazy agent that does nothing :(
-type noopAgent struct{}
+type noopAgentPool struct{}
 
-func NewNoopAgent() Agent { return &noopAgent{} }
+func (p *noopAgentPool) New() Agent { return &noopAgent{} }
+
+func NewNoopAgentPool() AgentPool { return &noopAgentPool{} }
+
+type noopAgent struct{}
 
-func (a *noopAgent) Trace(context.Context, Audit)  {}
-func (a *noopAgent) Cache() AuditCache             { return &SingleRecordCache{} }
-func (a *noopAgent) SetGlobalLabels(labels Labels) {}
+func (a *noopAgent) Trace(context.Context, Audit) {}
+func (a *noopAgent) Cache() AuditCache            { return &SingleRecordCache{} }

internal/audit/agent_log.go

@@ -21,30 +21,40 @@ import (
 	"github.com/rond-authz/rond/logging"
 )
 
-type logAgent struct {
-	l       logging.Logger
-	cache   AuditCache
-	globals map[string]any
+type agentPool struct {
+	l      logging.Logger
+	labels Labels
 }
 
-func NewLogAgent(l logging.Logger) Agent {
-	return &logAgent{
-		l:     l,
+func (p *agentPool) New() Agent {
+	agent := &logAgent{
+		l:     p.l,
 		cache: &SingleRecordCache{},
 	}
+
+	if p.labels != nil {
+		agent.cache.Store(p.labels)
+	}
+
+	return agent
 }
-func (a *logAgent) SetGlobalLabels(labels Labels) {
-	a.globals = labels
+
+func NewLogAgentPool(l logging.Logger, labels Labels) AgentPool {
+	return &agentPool{
+		l:      l,
+		labels: labels,
+	}
+}
+
+type logAgent struct {
+	l     logging.Logger
+	cache AuditCache
 }
 
 func (a *logAgent) Trace(_ context.Context, auditInput Audit) {
 	data := a.cache.Load()
 
 	auditData := auditInput.toPrint()
-	if a.globals != nil {
-		auditData.applyDataFromPolicy(a.globals)
-	}
-
 	if data != nil {
 		auditData.applyDataFromPolicy(data)
 	}

internal/audit/agent_log_test.go

@@ -24,9 +24,11 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-func TestLogAgent(t *testing.T) {
+func TestLogAgentPool(t *testing.T) {
 	l, hook := test.NewNullLogger()
-	agent := NewLogAgent(rondlogrus.NewLogger(l))
+	pool := NewLogAgentPool(rondlogrus.NewLogger(l), nil)
+
+	agent := pool.New()
 
 	agent.Cache().Store(Data{
 		"authorization.permission": "my-permission",
@@ -83,11 +85,13 @@ func TestLogAgent(t *testing.T) {
 
 func TestLogAgentSetGlobalLabels(t *testing.T) {
 	l, hook := test.NewNullLogger()
-	agent := NewLogAgent(rondlogrus.NewLogger(l))
-	agent.SetGlobalLabels(Labels{
+	pool := NewLogAgentPool(rondlogrus.NewLogger(l), Labels{
 		AuditAdditionalDataRequestTargetServiceKey: "some-service",
 		"some-label": "label_val",
 	})
+
+	agent := pool.New()
+
 	agent.Cache().Store(Data{
 		"authorization.permission": "my-permission",
 		"authorization.binding":    "my-binding",

internal/audit/agent_test.go

@@ -21,11 +21,12 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-func TestNoopAgentDoesnotBreakStuff(t *testing.T) {
-	a := NewNoopAgent()
-	require.NotNil(t, a)
+func TestNoopAgentDoesNotBreakStuff(t *testing.T) {
+	p := NewNoopAgentPool()
+	require.NotNil(t, p)
+
+	a := p.New()
 
-	a.SetGlobalLabels(map[string]any{"a": "b"})
 	a.Trace(context.Background(), Audit{})
 
 	require.NotNil(t, a.Cache())

internal/audit/cache.go

@@ -14,6 +14,8 @@
 
 package audit
 
+import "sync"
+
 type Data map[string]any
 
 type AuditCache interface {
@@ -22,10 +24,15 @@ type AuditCache interface {
 }
 
 type SingleRecordCache struct {
+	sync.RWMutex
+
 	data Data
 }
 
 func (c *SingleRecordCache) Store(d Data) {
+	c.Lock()
+	defer c.Unlock()
+
 	if c.data == nil {
 		c.data = make(Data)
 	}
@@ -36,5 +43,8 @@ func (c *SingleRecordCache) Store(d Data) {
 }
 
 func (c *SingleRecordCache) Load() Data {
+	c.RLock()
+	defer c.RUnlock()
+
 	return c.data
 }

logging/test/logger.go

@@ -42,8 +42,8 @@ type testLogger struct {
 }
 
 func (l *testLogger) setRecord(level string, msg any) {
-	l.mu.RLock()
-	defer l.mu.RUnlock()
+	l.mu.Lock()
+	defer l.mu.Unlock()
 
 	l.entry.records = append(l.entry.records, Record{
 		Fields:  l.Fields,
@@ -52,8 +52,8 @@ func (l *testLogger) setRecord(level string, msg any) {
 	})
 
 	if originalLogger := l.entry.originalLogger; originalLogger != nil {
-		originalLogger.mu.RLock()
-		defer originalLogger.mu.RUnlock()
+		originalLogger.mu.Lock()
+		defer originalLogger.mu.Unlock()
 
 		originalLogger.entry.records = append(originalLogger.entry.records, Record{
 			Fields:  l.Fields,
@@ -84,8 +84,8 @@ func (l *testLogger) Trace(msg any) {
 }
 
 func (l *testLogger) WithFields(fields map[string]any) logging.Logger {
-	l.mu.RLock()
-	defer l.mu.RUnlock()
+	l.mu.Lock()
+	defer l.mu.Unlock()
 
 	clonedFields := map[string]any{}
 	for k, v := range l.Fields {
@@ -115,8 +115,8 @@ func (l *testLogger) WithFields(fields map[string]any) logging.Logger {
 }
 
 func (l *testLogger) WithField(key string, value any) logging.Logger {
-	l.mu.RLock()
-	defer l.mu.RUnlock()
+	l.mu.Lock()
+	defer l.mu.Unlock()
 
 	clonedFields := map[string]any{}
 	for k, v := range l.Fields {
@@ -144,14 +144,14 @@ func (l *testLogger) WithField(key string, value any) logging.Logger {
 }
 
 func (e *entry) AllRecords() []Record {
-	e.mu.Lock()
-	defer e.mu.Unlock()
+	e.mu.RLock()
+	defer e.mu.RUnlock()
 	return e.records
 }
 
 func (l *testLogger) OriginalLogger() *entry {
-	l.mu.Lock()
-	defer l.mu.Unlock()
+	l.mu.RLock()
+	defer l.mu.RUnlock()
 	return l.entry
 }
 

sdk/audit_factory.go

@@ -19,13 +19,13 @@ import (
 	"github.com/rond-authz/rond/logging"
 )
 
-func buildAuditAgent(options *Options, logger logging.Logger) audit.Agent {
+func buildAuditAgent(options *Options, logger logging.Logger) audit.AgentPool {
 	if options == nil || options.EvaluatorOptions == nil {
-		return audit.NewNoopAgent()
+		return audit.NewNoopAgentPool()
 	}
 
 	if !options.EvaluatorOptions.EnableAuditTracing {
-		return audit.NewNoopAgent()
+		return audit.NewNoopAgentPool()
 	}
-	return audit.NewLogAgent(logger)
+	return audit.NewLogAgentPool(logger, options.AuditLabels)
 }

sdk/audit_factory_test.go

@@ -27,32 +27,42 @@ import (
 )
 
 func TestBuildAuditAgent(t *testing.T) {
-	res := buildAuditAgent(nil, nil)
-	require.NotNil(t, res)
-
-	res = buildAuditAgent(&Options{
-		EvaluatorOptions: nil,
-	}, nil)
-	require.NotNil(t, res)
-
-	res = buildAuditAgent(&Options{
-		EvaluatorOptions: &EvaluatorOptions{
-			EnableAuditTracing: false,
-		},
-	}, nil)
-	require.NotNil(t, res)
-
-	log, hook := test.NewNullLogger()
-	res = buildAuditAgent(&Options{
-		EvaluatorOptions: &EvaluatorOptions{
-			EnableAuditTracing: true,
-		},
-	}, rondlogrus.NewEntry(logrus.NewEntry(log)))
-	require.NotNil(t, res)
-
-	res.Cache().Store(audit.Data{"test": "123"})
-	res.Trace(context.Background(), audit.Audit{
-		AggregationID: "a-id",
+	t.Run("nil options", func(t *testing.T) {
+		res := buildAuditAgent(nil, nil)
+		require.NotNil(t, res)
+	})
+
+	t.Run("nil evaluator options", func(t *testing.T) {
+		res := buildAuditAgent(&Options{
+			EvaluatorOptions: nil,
+		}, nil)
+		require.NotNil(t, res)
+	})
+
+	t.Run("nil logger", func(t *testing.T) {
+		res := buildAuditAgent(&Options{
+			EvaluatorOptions: &EvaluatorOptions{
+				EnableAuditTracing: false,
+			},
+		}, nil)
+		require.NotNil(t, res)
+	})
+
+	t.Run("with setup trace", func(t *testing.T) {
+		log, hook := test.NewNullLogger()
+		res := buildAuditAgent(&Options{
+			EvaluatorOptions: &EvaluatorOptions{
+				EnableAuditTracing: true,
+			},
+		}, rondlogrus.NewEntry(logrus.NewEntry(log)))
+		require.NotNil(t, res)
+
+		agent := res.New()
+
+		agent.Cache().Store(audit.Data{"test": "123"})
+		agent.Trace(context.Background(), audit.Audit{
+			AggregationID: "a-id",
+		})
+		require.Len(t, hook.Entries, 1)
 	})
-	require.Len(t, hook.Entries, 1)
 }

sdk/evaluator.go

@@ -52,7 +52,7 @@ type evaluator struct {
 
 	evaluatorOptions        *EvaluatorOptions
 	policyEvaluationOptions *core.PolicyEvaluationOptions
-	auditAgent              audit.Agent
+	auditAgentPool          audit.AgentPool
 }
 
 func (e evaluator) Config() core.RondConfig {
@@ -91,7 +91,8 @@ func (e evaluator) EvaluateRequestPolicy(ctx context.Context, rondInput core.Inp
 
 	opaEvaluatorOptions := e.evaluatorOptions.opaEvaluatorOptions(logger)
 
-	ctx = audit.WithAuditCache(ctx, e.auditAgent)
+	auditAgent := e.auditAgentPool.New()
+	ctx = audit.WithAuditCache(ctx, auditAgent)
 	evaluatorAllowPolicy, err := e.partialResultEvaluators.GetEvaluatorFromPolicy(ctx, rondConfig.RequestFlow.PolicyName, opaEvaluatorOptions)
 	if err != nil {
 		return PolicyResult{}, err
@@ -108,7 +109,7 @@ func (e evaluator) EvaluateRequestPolicy(ctx context.Context, rondInput core.Inp
 			"message":    err.Error(),
 		}).Error("RBAC policy evaluation failed")
 
-		e.auditAgent.Trace(ctx, audit.Audit{
+		auditAgent.Trace(ctx, audit.Audit{
 			AggregationID: options.Audit.AggregationID,
 			Authorization: audit.AuthzInfo{
 				Allowed:    false,
@@ -139,7 +140,7 @@ func (e evaluator) EvaluateRequestPolicy(ctx context.Context, rondInput core.Inp
 		}
 	}
 
-	e.auditAgent.Trace(ctx, audit.Audit{
+	auditAgent.Trace(ctx, audit.Audit{
 		AggregationID: options.Audit.AggregationID,
 		Authorization: audit.AuthzInfo{
 			Allowed:    true,
@@ -178,15 +179,16 @@ func (e evaluator) EvaluateResponsePolicy(ctx context.Context, rondInput core.In
 
 	opaEvaluatorOptions := e.evaluatorOptions.opaEvaluatorOptions(logger)
 
-	ctx = audit.WithAuditCache(ctx, e.auditAgent)
+	auditAgent := e.auditAgentPool.New()
+	ctx = audit.WithAuditCache(ctx, auditAgent)
 	evaluator, err := e.partialResultEvaluators.GetEvaluatorFromPolicy(ctx, e.rondConfig.ResponseFlow.PolicyName, opaEvaluatorOptions)
 	if err != nil {
 		return nil, err
 	}
 
 	bodyToProxy, err := evaluator.Evaluate(logger, evalInput, e.policyEvaluationOptions)
 	if err != nil {
-		e.auditAgent.Trace(ctx, audit.Audit{
+		auditAgent.Trace(ctx, audit.Audit{
 			AggregationID: options.Audit.AggregationID,
 			Authorization: audit.AuthzInfo{
 				Allowed:    false,
@@ -205,7 +207,7 @@ func (e evaluator) EvaluateResponsePolicy(ctx context.Context, rondInput core.In
 		return nil, err
 	}
 
-	e.auditAgent.Trace(ctx, audit.Audit{
+	auditAgent.Trace(ctx, audit.Audit{
 		AggregationID: options.Audit.AggregationID,
 		Authorization: audit.AuthzInfo{
 			Allowed:    true,