diff --git a/cmd/watchdog/root.go b/cmd/watchdog/root.go index 1a4521b..96fdbd4 100644 --- a/cmd/watchdog/root.go +++ b/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 { diff --git a/internal/aggregate/custom_events.go b/internal/aggregate/custom_events.go index cc0e92a..4593c72 100644 --- a/internal/aggregate/custom_events.go +++ b/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) +} diff --git a/internal/normalize/referrer_registry.go b/internal/normalize/referrer_registry.go index daaf69c..4b7913c 100644 --- a/internal/normalize/referrer_registry.go +++ b/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) +}