Commit bd8701e

Antonio Mika  ·  2026-10-25 16:35:27 -0400 EDT
parent de92f37 
feat(metrics): count http requests and ssh sessions per user
13 files changed,  +485, -10
M cmd/pgs/cdn/main.go
M pkg/apps/pastes/api.go
M pkg/apps/pgs/auto_form.go
M pkg/apps/pgs/http_auth.go
A pkg/apps/pgs/site_metrics_test.go
M pkg/apps/pgs/web.go
M pkg/apps/pipe/api.go
M pkg/apps/prose/api.go
M pkg/pssh/server.go
A pkg/pssh/session_command_test.go
A pkg/shared/router/metrics.go
A pkg/shared/router/metrics_test.go
M pkg/shared/ssh.go
M cmd/pgs/cdn/main.go
+4, -3
......@@ -11,7 +11,6 @@ import (
1111 "strings"
1212
1313 "github.com/picosh/pico/pkg/apps/pgs"
14- "github.com/picosh/pico/pkg/httpcache"
1514 "github.com/picosh/pico/pkg/shared"
1615 "github.com/prometheus/client_golang/prometheus/promhttp"
1716 )
......@@ -29,9 +28,11 @@ func main() {
2928 cfg := pgs.NewConfig(logger, nil, nil, drain)
3029 proxy := newProxyServe(cfg.Logger)
3130 httpCache, cache := pgs.NewPgsHttpCache(cfg, proxy)
31+ // The CDN can't look users up, but origin fails every site route for a
32+ // site that doesn't exist, so a success proves the owner is real.
3233 cacher := &cachedHttp{
3334 Logger: cfg.Logger,
34- Cache: httpCache,
35+ Cache: pgs.NewSiteMetrics(cfg, "pgs-cdn", true).Handler(httpCache),
3536 }
3637
3738 go pgs.CacheMgmt(ctx, cfg, cache, nil)
......@@ -118,7 +119,7 @@ func (p *proxyServe) ServeHTTP(w http.ResponseWriter, req *http.Request) {
118119
119120 type cachedHttp struct {
120121 Logger *slog.Logger
121- Cache *httpcache.HttpCache
122+ Cache http.Handler
122123 }
123124
124125 func (c *cachedHttp) ServeHTTP(writer http.ResponseWriter, req *http.Request) {
M pkg/apps/pastes/api.go
+5, -0
......@@ -11,6 +11,7 @@ import (
1111 "github.com/picosh/pico/pkg/db/postgres"
1212 "github.com/picosh/pico/pkg/shared"
1313 "github.com/picosh/pico/pkg/shared/router"
14+ "github.com/prometheus/client_golang/prometheus"
1415 "github.com/prometheus/client_golang/prometheus/promhttp"
1516 )
1617
......@@ -84,6 +85,7 @@ func (web *webRouter) blogHandler(w http.ResponseWriter, r *http.Request) {
8485 http.Error(w, "user not found", http.StatusNotFound)
8586 return
8687 }
88+ router.SetUser(r, user.Name)
8789 logger = shared.LoggerWithUser(blogger, user)
8890
8991 pager, err := dbpool.FindPostsByUser(&db.Pager{Num: 1000, Page: 0}, user.ID, cfg.Space)
......@@ -169,6 +171,7 @@ func (web *webRouter) postHandler(w http.ResponseWriter, r *http.Request) {
169171 http.Error(w, "paste not found", http.StatusNotFound)
170172 return
171173 }
174+ router.SetUser(r, user.Name)
172175 logger = shared.LoggerWithUser(logger, user)
173176
174177 blogName := GetBlogName(username)
......@@ -266,6 +269,7 @@ func (web *webRouter) postHandlerRaw(w http.ResponseWriter, r *http.Request) {
266269 http.Error(w, "user not found", http.StatusNotFound)
267270 return
268271 }
272+ router.SetUser(r, user.Name)
269273 logger = shared.LoggerWithUser(blogger, user)
270274
271275 post, err := dbpool.FindPostWithSlug(slug, user.ID, cfg.Space)
......@@ -364,6 +368,7 @@ func StartApiServer() {
364368
365369 web := &webRouter{cfg: cfg, db: db}
366370 handler := router.NewHostRouter(cfg.Domain, cfg.Space, web.mainMux(), web.sitesMux())
371+ handler = router.NewUserMetrics("pastes", prometheus.DefaultRegisterer, router.UserMetricsOpts{}).Handler(handler)
367372
368373 portStr := fmt.Sprintf(":%s", cfg.Port)
369374 logger.Info(
M pkg/apps/pgs/auto_form.go
+1, -0
......@@ -43,6 +43,7 @@ func handleAutoForm(w http.ResponseWriter, r *http.Request, cfg *Config) {
4343 http.Error(w, "not found", http.StatusNotFound)
4444 return
4545 }
46+ router.SetUser(r, user.Name)
4647
4748 ff, err := cfg.DB.FindFeature(user.ID, "plus")
4849 if err != nil || (ff != nil && !ff.IsValid()) {
M pkg/apps/pgs/http_auth.go
+1, -0
......@@ -118,6 +118,7 @@ func handleLogin(w http.ResponseWriter, r *http.Request, cfg *Config) {
118118 http.Error(w, "not found", http.StatusNotFound)
119119 return
120120 }
121+ router.SetUser(r, user.Name)
121122
122123 project, err := cfg.DB.FindProjectByName(user.ID, projectName)
123124 if err != nil {
A pkg/apps/pgs/site_metrics_test.go
+61, -0
......@@ -0,0 +1,61 @@
1+package pgs
2+
3+import (
4+ "log/slog"
5+ "net/http/httptest"
6+ "strings"
7+ "testing"
8+
9+ "github.com/picosh/pico/pkg/shared"
10+ "github.com/picosh/pico/pkg/shared/router"
11+ "github.com/picosh/pico/pkg/storage"
12+ "github.com/prometheus/client_golang/prometheus"
13+)
14+
15+func TestSiteMetricsCountCacheHits(t *testing.T) {
16+ logger := slog.New(slog.DiscardHandler)
17+ dbpool := NewPgsDb(logger)
18+ user := dbpool.Users[0]
19+ st, err := storage.NewStorageMemory(map[string]map[string]string{
20+ shared.GetAssetBucketName(user.ID): {"/test/index.html": "hello"},
21+ })
22+ if err != nil {
23+ t.Fatal(err)
24+ }
25+ cfg := NewConfig(logger, dbpool, newTestStorage(st), discardPubsub{})
26+ cfg.Domain = "pgs.test"
27+ httpCache, _ := NewPgsHttpCache(cfg, newWebRouter(cfg))
28+ reg := prometheus.NewRegistry()
29+ h := router.NewUserMetrics("pgs", reg, router.UserMetricsOpts{SiteUser: siteUser(cfg)}).Handler(httpCache)
30+
31+ var statuses []string
32+ for range 2 {
33+ rec := httptest.NewRecorder()
34+ h.ServeHTTP(rec, httptest.NewRequest("GET", dbpool.mkpath("/"), nil))
35+ statuses = append(statuses, rec.Header().Get("cache-status"))
36+ }
37+ if !strings.Contains(statuses[1], "hit") {
38+ t.Fatalf("second request was not a cache hit: %q", statuses)
39+ }
40+
41+ families, err := reg.Gather()
42+ if err != nil {
43+ t.Fatal(err)
44+ }
45+ var count float64
46+ for _, family := range families {
47+ if family.GetName() != "pico_http_requests_total" {
48+ continue
49+ }
50+ for _, m := range family.GetMetric() {
51+ for _, l := range m.GetLabel() {
52+ if l.GetName() == "user" && l.GetValue() == user.Name {
53+ count += m.GetCounter().GetValue()
54+ }
55+ }
56+ }
57+ }
58+ if count != 2 {
59+ t.Errorf("requests for %s = %v, want 2", user.Name, count)
60+ }
61+}
M pkg/apps/pgs/web.go
+22, -1
......@@ -133,6 +133,26 @@ func NewPgsHttpCache(cfg *Config, upstream http.Handler) (*httpcache.HttpCache,
133133 return httpCache, cache
134134 }
135135
136+// NewSiteMetrics counts requests for each site owner. The owner comes from
137+// the host, so responses served from a cache count too.
138+func NewSiteMetrics(cfg *Config, app string, trustSuccess bool) *router.UserMetrics {
139+ return router.NewUserMetrics(app, prometheus.DefaultRegisterer, router.UserMetricsOpts{
140+ SiteUser: siteUser(cfg),
141+ TrustSuccess: trustSuccess,
142+ })
143+}
144+
145+func siteUser(cfg *Config) func(r *http.Request) string {
146+ return func(r *http.Request) string {
147+ subdomain := router.GetSubdomainFromRequest(r, cfg.Domain, cfg.TxtPrefix)
148+ if subdomain == "" {
149+ return ""
150+ }
151+ props, _ := router.GetProjectFromSubdomain(subdomain)
152+ return props.Username
153+ }
154+}
155+
136156 func StartApiServer(cfg *Config) {
137157 ctx := context.Background()
138158
......@@ -146,7 +166,7 @@ func StartApiServer(cfg *Config) {
146166 "port", cfg.WebPort,
147167 "domain", cfg.Domain,
148168 )
149- err := http.ListenAndServe(portStr, httpCache)
169+ err := http.ListenAndServe(portStr, NewSiteMetrics(cfg, "pgs", false).Handler(httpCache))
150170 cfg.Logger.Error(
151171 "listen and serve",
152172 "err", err.Error(),
......@@ -500,6 +520,7 @@ func (web *WebRouter) ServeAsset(fname string, opts *storage.ImgProcessOpts, has
500520 http.Error(w, msg, status)
501521 return
502522 }
523+ router.SetUser(r, site.user.Name)
503524 user, project, bucket := site.user, site.project, site.bucket
504525
505526 logger = logger.With(
M pkg/apps/pipe/api.go
+3, -0
......@@ -19,6 +19,7 @@ import (
1919 "github.com/picosh/pico/pkg/shared"
2020 "github.com/picosh/pico/pkg/shared/router"
2121 "github.com/picosh/utils/pipe"
22+ "github.com/prometheus/client_golang/prometheus"
2223 "github.com/prometheus/client_golang/prometheus/promhttp"
2324 "golang.org/x/crypto/ssh"
2425 )
......@@ -408,6 +409,7 @@ func rssHandler(cfg *Config, dbpool DB) http.HandlerFunc {
408409 http.Error(w, "invalid token", http.StatusNotFound)
409410 return
410411 }
412+ router.SetUser(r, user.Name)
411413 rss, err := MonitorRss(dbpool, user, cfg.Domain)
412414 if err != nil {
413415 cfg.Logger.Error(
......@@ -511,6 +513,7 @@ func StartApiServer() {
511513
512514 web := &webRouter{cfg: cfg, db: db, conn: newPipeConn(info)}
513515 handler := router.NewHostRouter(cfg.Domain, cfg.Space, web.mainMux(), web.sitesMux())
516+ handler = router.NewUserMetrics("pipe", prometheus.DefaultRegisterer, router.UserMetricsOpts{}).Handler(handler)
514517
515518 portStr := fmt.Sprintf(":%s", cfg.Port)
516519 logger.Info(
M pkg/apps/prose/api.go
+9, -0
......@@ -19,6 +19,7 @@ import (
1919 "github.com/picosh/pico/pkg/shared"
2020 "github.com/picosh/pico/pkg/shared/router"
2121 "github.com/picosh/pico/pkg/storage"
22+ "github.com/prometheus/client_golang/prometheus"
2223 "github.com/prometheus/client_golang/prometheus/promhttp"
2324 )
2425
......@@ -130,6 +131,7 @@ func (web *webRouter) blogStyleHandler(w http.ResponseWriter, r *http.Request) {
130131 http.Error(w, "blog not found", http.StatusNotFound)
131132 return
132133 }
134+ router.SetUser(r, user.Name)
133135 logger = shared.LoggerWithUser(logger, user)
134136
135137 styles, err := dbpool.FindPostWithFilename("_styles.css", user.ID, cfg.Space)
......@@ -160,6 +162,7 @@ func (web *webRouter) blogHandler(w http.ResponseWriter, r *http.Request) {
160162 http.Error(w, "blog not found", http.StatusNotFound)
161163 return
162164 }
165+ router.SetUser(r, user.Name)
163166 logger = shared.LoggerWithUser(logger, user)
164167
165168 tag := r.URL.Query().Get("tag")
......@@ -291,6 +294,7 @@ func (web *webRouter) postRawHandler(w http.ResponseWriter, r *http.Request) {
291294 http.Error(w, "blog not found", http.StatusNotFound)
292295 return
293296 }
297+ router.SetUser(r, user.Name)
294298 logger = shared.LoggerWithUser(logger, user)
295299
296300 post, err := dbpool.FindPostWithSlug(slug, user.ID, cfg.Space)
......@@ -320,6 +324,7 @@ func (web *webRouter) robotsHandler(w http.ResponseWriter, r *http.Request) {
320324 http.Error(w, "blog not found", http.StatusNotFound)
321325 return
322326 }
327+ router.SetUser(r, user.Name)
323328 logger = shared.LoggerWithUser(logger, user)
324329 w.Header().Add("Content-Type", "text/plain")
325330
......@@ -350,6 +355,7 @@ func (web *webRouter) postHandler(w http.ResponseWriter, r *http.Request) {
350355 http.Error(w, "blog not found", http.StatusNotFound)
351356 return
352357 }
358+ router.SetUser(r, user.Name)
353359
354360 logger = shared.LoggerWithUser(logger, user)
355361 logger = logger.With("slug", slug)
......@@ -591,6 +597,7 @@ func (web *webRouter) rssBlogHandler(w http.ResponseWriter, r *http.Request) {
591597 http.Error(w, "rss feed not found", http.StatusNotFound)
592598 return
593599 }
600+ router.SetUser(r, user.Name)
594601 logger = shared.LoggerWithUser(logger, user)
595602 logger.Info("fetching blog rss")
596603
......@@ -852,6 +859,7 @@ func (web *webRouter) imgRequest(w http.ResponseWriter, r *http.Request) {
852859 http.Error(w, "could find user", http.StatusNotFound)
853860 return
854861 }
862+ router.SetUser(r, user.Name)
855863 logger = shared.LoggerWithUser(logger, user)
856864
857865 rawname := r.PathValue("file")
......@@ -983,6 +991,7 @@ func StartApiServer() {
983991
984992 web := &webRouter{cfg: cfg, db: dbpool, storage: st}
985993 handler := router.NewHostRouter(cfg.Domain, cfg.Space, web.mainMux(), web.sitesMux())
994+ handler = router.NewUserMetrics("prose", prometheus.DefaultRegisterer, router.UserMetricsOpts{}).Handler(handler)
986995
987996 portStr := fmt.Sprintf(":%s", cfg.Port)
988997 logger.Info(
M pkg/pssh/server.go
+25, -6
......@@ -16,6 +16,7 @@ import (
1616 "net/http"
1717 "os"
1818 "path"
19+ "strings"
1920 "sync"
2021 "time"
2122 "unicode/utf8"
......@@ -349,7 +350,7 @@ func (s *SSHServer) ListenAndServe() error {
349350 ConstLabels: prometheus.Labels{
350351 "app": s.Config.App,
351352 },
352- }, []string{"command"})
353+ }, []string{"command", "user"})
353354
354355 s.SessionsFinished = promauto.With(prometheus.DefaultRegisterer).NewCounterVec(prometheus.CounterOpts{
355356 Name: "pssh_sessions_finished_total",
......@@ -357,7 +358,7 @@ func (s *SSHServer) ListenAndServe() error {
357358 ConstLabels: prometheus.Labels{
358359 "app": s.Config.App,
359360 },
360- }, []string{"command"})
361+ }, []string{"command", "user"})
361362
362363 s.SessionsDuration = promauto.With(prometheus.DefaultRegisterer).NewCounterVec(prometheus.CounterOpts{
363364 Name: "pssh_sessions_duration_seconds",
......@@ -365,7 +366,7 @@ func (s *SSHServer) ListenAndServe() error {
365366 ConstLabels: prometheus.Labels{
366367 "app": s.Config.App,
367368 },
368- }, []string{"command"})
369+ }, []string{"command", "user"})
369370
370371 // The handler reads the default registerer, so build it before the
371372 // listener is published rather than in the goroutine.
......@@ -618,10 +619,11 @@ func NewSSHServer(ctx context.Context, logger *slog.Logger, config *SSHServerCon
618619 }
619620
620621 if sc.SSHServer.Config.PromListenAddr != "" {
621- sc.SSHServer.SessionsCreated.WithLabelValues(command).Inc()
622+ labels := []string{sessionCommand(command), sessionUser(sesh)}
623+ sc.SSHServer.SessionsCreated.WithLabelValues(labels...).Inc()
622624 defer func() {
623- sc.SSHServer.SessionsFinished.WithLabelValues(command).Inc()
624- sc.SSHServer.SessionsDuration.WithLabelValues(command).Add(time.Since(sc.Start).Seconds())
625+ sc.SSHServer.SessionsFinished.WithLabelValues(labels...).Inc()
626+ sc.SSHServer.SessionsDuration.WithLabelValues(labels...).Add(time.Since(sc.Start).Seconds())
625627 }()
626628 }
627629
......@@ -695,6 +697,23 @@ func NewSSHServer(ctx context.Context, logger *slog.Logger, config *SSHServerCon
695697 return server
696698 }
697699
700+// sessionCommand is the subcommand alone, since arguments would add a series
701+// for every distinct command line.
702+func sessionCommand(command string) string {
703+ if fields := strings.Fields(command); len(fields) > 0 {
704+ return fields[0]
705+ }
706+ return command
707+}
708+
709+// sessionUser is the name the auth handler confirmed, or "" without one.
710+func sessionUser(s *SSHServerConnSession) string {
711+ if perms := s.Permissions(); perms != nil {
712+ return perms.Extensions["user"]
713+ }
714+ return ""
715+}
716+
698717 type PubKeyAuthHandler func(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error)
699718
700719 func NewSSHServerWithConfig(
A pkg/pssh/session_command_test.go
+16, -0
......@@ -0,0 +1,16 @@
1+package pssh
2+
3+import "testing"
4+
5+func TestSessionCommand(t *testing.T) {
6+ for command, want := range map[string]string{
7+ "shell": "shell",
8+ "pub my-topic -b=false": "pub",
9+ " ls ": "ls",
10+ "": "",
11+ } {
12+ if got := sessionCommand(command); got != want {
13+ t.Errorf("sessionCommand(%q) = %q, want %q", command, got, want)
14+ }
15+ }
16+}
A pkg/shared/router/metrics.go
+176, -0
......@@ -0,0 +1,176 @@
1+package router
2+
3+import (
4+ "bufio"
5+ "context"
6+ "net"
7+ "net/http"
8+ "strconv"
9+ "time"
10+
11+ lru "github.com/hashicorp/golang-lru/v2"
12+ "github.com/prometheus/client_golang/prometheus"
13+ "github.com/prometheus/client_golang/prometheus/promauto"
14+)
15+
16+// UserMetricsOpts tells UserMetrics who a request belongs to when no
17+// handler called SetUser, such as a response served from a cache.
18+type UserMetricsOpts struct {
19+ // SiteUser names the owner of the site a request is for, without
20+ // checking that the user exists.
21+ SiteUser func(r *http.Request) string
22+ // TrustSuccess counts a SiteUser as real once its site returns a
23+ // success, for a proxy that can't look users up. Only safe when every
24+ // route on a site's host fails for a site that doesn't exist.
25+ TrustSuccess bool
26+}
27+
28+// UserMetrics counts HTTP requests, latency and response bytes for each
29+// site owner. Every app shares the metric names and sets the app label, so
30+// queries can aggregate across services.
31+type UserMetrics struct {
32+ requests *prometheus.CounterVec
33+ duration *prometheus.HistogramVec
34+ bytes *prometheus.CounterVec
35+ opts UserMetricsOpts
36+ // known holds the users a handler or a successful response confirmed.
37+ // A name from the Host header alone never becomes a label, so made-up
38+ // subdomains can't add series.
39+ known *lru.Cache[string, struct{}]
40+}
41+
42+func NewUserMetrics(app string, reg prometheus.Registerer, opts UserMetricsOpts) *UserMetrics {
43+ auto := promauto.With(reg)
44+ known, _ := lru.New[string, struct{}](16384)
45+ return &UserMetrics{
46+ requests: auto.NewCounterVec(prometheus.CounterOpts{
47+ Namespace: "pico",
48+ Subsystem: "http",
49+ ConstLabels: prometheus.Labels{"app": app},
50+ Name: "requests_total",
51+ Help: "Number of HTTP requests handled.",
52+ }, []string{"user", "status", "method"}),
53+ duration: auto.NewHistogramVec(prometheus.HistogramOpts{
54+ Namespace: "pico",
55+ Subsystem: "http",
56+ ConstLabels: prometheus.Labels{"app": app},
57+ Name: "request_duration_seconds",
58+ Help: "Duration of HTTP requests in seconds.",
59+ Buckets: prometheus.DefBuckets,
60+ }, []string{"user", "status", "method"}),
61+ bytes: auto.NewCounterVec(prometheus.CounterOpts{
62+ Namespace: "pico",
63+ Subsystem: "http",
64+ ConstLabels: prometheus.Labels{"app": app},
65+ Name: "response_bytes_total",
66+ Help: "Number of response body bytes sent.",
67+ }, []string{"user"}),
68+ opts: opts,
69+ known: known,
70+ }
71+}
72+
73+type ctxUserKey struct{}
74+
75+// SetUser records name as the user a request belongs to. Handlers call it
76+// once they've found the user in the database.
77+func SetUser(r *http.Request, name string) {
78+ if user, ok := r.Context().Value(ctxUserKey{}).(*string); ok {
79+ *user = name
80+ }
81+}
82+
83+func (m *UserMetrics) Handler(next http.Handler) http.Handler {
84+ return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
85+ start := time.Now()
86+ var user string
87+ mw := &metricsWriter{ResponseWriter: w}
88+ next.ServeHTTP(mw, r.WithContext(context.WithValue(r.Context(), ctxUserKey{}, &user)))
89+
90+ status := mw.Status()
91+ user = m.resolveUser(r, user, status)
92+ labels := []string{user, strconv.Itoa(status), metricMethod(r.Method)}
93+ m.requests.WithLabelValues(labels...).Inc()
94+ m.duration.WithLabelValues(labels...).Observe(time.Since(start).Seconds())
95+ m.bytes.WithLabelValues(user).Add(float64(mw.bytes))
96+ })
97+}
98+
99+func (m *UserMetrics) resolveUser(r *http.Request, user string, status int) string {
100+ if user != "" {
101+ m.known.Add(user, struct{}{})
102+ return user
103+ }
104+ if m.opts.SiteUser == nil {
105+ return ""
106+ }
107+ user = m.opts.SiteUser(r)
108+ if user == "" {
109+ return ""
110+ }
111+ if m.known.Contains(user) {
112+ return user
113+ }
114+ if m.opts.TrustSuccess && (status >= 200 && status < 300 || status == http.StatusNotModified) {
115+ m.known.Add(user, struct{}{})
116+ return user
117+ }
118+ return ""
119+}
120+
121+// metricMethod keeps clients from adding series with made-up methods.
122+func metricMethod(method string) string {
123+ switch method {
124+ case http.MethodGet, http.MethodHead, http.MethodPost, http.MethodPut, http.MethodPatch,
125+ http.MethodDelete, http.MethodOptions, http.MethodConnect, http.MethodTrace:
126+ return method
127+ }
128+ return "OTHER"
129+}
130+
131+type metricsWriter struct {
132+ http.ResponseWriter
133+ status int
134+ bytes int64
135+}
136+
137+func (w *metricsWriter) WriteHeader(code int) {
138+ // Informational responses like 103 Early Hints precede the real status.
139+ if w.status == 0 && (code >= 200 || code == http.StatusSwitchingProtocols) {
140+ w.status = code
141+ }
142+ w.ResponseWriter.WriteHeader(code)
143+}
144+
145+func (w *metricsWriter) Write(b []byte) (int, error) {
146+ if w.status == 0 {
147+ w.status = http.StatusOK
148+ }
149+ n, err := w.ResponseWriter.Write(b)
150+ w.bytes += int64(n)
151+ return n, err
152+}
153+
154+func (w *metricsWriter) Status() int {
155+ if w.status == 0 {
156+ return http.StatusOK
157+ }
158+ return w.status
159+}
160+
161+func (w *metricsWriter) Flush() {
162+ _ = http.NewResponseController(w.ResponseWriter).Flush()
163+}
164+
165+// Hijack lets websockets upgrade through the middleware.
166+func (w *metricsWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
167+ conn, rw, err := http.NewResponseController(w.ResponseWriter).Hijack()
168+ if err == nil && w.status == 0 {
169+ w.status = http.StatusSwitchingProtocols
170+ }
171+ return conn, rw, err
172+}
173+
174+func (w *metricsWriter) Unwrap() http.ResponseWriter {
175+ return w.ResponseWriter
176+}
A pkg/shared/router/metrics_test.go
+161, -0
......@@ -0,0 +1,161 @@
1+package router
2+
3+import (
4+ "net/http"
5+ "net/http/httptest"
6+ "strings"
7+ "testing"
8+
9+ "github.com/prometheus/client_golang/prometheus"
10+)
11+
12+// metricValue reads a counter, or a histogram's count, from reg.
13+func metricValue(t *testing.T, reg *prometheus.Registry, name string, labels map[string]string) float64 {
14+ t.Helper()
15+ families, err := reg.Gather()
16+ if err != nil {
17+ t.Fatalf("gather: %v", err)
18+ }
19+ for _, family := range families {
20+ if family.GetName() != name {
21+ continue
22+ }
23+ metrics:
24+ for _, m := range family.GetMetric() {
25+ for _, l := range m.GetLabel() {
26+ if want, ok := labels[l.GetName()]; ok && want != l.GetValue() {
27+ continue metrics
28+ }
29+ }
30+ if h := m.GetHistogram(); h != nil {
31+ return float64(h.GetSampleCount())
32+ }
33+ return m.GetCounter().GetValue()
34+ }
35+ }
36+ return 0
37+}
38+
39+func serve(h http.Handler, method, target string) *httptest.ResponseRecorder {
40+ rec := httptest.NewRecorder()
41+ h.ServeHTTP(rec, httptest.NewRequest(method, target, nil))
42+ return rec
43+}
44+
45+func TestUserMetricsSetUser(t *testing.T) {
46+ reg := prometheus.NewRegistry()
47+ m := NewUserMetrics("test", reg, UserMetricsOpts{})
48+ h := m.Handler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
49+ SetUser(r, "alice")
50+ w.WriteHeader(http.StatusCreated)
51+ _, _ = w.Write([]byte("hello"))
52+ }))
53+
54+ serve(h, http.MethodPost, "/")
55+
56+ labels := map[string]string{"app": "test", "user": "alice", "status": "201", "method": "POST"}
57+ if got := metricValue(t, reg, "pico_http_requests_total", labels); got != 1 {
58+ t.Errorf("requests = %v, want 1", got)
59+ }
60+ if got := metricValue(t, reg, "pico_http_request_duration_seconds", labels); got != 1 {
61+ t.Errorf("duration samples = %v, want 1", got)
62+ }
63+ if got := metricValue(t, reg, "pico_http_response_bytes_total", map[string]string{"user": "alice"}); got != 5 {
64+ t.Errorf("bytes = %v, want 5", got)
65+ }
66+}
67+
68+func TestUserMetricsSiteUser(t *testing.T) {
69+ reg := prometheus.NewRegistry()
70+ m := NewUserMetrics("test", reg, UserMetricsOpts{
71+ SiteUser: func(r *http.Request) string { return strings.Split(r.Host, ".")[0] },
72+ })
73+ // The handler confirms alice on her first request; later ones stand in
74+ // for cache hits that never reach it.
75+ confirmed := false
76+ h := m.Handler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
77+ if !confirmed && r.Host == "alice.example.com" {
78+ SetUser(r, "alice")
79+ confirmed = true
80+ }
81+ }))
82+
83+ serve(h, http.MethodGet, "http://ghost.example.com/")
84+ serve(h, http.MethodGet, "http://alice.example.com/")
85+ serve(h, http.MethodGet, "http://alice.example.com/")
86+
87+ if got := metricValue(t, reg, "pico_http_requests_total", map[string]string{"user": "alice"}); got != 2 {
88+ t.Errorf("alice requests = %v, want 2", got)
89+ }
90+ if got := metricValue(t, reg, "pico_http_requests_total", map[string]string{"user": ""}); got != 1 {
91+ t.Errorf("unlabeled requests = %v, want 1", got)
92+ }
93+ if got := metricValue(t, reg, "pico_http_requests_total", map[string]string{"user": "ghost"}); got != 0 {
94+ t.Errorf("ghost requests = %v, want 0", got)
95+ }
96+}
97+
98+func TestUserMetricsTrustSuccess(t *testing.T) {
99+ reg := prometheus.NewRegistry()
100+ m := NewUserMetrics("test", reg, UserMetricsOpts{
101+ SiteUser: func(r *http.Request) string { return strings.Split(r.Host, ".")[0] },
102+ TrustSuccess: true,
103+ })
104+ h := m.Handler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
105+ if r.Host == "ghost.example.com" {
106+ http.NotFound(w, r)
107+ }
108+ }))
109+
110+ serve(h, http.MethodGet, "http://bob.example.com/")
111+ serve(h, http.MethodGet, "http://ghost.example.com/")
112+
113+ if got := metricValue(t, reg, "pico_http_requests_total", map[string]string{"user": "bob", "status": "200"}); got != 1 {
114+ t.Errorf("bob requests = %v, want 1", got)
115+ }
116+ if got := metricValue(t, reg, "pico_http_requests_total", map[string]string{"user": "", "status": "404"}); got != 1 {
117+ t.Errorf("unlabeled 404s = %v, want 1", got)
118+ }
119+}
120+
121+func TestUserMetricsMethod(t *testing.T) {
122+ reg := prometheus.NewRegistry()
123+ h := NewUserMetrics("test", reg, UserMetricsOpts{}).Handler(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}))
124+
125+ serve(h, "MADEUP", "/")
126+
127+ if got := metricValue(t, reg, "pico_http_requests_total", map[string]string{"method": "OTHER"}); got != 1 {
128+ t.Errorf("OTHER requests = %v, want 1", got)
129+ }
130+}
131+
132+func TestUserMetricsHijack(t *testing.T) {
133+ reg := prometheus.NewRegistry()
134+ h := NewUserMetrics("test", reg, UserMetricsOpts{}).Handler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
135+ hj, ok := w.(http.Hijacker)
136+ if !ok {
137+ t.Error("writer does not implement http.Hijacker")
138+ return
139+ }
140+ conn, rw, err := hj.Hijack()
141+ if err != nil {
142+ t.Errorf("hijack: %v", err)
143+ return
144+ }
145+ _, _ = rw.WriteString("HTTP/1.1 101 Switching Protocols\r\n\r\n")
146+ _ = rw.Flush()
147+ _ = conn.Close()
148+ }))
149+ srv := httptest.NewServer(h)
150+ defer srv.Close()
151+
152+ resp, err := http.Get(srv.URL)
153+ if err != nil {
154+ t.Fatalf("get: %v", err)
155+ }
156+ _ = resp.Body.Close()
157+
158+ if got := metricValue(t, reg, "pico_http_requests_total", map[string]string{"status": "101"}); got != 1 {
159+ t.Errorf("101 requests = %v, want 1", got)
160+ }
161+}
M pkg/shared/ssh.go
+1, -0
......@@ -132,6 +132,7 @@ func (r *SshAuthHandler) PubkeyAuthHandler(conn ssh.ConnMetadata, key ssh.Public
132132 perms := &ssh.Permissions{
133133 Extensions: map[string]string{
134134 "user_id": user.ID,
135+ "user": user.Name,
135136 "pubkey": authed.Pubkey,
136137 "identity": authed.Identity,
137138 },