diff --git a/internal/service/relational/risks/scores.go b/internal/service/relational/risks/scores.go index dd2b750f..394d4d7d 100644 --- a/internal/service/relational/risks/scores.go +++ b/internal/service/relational/risks/scores.go @@ -149,7 +149,7 @@ func (s *RiskService) RecordRiskScoreSnapshot(tx *gorm.DB, riskID uuid.UUID, sou status = RiskScoreStatusDeleted openBaselineScore = 0 openResidualScore = 0 - } else if isTerminalRiskStatus(risk.Status) { + } else if IsTerminalRiskStatus(risk.Status) { openBaselineScore = 0 openResidualScore = 0 } @@ -180,7 +180,8 @@ func isScoreValueEvent(eventType RiskEventType) bool { } } -func isTerminalRiskStatus(status string) bool { +// IsTerminalRiskStatus reports whether status is closed or remediated. +func IsTerminalRiskStatus(status string) bool { return status == string(RiskStatusClosed) || status == string(RiskStatusRemediated) } diff --git a/internal/service/worker/risk_evidence_integration_helpers_test.go b/internal/service/worker/risk_evidence_integration_helpers_test.go new file mode 100644 index 00000000..34183884 --- /dev/null +++ b/internal/service/worker/risk_evidence_integration_helpers_test.go @@ -0,0 +1,31 @@ +//go:build integration + +package worker + +import ( + "testing" + "time" + + "github.com/compliance-framework/api/internal/service/relational" + oscalTypes_1_1_3 "github.com/defenseunicorns/go-oscal/src/types/oscal-1-1-3" + "github.com/google/uuid" + "github.com/stretchr/testify/require" + "gorm.io/datatypes" + "gorm.io/gorm" +) + +// createStreamEvidence creates an evidence of the given stream with the given state and end. +func createStreamEvidence(t *testing.T, db *gorm.DB, stream uuid.UUID, state string, end time.Time) *relational.Evidence { + t.Helper() + id := uuid.New() + evidence := &relational.Evidence{ + UUIDModel: relational.UUIDModel{ID: &id}, + UUID: stream, + Title: "Test Evidence", + Start: end.Add(-time.Hour), + End: end, + Status: datatypes.NewJSONType(oscalTypes_1_1_3.ObjectiveStatus{State: state}), + } + require.NoError(t, db.Create(evidence).Error) + return evidence +} diff --git a/internal/service/worker/risk_evidence_worker.go b/internal/service/worker/risk_evidence_worker.go index f1fedfa1..b8207c63 100644 --- a/internal/service/worker/risk_evidence_worker.go +++ b/internal/service/worker/risk_evidence_worker.go @@ -1089,7 +1089,7 @@ func (w *RiskEvidenceWorker) handleEvidenceResolution(ctx context.Context, evide continue } // Skip risks that are already closed or remediated — nothing to resolve. - if risk.Status == string(risks.RiskStatusClosed) || risk.Status == string(risks.RiskStatusRemediated) { + if risks.IsTerminalRiskStatus(risk.Status) { continue } @@ -1098,7 +1098,7 @@ func (w *RiskEvidenceWorker) handleEvidenceResolution(ctx context.Context, evide continue } - if err := w.resolveRiskEvidenceLink(ctx, risk, evidence.UUID); err != nil { + if err := w.resolveRiskEvidenceLink(ctx, *risk.ID, evidence.UUID); err != nil { w.logger.Errorw("Failed to resolve risk evidence link", "error", err, "risk_id", risk.ID, "evidence_id", evidence.UUID) errs = append(errs, err) @@ -1138,8 +1138,23 @@ func (w *RiskEvidenceWorker) shouldRemoveEvidenceLink(risk *risks.Risk, evidence // resolveRiskEvidenceLink removes a single evidence link from a risk and, if no links remain, // transitions the risk to "remediated" (or emits an audit event for "risk-accepted" risks). -func (w *RiskEvidenceWorker) resolveRiskEvidenceLink(ctx context.Context, risk *risks.Risk, evidenceStreamID uuid.UUID) error { +// It locks the risk first, so resolutions of the risk's other streams count its links only +// after this one commits, and it acts on the status read under the lock. +func (w *RiskEvidenceWorker) resolveRiskEvidenceLink(ctx context.Context, riskID uuid.UUID, evidenceStreamID uuid.UUID) error { return w.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { + var risk risks.Risk + err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Select("id", "status").First(&risk, "id = ?", riskID).Error + if errors.Is(err, gorm.ErrRecordNotFound) { + return nil + } + if err != nil { + return fmt.Errorf("failed to lock risk: %w", err) + } + // Closed or remediated since the caller read it: nothing to resolve. + if risks.IsTerminalRiskStatus(risk.Status) { + return nil + } + // Delete the evidence link. result := tx.Delete(&risks.RiskEvidenceLink{}, "risk_id = ? AND evidence_id = ?", *risk.ID, evidenceStreamID) if result.Error != nil { @@ -1187,7 +1202,7 @@ func (w *RiskEvidenceWorker) resolveRiskEvidenceLink(ctx context.Context, risk * // Transition to remediated. oldStatus := risk.Status occurredAt := time.Now().UTC() - if err := tx.Model(risk).Update("status", string(risks.RiskStatusRemediated)).Error; err != nil { + if err := tx.Model(&risk).Update("status", string(risks.RiskStatusRemediated)).Error; err != nil { return fmt.Errorf("failed to transition risk to remediated: %w", err) } if err := w.emitRiskEventAt(ctx, tx, *risk.ID, string(risks.RiskEventTypeStatusChange), map[string]interface{}{ diff --git a/internal/service/worker/risk_resolution_race_integration_test.go b/internal/service/worker/risk_resolution_race_integration_test.go new file mode 100644 index 00000000..a66f3f57 --- /dev/null +++ b/internal/service/worker/risk_resolution_race_integration_test.go @@ -0,0 +1,124 @@ +//go:build integration + +package worker + +import ( + "context" + "sync/atomic" + "testing" + "time" + + "github.com/compliance-framework/api/internal/service/relational" + "github.com/compliance-framework/api/internal/service/relational/risks" + "github.com/compliance-framework/api/internal/tests" + "github.com/google/uuid" + "github.com/stretchr/testify/suite" + "go.uber.org/zap" + "gorm.io/gorm" +) + +type RiskResolutionRaceIntegrationSuite struct { + tests.IntegrationTestSuite +} + +func TestRiskResolutionRaceIntegrationSuite(t *testing.T) { + suite.Run(t, new(RiskResolutionRaceIntegrationSuite)) +} + +func (s *RiskResolutionRaceIntegrationSuite) SetupTest() { + s.Require().NoError(s.Migrator.Refresh()) +} + +// Two streams of a risk pass at once: the second job runs between the first's count of the +// risk's remaining links and its commit, and the risk must still end up remediated. +func (s *RiskResolutionRaceIntegrationSuite) TestConcurrentPassingStreamsRemediateTheRisk() { + ctx := context.Background() + worker := NewRiskEvidenceWorker(s.DB, zap.NewNop().Sugar()) + first := createStreamEvidence(s.T(), s.DB, uuid.New(), relational.EvidenceStatusSatisfied, time.Now()) + second := createStreamEvidence(s.T(), s.DB, uuid.New(), relational.EvidenceStatusSatisfied, time.Now()) + risk := createRiskWithEvidenceLink(s.T(), s.DB, string(risks.RiskStatusOpen), first.UUID, nil) + s.Require().NoError(s.DB.Create(&risks.RiskEvidenceLink{RiskID: *risk.ID, EvidenceID: second.UUID}).Error) + + const cb = "test:second_stream_resolves" + var fired atomic.Bool + var secondErr error + secondDone := make(chan struct{}) + s.Require().NoError(s.DB.Callback().Query().After("gorm:query").Register(cb, func(db *gorm.DB) { + _, isCount := db.Statement.Dest.(*int64) + if !isCount || db.Statement.Table != (risks.RiskEvidenceLink{}).TableName() || !fired.CompareAndSwap(false, true) { + return + } + go func() { + defer close(secondDone) + secondErr = worker.Work(ctx, makeWorkerJob(RiskProcessEvidenceArgs{EvidenceID: *second.ID})) + }() + // Let the second job finish, or block on a lock the first job holds. + s.Eventually(func() bool { + select { + case <-secondDone: + return true + default: + } + var waiting int64 + err := s.DB.Raw(`SELECT count(*) FROM pg_stat_activity + WHERE datname = current_database() AND wait_event_type = 'Lock'`).Scan(&waiting).Error + return err == nil && waiting > 0 + }, 10*time.Second, 10*time.Millisecond) + })) + defer func() { s.Require().NoError(s.DB.Callback().Query().Remove(cb)) }() + + s.Require().NoError(worker.Work(ctx, makeWorkerJob(RiskProcessEvidenceArgs{EvidenceID: *first.ID}))) + s.Require().True(fired.Load()) + <-secondDone + s.Require().NoError(secondErr) + + s.assertRisk(*risk.ID, risks.RiskStatusRemediated, 0) +} + +// A user closes the risk after resolution read it: the link stays and the risk stays closed. +func (s *RiskResolutionRaceIntegrationSuite) TestResolutionLeavesRiskClosedMeanwhile() { + riskID := s.resolveWhileStatusChanges(risks.RiskStatusClosed) + s.assertRisk(riskID, risks.RiskStatusClosed, 1) +} + +// A user accepts the risk after resolution read it: the link goes and the risk stays +// accepted, with an evidence_recovered event in place of remediation. +func (s *RiskResolutionRaceIntegrationSuite) TestResolutionKeepsRiskAcceptedMeanwhile() { + riskID := s.resolveWhileStatusChanges(risks.RiskStatusRiskAccepted) + s.assertRisk(riskID, risks.RiskStatusRiskAccepted, 0) + var recovered int64 + s.Require().NoError(s.DB.Model(&risks.RiskEvent{}).Where("risk_id = ? AND event_type = ?", + riskID, string(risks.RiskEventTypeEvidenceRecovered)).Count(&recovered).Error) + s.Equal(int64(1), recovered) +} + +// resolveWhileStatusChanges resolves a passing evidence of an open risk, setting the risk's +// status right after resolution's unlocked read of the linked risks. +func (s *RiskResolutionRaceIntegrationSuite) resolveWhileStatusChanges(status risks.RiskStatus) uuid.UUID { + worker := NewRiskEvidenceWorker(s.DB, zap.NewNop().Sugar()) + evidence := createStreamEvidence(s.T(), s.DB, uuid.New(), relational.EvidenceStatusSatisfied, time.Now()) + risk := createRiskWithEvidenceLink(s.T(), s.DB, string(risks.RiskStatusOpen), evidence.UUID, nil) + + const cb = "test:status_changes_after_read" + var changed atomic.Bool + s.Require().NoError(s.DB.Callback().Query().After("gorm:query").Register(cb, func(db *gorm.DB) { + if _, ok := db.Statement.Dest.(*[]risks.Risk); ok && changed.CompareAndSwap(false, true) { + s.Require().NoError(s.DB.Model(risk).Update("status", string(status)).Error) + } + })) + defer func() { s.Require().NoError(s.DB.Callback().Query().Remove(cb)) }() + + s.Require().NoError(worker.handleEvidenceResolution(context.Background(), evidence)) + s.Require().True(changed.Load()) + return *risk.ID +} + +// assertRisk checks the stored risk's status and how many evidence links it has left. +func (s *RiskResolutionRaceIntegrationSuite) assertRisk(id uuid.UUID, status risks.RiskStatus, links int64) { + var stored risks.Risk + s.Require().NoError(s.DB.First(&stored, "id = ?", id).Error) + s.Equal(string(status), stored.Status) + var count int64 + s.Require().NoError(s.DB.Model(&risks.RiskEvidenceLink{}).Where("risk_id = ?", id).Count(&count).Error) + s.Equal(links, count) +}