Add unit tests for policy output validation

cmd/policy/main.go

@@ -166,7 +166,7 @@ func main() {
 		healthSvc goahealth.Service
 	)
 	{
-		policySvc = policy.New(storage, regocache, cache, signer, logger)
+		policySvc = policy.New(storage, regocache, cache, signer, cfg.Policy.LockOnValidationFailure, logger)
 		healthSvc = health.New(Version)
 	}
 

internal/config/config.go

@@ -64,6 +64,10 @@ type policyConfig struct {
 	// needed policies. If present, only policies inside this folder
 	// are going to be fetched and used for evaluation.
 	Folder string `envconfig:"POLICY_REPOSITORY_FOLDER"`
+
+	// LockOnValidationFailure indicates whether a policy must be locked for execution
+	// if the policy output fails the schema validation.
+	LockOnValidationFailure bool `envconfig:"POLICY_LOCK_ON_VALIDATION_FAILURE" default:"false"`
 }
 
 type metricsConfig struct {

internal/service/policy/bundle_test.go

@@ -47,7 +47,7 @@ var testMetadata = Metadata{
 }
 
 func TestPolicy_createPolicyBundle(t *testing.T) {
-	svc := New(nil, nil, nil, nil, zap.NewNop())
+	svc := New(nil, nil, nil, nil, false, zap.NewNop())
 	bundle, err := svc.createPolicyBundle(testPolicy)
 	assert.NoError(t, err)
 	assert.NotNil(t, bundle)

internal/service/policy/service.go

@@ -44,14 +44,16 @@ type Signer interface {
 }
 
 type Service struct {
-	storage     Storage
-	policyCache RegoCache
-	cache       Cache
-	signer      Signer
-	logger      *zap.Logger
+	storage        Storage
+	policyCache    RegoCache
+	cache          Cache
+	signer         Signer
+	validationLock bool
+
+	logger *zap.Logger
 }
 
-func New(storage Storage, policyCache RegoCache, cache Cache, signer Signer, logger *zap.Logger) *Service {
+func New(storage Storage, policyCache RegoCache, cache Cache, signer Signer, validationLock bool, logger *zap.Logger) *Service {
 	signerFactory := func(signer Signer) func() (jws.Signer, error) {
 		return func() (jws.Signer, error) {
 			return &signAdapter{signer: signer}, nil
@@ -70,11 +72,12 @@ func New(storage Storage, policyCache RegoCache, cache Cache, signer Signer, log
 	jws.RegisterSigner(JwaVaultSignature, jws.SignerFactoryFn(signerFactory(signer)))
 
 	return &Service{
-		storage:     storage,
-		policyCache: policyCache,
-		cache:       cache,
-		signer:      signer,
-		logger:      logger,
+		storage:        storage,
+		policyCache:    policyCache,
+		cache:          cache,
+		signer:         signer,
+		validationLock: validationLock,
+		logger:         logger,
 	}
 }
 
@@ -205,8 +208,15 @@ func (s *Service) Validate(ctx context.Context, req *policy.EvaluateRequest) (*p
 
 	// validate the policy output
 	if err := sch.Validate(res.Result); err != nil {
-		logger.Error("invalid policy output schema", zap.Error(err))
-		return nil, errors.New("invalid policy output schema", err)
+		// lock the policy for execution if configured
+		if s.validationLock {
+			if err := s.lock(ctx, pol); err != nil {
+				logger.Error("error locking policy after validation failure", zap.Error(err))
+			}
+		}
+
+		logger.Error("policy output schema validation failed", zap.Error(err))
+		return nil, errors.New(errors.Unknown, "policy output schema validation failed", err)
 	}
 
 	return res, nil
@@ -231,17 +241,25 @@ func (s *Service) Lock(ctx context.Context, req *policy.LockRequest) error {
 		return errors.New("error locking policy", err)
 	}
 
-	if pol.Locked {
+	if err := s.lock(ctx, pol); err != nil {
+		logger.Error("error locking policy", zap.Error(err))
+		return err
+	}
+
+	logger.Debug("policy is locked")
+
+	return nil
+}
+
+func (s *Service) lock(ctx context.Context, p *storage.Policy) error {
+	if p.Locked {
 		return errors.New(errors.Forbidden, "policy is already locked")
 	}
 
-	if err := s.storage.SetPolicyLock(ctx, req.Repository, req.Group, req.PolicyName, req.Version, true); err != nil {
-		logger.Error("error locking policy", zap.Error(err))
+	if err := s.storage.SetPolicyLock(ctx, p.Repository, p.Group, p.Name, p.Version, true); err != nil {
 		return errors.New("error locking policy", err)
 	}
 
-	logger.Debug("policy is locked")
-
 	return nil
 }
 

internal/service/policy/service_test.go

@@ -26,10 +26,25 @@ import (
 )
 
 func TestNew(t *testing.T) {
-	svc := policy.New(nil, nil, nil, nil, zap.NewNop())
+	svc := policy.New(nil, nil, nil, nil, false, zap.NewNop())
 	assert.Implements(t, (*goapolicy.Service)(nil), svc)
 }
 
+// testReq prepares test request to be used in tests
+func testReq() *goapolicy.EvaluateRequest {
+	input := map[string]interface{}{"msg": "yes"}
+	var body interface{} = input
+
+	return &goapolicy.EvaluateRequest{
+		Repository: "policies",
+		Group:      "testgroup",
+		PolicyName: "example",
+		Version:    "1.0",
+		Input:      &body,
+		TTL:        ptr.Int(30),
+	}
+}
+
 func TestService_Evaluate(t *testing.T) {
 	testPolicy := &storage.Policy{
 		Repository: "policies",
@@ -52,21 +67,6 @@ func TestService_Evaluate(t *testing.T) {
 	// prepare test policy accessing headers during evaluation
 	testPolicyAccessingHeaders := `package testgroup.example token := external.http.header("Authorization")`
 
-	// prepare test request to be used in tests
-	testReq := func() *goapolicy.EvaluateRequest {
-		input := map[string]interface{}{"msg": "yes"}
-		var body interface{} = input
-
-		return &goapolicy.EvaluateRequest{
-			Repository: "policies",
-			Group:      "testgroup",
-			PolicyName: "example",
-			Version:    "1.0",
-			Input:      &body,
-			TTL:        ptr.Int(30),
-		}
-	}
-
 	// prepare test request with empty body
 	testEmptyReq := func() *goapolicy.EvaluateRequest {
 		var body interface{}
@@ -373,7 +373,7 @@ func TestService_Evaluate(t *testing.T) {
 
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
-			svc := policy.New(test.storage, test.regocache, test.cache, nil, zap.NewNop())
+			svc := policy.New(test.storage, test.regocache, test.cache, nil, false, zap.NewNop())
 			ctx := context.Background()
 			if test.ctx != nil {
 				ctx = test.ctx
@@ -397,6 +397,227 @@ func TestService_Evaluate(t *testing.T) {
 	}
 }
 
+func TestService_Validate(t *testing.T) {
+	// prepare basic JSON schema
+	jsonSchema := `
+		{
+		  "type": "object",
+		  "properties": {
+			"foo": {
+			  "type": "string",
+			  "minLength": 5
+			}
+		  },
+		  "required": [
+			"foo"
+		  ]
+		}
+	`
+	// prepare schema with specified $schema property
+	jsonSchemaWithSchemaProperty := `
+		{
+		  "$schema": "http://json-schema.org/draft-04/schema#",
+		  "type": "object",
+		  "properties": {
+			"foo": {
+			  "type": "string",
+			  "minLength": 5
+			}
+		  },
+		  "required": [
+			"foo"
+		  ]
+		}
+	`
+
+	tests := []struct {
+		name      string
+		req       *goapolicy.EvaluateRequest
+		storage   policy.Storage
+		regocache policy.RegoCache
+		cache     policy.Cache
+		// expected result
+		evalRes *goapolicy.EvaluateResult
+		errkind errors.Kind
+		errtext string
+	}{
+		{
+			name: "output validation schema is empty",
+			req:  testReq(),
+			regocache: &policyfakes.FakeRegoCache{
+				GetStub: func(key string) (*storage.Policy, bool) {
+					return nil, false
+				},
+			},
+			storage: &policyfakes.FakeStorage{
+				PolicyStub: func(ctx context.Context, s string, s2 string, s3 string, s4 string) (*storage.Policy, error) {
+					return &storage.Policy{
+						Repository:   "policies",
+						Name:         "example",
+						Group:        "testgroup",
+						Version:      "1.0",
+						Rego:         `package testgroup.example _ = {"hello":"world"}`,
+						Locked:       false,
+						OutputSchema: "",
+						LastUpdate:   time.Now(),
+					}, nil
+				},
+			},
+			cache: &policyfakes.FakeCache{
+				SetStub: func(ctx context.Context, s string, s2 string, s3 string, bytes []byte, i int) error {
+					return nil
+				},
+			},
+			errtext: "validation schema for policy output is not found",
+			errkind: errors.BadRequest,
+		},
+		{
+			name: "output validation schema is invalid JSON schema",
+			req:  testReq(),
+			regocache: &policyfakes.FakeRegoCache{
+				GetStub: func(key string) (*storage.Policy, bool) {
+					return nil, false
+				},
+			},
+			storage: &policyfakes.FakeStorage{
+				PolicyStub: func(ctx context.Context, s string, s2 string, s3 string, s4 string) (*storage.Policy, error) {
+					return &storage.Policy{
+						Repository:   "policies",
+						Name:         "example",
+						Group:        "testgroup",
+						Version:      "1.0",
+						Rego:         `package testgroup.example _ = {"hello":"world"}`,
+						Locked:       false,
+						OutputSchema: "invalid JSON schema",
+						LastUpdate:   time.Now(),
+					}, nil
+				},
+			},
+			cache: &policyfakes.FakeCache{
+				SetStub: func(ctx context.Context, s string, s2 string, s3 string, bytes []byte, i int) error {
+					return nil
+				},
+			},
+			errtext: "error compiling output validation schema",
+			errkind: errors.Unknown,
+		},
+		{
+			name: "policy output schema validation fails",
+			req:  testReq(),
+			regocache: &policyfakes.FakeRegoCache{
+				GetStub: func(key string) (*storage.Policy, bool) {
+					return nil, false
+				},
+			},
+			storage: &policyfakes.FakeStorage{
+				PolicyStub: func(ctx context.Context, s string, s2 string, s3 string, s4 string) (*storage.Policy, error) {
+					return &storage.Policy{
+						Repository:   "policies",
+						Name:         "example",
+						Group:        "testgroup",
+						Version:      "1.0",
+						Rego:         `package testgroup.example _ = {"foo":"bar"}`,
+						Locked:       false,
+						OutputSchema: jsonSchema,
+						LastUpdate:   time.Now(),
+					}, nil
+				},
+			},
+			cache: &policyfakes.FakeCache{
+				SetStub: func(ctx context.Context, s string, s2 string, s3 string, bytes []byte, i int) error {
+					return nil
+				},
+			},
+			errtext: "policy output schema validation failed",
+			errkind: errors.Unknown,
+		},
+		{
+			name: "policy output validation is successful",
+			req:  testReq(),
+			regocache: &policyfakes.FakeRegoCache{
+				GetStub: func(key string) (*storage.Policy, bool) {
+					return nil, false
+				},
+			},
+			storage: &policyfakes.FakeStorage{
+				PolicyStub: func(ctx context.Context, s string, s2 string, s3 string, s4 string) (*storage.Policy, error) {
+					return &storage.Policy{
+						Repository:   "policies",
+						Name:         "example",
+						Group:        "testgroup",
+						Version:      "1.0",
+						Rego:         `package testgroup.example _ = {"foo":"barbaz"}`,
+						Locked:       false,
+						OutputSchema: jsonSchema,
+						LastUpdate:   time.Now(),
+					}, nil
+				},
+			},
+			cache: &policyfakes.FakeCache{
+				SetStub: func(ctx context.Context, s string, s2 string, s3 string, bytes []byte, i int) error {
+					return nil
+				},
+			},
+			evalRes: &goapolicy.EvaluateResult{
+				Result: map[string]interface{}{"foo": "barbaz"},
+			},
+		},
+		{
+			name: "policy output validation using explicit schema draft version is successful ",
+			req:  testReq(),
+			regocache: &policyfakes.FakeRegoCache{
+				GetStub: func(key string) (*storage.Policy, bool) {
+					return nil, false
+				},
+			},
+			storage: &policyfakes.FakeStorage{
+				PolicyStub: func(ctx context.Context, s string, s2 string, s3 string, s4 string) (*storage.Policy, error) {
+					return &storage.Policy{
+						Repository:   "policies",
+						Name:         "example",
+						Group:        "testgroup",
+						Version:      "1.0",
+						Rego:         `package testgroup.example _ = {"foo":"barbaz"}`,
+						Locked:       false,
+						OutputSchema: jsonSchemaWithSchemaProperty,
+						LastUpdate:   time.Now(),
+					}, nil
+				},
+			},
+			cache: &policyfakes.FakeCache{
+				SetStub: func(ctx context.Context, s string, s2 string, s3 string, bytes []byte, i int) error {
+					return nil
+				},
+			},
+			evalRes: &goapolicy.EvaluateResult{
+				Result: map[string]interface{}{"foo": "barbaz"},
+			},
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			svc := policy.New(test.storage, test.regocache, test.cache, nil, false, zap.NewNop())
+
+			res, err := svc.Validate(context.Background(), test.req)
+			if err == nil {
+				assert.Empty(t, test.errtext)
+				assert.NotNil(t, res)
+
+				assert.Equal(t, test.evalRes.Result, res.Result)
+				assert.NotEmpty(t, res.ETag)
+			} else {
+				e, ok := err.(*errors.Error)
+				assert.True(t, ok)
+
+				assert.Contains(t, e.Error(), test.errtext)
+				assert.Equal(t, test.errkind, e.Kind)
+				assert.Equal(t, test.evalRes, res)
+			}
+		})
+	}
+}
+
 func TestService_Lock(t *testing.T) {
 	// prepare test request to be used in tests
 	testReq := func() *goapolicy.LockRequest {
@@ -479,7 +700,7 @@ func TestService_Lock(t *testing.T) {
 
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
-			svc := policy.New(test.storage, nil, nil, nil, zap.NewNop())
+			svc := policy.New(test.storage, nil, nil, nil, false, zap.NewNop())
 			err := svc.Lock(context.Background(), test.req)
 			if err == nil {
 				assert.Empty(t, test.errtext)
@@ -577,7 +798,7 @@ func TestService_Unlock(t *testing.T) {
 
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
-			svc := policy.New(test.storage, nil, nil, nil, zap.NewNop())
+			svc := policy.New(test.storage, nil, nil, nil, false, zap.NewNop())
 			err := svc.Unlock(context.Background(), test.req)
 			if err == nil {
 				assert.Empty(t, test.errtext)
@@ -809,7 +1030,7 @@ func TestService_ListPolicies(t *testing.T) {
 
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
-			svc := policy.New(test.storage, nil, nil, nil, zap.NewNop())
+			svc := policy.New(test.storage, nil, nil, nil, false, zap.NewNop())
 			result, err := svc.ListPolicies(context.Background(), test.request)
 
 			if test.errText != "" {
@@ -874,7 +1095,7 @@ func TestService_SubscribeForPolicyChange(t *testing.T) {
 
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
-			svc := policy.New(test.storage, nil, nil, nil, zap.NewNop())
+			svc := policy.New(test.storage, nil, nil, nil, false, zap.NewNop())
 			res, err := svc.SubscribeForPolicyChange(context.Background(), test.request)
 			if test.errText != "" {
 				assert.ErrorContains(t, err, test.errText)
@@ -894,7 +1115,7 @@ func TestService_ExportBundleError(t *testing.T) {
 				return nil, errors.New(errors.NotFound, "policy not found")
 			},
 		}
-		svc := policy.New(storage, nil, nil, nil, zap.NewNop())
+		svc := policy.New(storage, nil, nil, nil, false, zap.NewNop())
 		res, reader, err := svc.ExportBundle(context.Background(), &goapolicy.ExportBundleRequest{})
 		assert.Nil(t, res)
 		assert.Nil(t, reader)
@@ -911,7 +1132,7 @@ func TestService_ExportBundleError(t *testing.T) {
 				return nil, errors.New("unexpected error")
 			},
 		}
-		svc := policy.New(storage, nil, nil, nil, zap.NewNop())
+		svc := policy.New(storage, nil, nil, nil, false, zap.NewNop())
 		res, reader, err := svc.ExportBundle(context.Background(), &goapolicy.ExportBundleRequest{})
 		assert.Nil(t, res)
 		assert.Nil(t, reader)
@@ -945,7 +1166,7 @@ func TestService_ExportBundleError(t *testing.T) {
 			},
 		}
 
-		svc := policy.New(storage, nil, nil, signer, zap.NewNop())
+		svc := policy.New(storage, nil, nil, signer, false, zap.NewNop())
 		res, reader, err := svc.ExportBundle(context.Background(), &goapolicy.ExportBundleRequest{})
 		assert.Nil(t, res)
 		assert.Nil(t, reader)
@@ -977,7 +1198,7 @@ func TestService_ExportBundleSuccess(t *testing.T) {
 		},
 	}
 
-	svc := policy.New(storage, nil, nil, signer, zap.NewNop())
+	svc := policy.New(storage, nil, nil, signer, false, zap.NewNop())
 	res, reader, err := svc.ExportBundle(context.Background(), &goapolicy.ExportBundleRequest{})
 	require.NoError(t, err)
 	require.NotNil(t, res)