From bf4494e8fd0c5ca3beadbfae98c78038970f66c2 Mon Sep 17 00:00:00 2001
From: Catherine Fang <yinghongfang@google.com>
Date: Sat, 12 Aug 2023 17:10:50 -0400
Subject: [PATCH] Add HandlerWithAllowlist which can control exposed metrics

---
 prometheus/promhttp/http.go | 31 +++++++++++++++++++++++++++++++
 1 file changed, 31 insertions(+)

diff --git a/prometheus/promhttp/http.go b/prometheus/promhttp/http.go
index 09b8d2f..fc0d5b3 100644
--- a/prometheus/promhttp/http.go
+++ b/prometheus/promhttp/http.go
@@ -37,6 +37,7 @@ import (
 	"fmt"
 	"io"
 	"net/http"
+	"regexp"
 	"strconv"
 	"strings"
 	"sync"
@@ -45,6 +46,8 @@ import (
 	"github.com/prometheus/common/expfmt"
 
 	"github.com/prometheus/client_golang/prometheus"
+
+	dto "github.com/prometheus/client_model/go"
 )
 
 const (
@@ -90,6 +93,13 @@ func HandlerFor(reg prometheus.Gatherer, opts HandlerOpts) http.Handler {
 	return HandlerForTransactional(prometheus.ToTransactionalGatherer(reg), opts)
 }
 
+// HandlerWithAllowlist remove metrics whose name can't match with metricNameMatcher
+func HandlerWithAllowlist(metricNameMatcher *regexp.Regexp) http.Handler {
+	return InstrumentMetricHandler(
+		prometheus.DefaultRegisterer, HandlerFor(&FilteredGatherer{gather: prometheus.DefaultGatherer, metricNameMatcher: metricNameMatcher}, HandlerOpts{}),
+	)
+}
+
 // HandlerForTransactional is like HandlerFor, but it uses transactional gather, which
 // can safely change in-place returned *dto.MetricFamily before call to `Gather` and after
 // call to `done` of that `Gather`.
@@ -406,3 +416,24 @@ func httpError(rsp http.ResponseWriter, err error) {
 		http.StatusInternalServerError,
 	)
 }
+
+type FilteredGatherer struct {
+	gather            prometheus.Gatherer
+	metricNameMatcher *regexp.Regexp
+}
+
+func (f *FilteredGatherer) Gather() ([]*dto.MetricFamily, error) {
+	results := []*dto.MetricFamily{}
+	metrics, err := f.gather.Gather()
+	if err != nil {
+		return nil, err
+	}
+	for _, m := range metrics {
+
+		if f.metricNameMatcher.MatchString(*m.Name) {
+			results = append(results, m)
+		}
+
+	}
+	return results, nil
+}