1package router
  2
  3import (
  4	"bufio"
  5	"context"
  6	"net"
  7	"net/http"
  8	"strconv"
  9	"strings"
 10	"time"
 11
 12	lru "github.com/hashicorp/golang-lru/v2"
 13	"github.com/prometheus/client_golang/prometheus"
 14	"github.com/prometheus/client_golang/prometheus/promauto"
 15)
 16
 17// UserMetricsOpts tells UserMetrics who a request belongs to when no
 18// handler called SetUser, such as a response served from a cache.
 19type UserMetricsOpts struct {
 20	// Owner names the user and resource a request is for, read from the
 21	// host without checking that either exists.
 22	Owner func(r *http.Request) (user, resource string)
 23	// TrustSuccess counts an Owner as real once it returns a success, for
 24	// a proxy that can't look users up. Only safe when every route on a
 25	// site's host fails for a site that doesn't exist.
 26	TrustSuccess bool
 27}
 28
 29// UserMetrics counts HTTP requests, latency and response bytes for each
 30// site owner. Requests and bytes also carry the resource they were for,
 31// such as a pgs project, which the latency histogram leaves out to keep its
 32// series down. Every app shares the metric names and sets the app label, so
 33// queries can aggregate across services.
 34type UserMetrics struct {
 35	requests *prometheus.CounterVec
 36	duration *prometheus.HistogramVec
 37	bytes    *prometheus.CounterVec
 38	opts     UserMetricsOpts
 39	// known holds the users and resources a handler or a successful
 40	// response confirmed. A name from the Host header alone never becomes a
 41	// label, so made-up subdomains can't add series.
 42	known *lru.Cache[string, struct{}]
 43}
 44
 45func NewUserMetrics(app string, reg prometheus.Registerer, opts UserMetricsOpts) *UserMetrics {
 46	auto := promauto.With(reg)
 47	known, _ := lru.New[string, struct{}](16384)
 48	return &UserMetrics{
 49		requests: auto.NewCounterVec(prometheus.CounterOpts{
 50			Namespace:   "pico",
 51			Subsystem:   "http",
 52			ConstLabels: prometheus.Labels{"app": app},
 53			Name:        "requests_total",
 54			Help:        "Number of HTTP requests handled.",
 55		}, []string{"user", "resource", "status", "method"}),
 56		duration: auto.NewHistogramVec(prometheus.HistogramOpts{
 57			Namespace:   "pico",
 58			Subsystem:   "http",
 59			ConstLabels: prometheus.Labels{"app": app},
 60			Name:        "request_duration_seconds",
 61			Help:        "Duration of HTTP requests in seconds.",
 62			Buckets:     prometheus.DefBuckets,
 63		}, []string{"user", "status", "method"}),
 64		bytes: auto.NewCounterVec(prometheus.CounterOpts{
 65			Namespace:   "pico",
 66			Subsystem:   "http",
 67			ConstLabels: prometheus.Labels{"app": app},
 68			Name:        "response_bytes_total",
 69			Help:        "Number of response body bytes sent.",
 70		}, []string{"user", "resource"}),
 71		opts:  opts,
 72		known: known,
 73	}
 74}
 75
 76type owner struct {
 77	user     string
 78	resource string
 79}
 80
 81type ctxOwnerKey struct{}
 82
 83// SetUser records name as the user a request belongs to. Handlers call it
 84// once they've found the user in the database.
 85func SetUser(r *http.Request, name string) {
 86	if o, ok := r.Context().Value(ctxOwnerKey{}).(*owner); ok {
 87		o.user = name
 88	}
 89}
 90
 91// SetResource records name as what a request is for within the user's
 92// account, such as a pgs project. Handlers call it alongside SetUser once
 93// they've found it.
 94func SetResource(r *http.Request, name string) {
 95	if o, ok := r.Context().Value(ctxOwnerKey{}).(*owner); ok {
 96		o.resource = name
 97	}
 98}
 99
100func (m *UserMetrics) Handler(next http.Handler) http.Handler {
101	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
102		start := time.Now()
103		var o owner
104		mw := &metricsWriter{ResponseWriter: w}
105		next.ServeHTTP(mw, r.WithContext(context.WithValue(r.Context(), ctxOwnerKey{}, &o)))
106
107		status := mw.Status()
108		user, resource := m.resolve(r, o, status)
109		m.Observe(user, resource, status, r.Method, time.Since(start), mw.bytes)
110	})
111}
112
113// Observe records one request for servers that can't use Handler, such as
114// tuns's gin logger. user and resource must already be confirmed to
115// exist, and resource may be empty. A negative size, which gin reports when
116// nothing was written, counts as zero bytes.
117func (m *UserMetrics) Observe(user, resource string, status int, method string, elapsed time.Duration, bytes int64) {
118	if user == "" {
119		resource = ""
120	}
121	code, verb := strconv.Itoa(status), metricMethod(method)
122	m.requests.WithLabelValues(user, resource, code, verb).Inc()
123	m.duration.WithLabelValues(user, code, verb).Observe(elapsed.Seconds())
124	m.bytes.WithLabelValues(user, resource).Add(float64(max(bytes, 0)))
125}
126
127func (m *UserMetrics) resolve(r *http.Request, o owner, status int) (string, string) {
128	var hostUser, hostResource string
129	if m.opts.Owner != nil {
130		hostUser, hostResource = m.opts.Owner(r)
131	}
132	trusted := m.opts.TrustSuccess && (status >= 200 && status < 300 || status == http.StatusNotModified)
133
134	user := m.confirm(o.user, hostUser, trusted)
135	if user == "" {
136		return "", ""
137	}
138	// Resources are keyed under their owner, and user names can't hold a
139	// slash.
140	prefix := user + "/"
141	set, fromHost := "", ""
142	if o.resource != "" {
143		set = prefix + o.resource
144	}
145	if hostUser == user && hostResource != "" {
146		fromHost = prefix + hostResource
147	}
148	resource := strings.TrimPrefix(m.confirm(set, fromHost, trusted), prefix)
149	return user, resource
150}
151
152// confirm returns the name a handler set, or else the name from the host
153// once it's known or trusted.
154func (m *UserMetrics) confirm(set, fromHost string, trusted bool) string {
155	if set != "" {
156		m.known.Add(set, struct{}{})
157		return set
158	}
159	if fromHost == "" {
160		return ""
161	}
162	if m.known.Contains(fromHost) {
163		return fromHost
164	}
165	if trusted {
166		m.known.Add(fromHost, struct{}{})
167		return fromHost
168	}
169	return ""
170}
171
172// metricMethod keeps clients from adding series with made-up methods.
173func metricMethod(method string) string {
174	switch method {
175	case http.MethodGet, http.MethodHead, http.MethodPost, http.MethodPut, http.MethodPatch,
176		http.MethodDelete, http.MethodOptions, http.MethodConnect, http.MethodTrace:
177		return method
178	}
179	return "OTHER"
180}
181
182type metricsWriter struct {
183	http.ResponseWriter
184	status int
185	bytes  int64
186}
187
188func (w *metricsWriter) WriteHeader(code int) {
189	// Informational responses like 103 Early Hints precede the real status.
190	if w.status == 0 && (code >= 200 || code == http.StatusSwitchingProtocols) {
191		w.status = code
192	}
193	w.ResponseWriter.WriteHeader(code)
194}
195
196func (w *metricsWriter) Write(b []byte) (int, error) {
197	if w.status == 0 {
198		w.status = http.StatusOK
199	}
200	n, err := w.ResponseWriter.Write(b)
201	w.bytes += int64(n)
202	return n, err
203}
204
205func (w *metricsWriter) Status() int {
206	if w.status == 0 {
207		return http.StatusOK
208	}
209	return w.status
210}
211
212func (w *metricsWriter) Flush() {
213	_ = http.NewResponseController(w.ResponseWriter).Flush()
214}
215
216// Hijack lets websockets upgrade through the middleware.
217func (w *metricsWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
218	conn, rw, err := http.NewResponseController(w.ResponseWriter).Hijack()
219	if err == nil && w.status == 0 {
220		w.status = http.StatusSwitchingProtocols
221	}
222	return conn, rw, err
223}
224
225func (w *metricsWriter) Unwrap() http.ResponseWriter {
226	return w.ResponseWriter
227}