Harden diagnostic handling and warning presentation

internal/framework/diagnostics/aggregator.go

@@ -17,9 +17,10 @@ const (
 type Aggregator struct {
 	groups                   []contracts.DiagnosticGroup
 	indices                  map[groupKey]int
-	omitted                  map[groupKey]struct{}
 	warningGroups            int
 	nonWarningGroups         int
+	warningOccurrences       int
+	nonWarningOccurrences    int
 	unrepresentedOccurrences int
 }
 
@@ -32,31 +33,62 @@ func (aggregator *Aggregator) Add(group contracts.DiagnosticGroup) error {
 	}
 	if aggregator.indices == nil {
 		aggregator.indices = make(map[groupKey]int)
-		aggregator.omitted = make(map[groupKey]struct{})
 	}
 	key := groupKeyFromGroup(group)
 	if index, exists := aggregator.indices[key]; exists {
-		return aggregator.merge(index, group)
-	}
-	if _, exists := aggregator.omitted[key]; exists {
-		return aggregator.addUnrepresented(group.OccurrenceCount)
+		if err := aggregator.checkOccurrenceTotal(group.Disposition, group.OccurrenceCount); err != nil {
+			return err
+		}
+		if err := aggregator.merge(index, group); err != nil {
+			return err
+		}
+		aggregator.addOccurrenceTotal(group.Disposition, group.OccurrenceCount)
+		return nil
 	}
 	if group.Disposition == contracts.DiagnosticDispositionWarning {
 		if aggregator.warningGroups >= MaxWarningGroups {
 			return errors.New("diagnostic warning groups exceed maximum count")
 		}
-		aggregator.warningGroups++
 	} else if aggregator.nonWarningGroups >= MaxNonWarningGroups {
-		aggregator.omitted[key] = struct{}{}
+		if err := aggregator.checkOccurrenceTotal(group.Disposition, group.OccurrenceCount); err != nil {
+			return err
+		}
+		aggregator.addOccurrenceTotal(group.Disposition, group.OccurrenceCount)
 		return aggregator.addUnrepresented(group.OccurrenceCount)
+	}
+	if err := aggregator.checkOccurrenceTotal(group.Disposition, group.OccurrenceCount); err != nil {
+		return err
+	}
+	if group.Disposition == contracts.DiagnosticDispositionWarning {
+		aggregator.warningGroups++
 	} else {
 		aggregator.nonWarningGroups++
 	}
+	aggregator.addOccurrenceTotal(group.Disposition, group.OccurrenceCount)
 	aggregator.indices[key] = len(aggregator.groups)
 	aggregator.groups = append(aggregator.groups, contracts.CloneDiagnosticCollection(contracts.DiagnosticCollection{Groups: []contracts.DiagnosticGroup{group}}).Groups[0])
 	return nil
 }
 
+func (aggregator *Aggregator) checkOccurrenceTotal(disposition contracts.DiagnosticDisposition, count int) error {
+	current := aggregator.nonWarningOccurrences
+	if disposition == contracts.DiagnosticDispositionWarning {
+		current = aggregator.warningOccurrences
+	}
+	if count > maximumInt()-current {
+		return errors.New("diagnostic occurrence count overflow")
+	}
+	return nil
+}
+
+func (aggregator *Aggregator) addOccurrenceTotal(disposition contracts.DiagnosticDisposition, count int) {
+	if disposition == contracts.DiagnosticDispositionWarning {
+		aggregator.warningOccurrences += count
+		return
+	}
+	aggregator.nonWarningOccurrences += count
+}
+
 // Collection returns an independently owned grouped result in first-occurrence
 // order.
 func (aggregator *Aggregator) Collection() contracts.DiagnosticCollection {

internal/framework/diagnostics/aggregator_test.go

@@ -46,7 +46,7 @@ func TestAggregatorEnforcesWarningBoundAndTruncatesOnlyNonWarnings(t *testing.T)
 	}
 
 	nonWarnings := Aggregator{}
-	for index := 0; index < MaxNonWarningGroups+1; index++ {
+	for index := 0; index < MaxNonWarningGroups+3; index++ {
 		group := groupForAggregation("scope", "message", contracts.DiagnosticDispositionAdvisory, contracts.DiagnosticCategoryDataQuality, "module")
 		group.ReasonCode = "advisory-" + string(rune('a'+index))
 		if err := nonWarnings.Add(group); err != nil {
@@ -54,8 +54,26 @@ func TestAggregatorEnforcesWarningBoundAndTruncatesOnlyNonWarnings(t *testing.T)
 		}
 	}
 	collection := nonWarnings.Collection()
-	if len(collection.Groups) != MaxNonWarningGroups || !collection.Truncated || collection.UnrepresentedOccurrenceCount != 1 {
-		t.Fatalf("collection = %#v, want bounded non-warning groups and one unrepresented occurrence", collection)
+	if len(collection.Groups) != MaxNonWarningGroups || !collection.Truncated || collection.UnrepresentedOccurrenceCount != 3 {
+		t.Fatalf("collection = %#v, want bounded non-warning groups and three unrepresented occurrences", collection)
+	}
+}
+
+func TestAggregatorRejectsOccurrenceTotalOverflow(t *testing.T) {
+	aggregator := Aggregator{}
+	first := groupForAggregation("first", "first", contracts.DiagnosticDispositionAdvisory, contracts.DiagnosticCategoryDataQuality, "first")
+	first.OccurrenceCount = maximumInt()
+	first.OmittedSampleCount = maximumInt() - 1
+	if err := aggregator.Add(first); err != nil {
+		t.Fatal(err)
+	}
+	second := groupForAggregation("second", "second", contracts.DiagnosticDispositionAdvisory, contracts.DiagnosticCategoryDataQuality, "second")
+	if err := aggregator.Add(second); err == nil {
+		t.Fatal("Add() occurrence total overflow error = nil")
+	}
+	collection := aggregator.Collection()
+	if len(collection.Groups) != 1 || collection.Groups[0].OccurrenceCount != maximumInt() {
+		t.Fatalf("collection changed after rejected overflow = %#v", collection)
 	}
 }