various: standardize registry APIs; truncate metrics responses at 10MB

Signed-off-by: NotAShelf <raf@notashelf.dev>
Change-Id: I428255e61f8d2211fec0c320527b8e066a6a6964

cmd/watchdog/root.go

@@ -100,6 +100,10 @@ func Run(cfg *config.Config) error {
 		metricsHandler = rateLimitMiddleware(metricsHandler, metricsRateLimiter)
 	}
 
+	// Add response size limit to metrics endpoint (10MB max)
+	const maxMetricsResponseSize = 10 * 1024 * 1024 // 10MB
+	metricsHandler = responseSizeLimitMiddleware(metricsHandler, maxMetricsResponseSize)
+
 	mux.Handle(cfg.Server.MetricsPath, metricsHandler)
 
 	// Ingestion endpoint
@@ -191,6 +195,41 @@ func rateLimitMiddleware(next http.Handler, limiter *ratelimit.TokenBucket) http
 	})
 }
 
+// Wraps http.ResponseWriter to enforce max response size
+type limitedResponseWriter struct {
+	http.ResponseWriter
+	maxSize       int
+	written       int
+	limitExceeded bool
+}
+
+func (w *limitedResponseWriter) Write(p []byte) (int, error) {
+	if w.limitExceeded {
+		return 0, fmt.Errorf("response size limit exceeded")
+	}
+
+	if w.written+len(p) > w.maxSize {
+		w.limitExceeded = true
+		w.Header().Set("X-Response-Truncated", "true")
+		http.Error(w.ResponseWriter, "Response size limit exceeded", http.StatusInternalServerError)
+		return 0, fmt.Errorf("response size limit exceeded: %d bytes", w.maxSize)
+	}
+	n, err := w.ResponseWriter.Write(p)
+	w.written += n
+	return n, err
+}
+
+// Wraps a handler with response size limiting
+func responseSizeLimitMiddleware(next http.Handler, maxSize int) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		limited := &limitedResponseWriter{
+			ResponseWriter: w,
+			maxSize:        maxSize,
+		}
+		next.ServeHTTP(limited, r)
+	})
+}
+
 // Sanitizes a path for logging to prevent log injection attacks. Uses `strconv.Quote`
 // to properly escape control characters and special bytes.
 func sanitizePathForLog(path string) string {

internal/aggregate/custom_events.go

@@ -54,3 +54,20 @@ func (r *CustomEventRegistry) OverflowCount() int {
 	defer r.mu.RUnlock()
 	return r.overflowCount
 }
+
+// Contains checks if an event name exists in the registry.
+func (r *CustomEventRegistry) Contains(eventName string) bool {
+	r.mu.RLock()
+	defer r.mu.RUnlock()
+
+	_, exists := r.events[eventName]
+	return exists
+}
+
+// Count returns the number of unique events in the registry.
+func (r *CustomEventRegistry) Count() int {
+	r.mu.RLock()
+	defer r.mu.RUnlock()
+
+	return len(r.events)
+}

internal/normalize/referrer_registry.go

@@ -56,3 +56,20 @@ func (r *ReferrerRegistry) OverflowCount() int {
 	defer r.mu.RUnlock()
 	return r.overflowCount
 }
+
+// Contains checks if a source exists in the registry.
+func (r *ReferrerRegistry) Contains(source string) bool {
+	r.mu.RLock()
+	defer r.mu.RUnlock()
+
+	_, exists := r.sources[source]
+	return exists
+}
+
+// Count returns the number of unique sources in the registry.
+func (r *ReferrerRegistry) Count() int {
+	r.mu.RLock()
+	defer r.mu.RUnlock()
+
+	return len(r.sources)
+}