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 {