coder · Emyrk · Jan 10, 2024 · Jan 9, 2024 · Jan 9, 2024 · Jan 10, 2024
diff --git a/cli/server.go b/cli/server.go
@@ -1802,7 +1802,7 @@ func configureGithubOAuth2(instrument *promoauth.Factory, accessURL *url.URL, cl
 	}
 
 	return &coderd.GithubOAuth2Config{
-		OAuth2Config: instrument.New("github-login", &oauth2.Config{
+		OAuth2Config: instrument.NewGithub("github-login", &oauth2.Config{
 			ClientID:     clientID,
 			ClientSecret: clientSecret,
 			Endpoint:     endpoint,

diff --git a/coderd/coderdtest/oidctest/idp.go b/coderd/coderdtest/oidctest/idp.go
@@ -85,6 +85,8 @@ type FakeIDP struct {
 	// to test something like PKI auth vs a client_secret.
 	hookAuthenticateClient func(t testing.TB, req *http.Request) (url.Values, error)
 	serve                  bool
+	// optional middlewares
+	middlewares chi.Middlewares
 }
 
 func StatusError(code int, err error) error {
@@ -115,6 +117,12 @@ func WithAuthorizedRedirectURL(hook func(redirectURL string) error) func(*FakeID
 	}
 }
 
+func WithMiddlewares(mws ...func(http.Handler) http.Handler) func(*FakeIDP) {
+	return func(f *FakeIDP) {
+		f.middlewares = append(f.middlewares, mws...)
+	}
+}
+
 // WithRefresh is called when a refresh token is used. The email is
 // the email of the user that is being refreshed assuming the claims are correct.
 func WithRefresh(hook func(email string) error) func(*FakeIDP) {
@@ -570,6 +578,7 @@ func (f *FakeIDP) httpHandler(t testing.TB) http.Handler {
 	t.Helper()
 
 	mux := chi.NewMux()
+	mux.Use(f.middlewares...)
 	// This endpoint is required to initialize the OIDC provider.
 	// It is used to get the OIDC configuration.
 	mux.Get("/.well-known/openid-configuration", func(rw http.ResponseWriter, r *http.Request) {

diff --git a/coderd/externalauth/externalauth.go b/coderd/externalauth/externalauth.go
@@ -464,8 +464,13 @@ func ConvertConfig(instrument *promoauth.Factory, entries []codersdk.ExternalAut
 			oauthConfig = &exchangeWithClientSecret{oc}
 		}
 
+		instrumented := instrument.New(entry.ID, oauthConfig)
+		if strings.EqualFold(entry.Type, string(codersdk.EnhancedExternalAuthProviderGitHub)) {
+			instrumented = instrument.NewGithub(entry.ID, oauthConfig)
+		}
+
 		cfg := &Config{
-			InstrumentedOAuth2Config: instrument.New(entry.ID, oauthConfig),
+			InstrumentedOAuth2Config: instrumented,
 			ID:                       entry.ID,
 			Regex:                    regex,
 			Type:                     entry.Type,

diff --git a/coderd/promoauth/github.go b/coderd/promoauth/github.go
@@ -0,0 +1,100 @@
+package promoauth
+
+import (
+	"fmt"
+	"net/http"
+	"strconv"
+	"time"
+)
+
+type rateLimits struct {
+	Limit     int
+	Remaining int
+	Used      int
+	Reset     time.Time
+	Resource  string
+}
+
+// githubRateLimits checks the returned response headers and
+func githubRateLimits(resp *http.Response, err error) (rateLimits, bool) {
+	if err != nil || resp == nil {
+		return rateLimits{}, false
+	}
+
+	p := headerParser{header: resp.Header}
+	// See
+	// https://docs.github.com/en/rest/using-the-rest-api/rate-limits-for-the-rest-api?apiVersion=2022-11-28#checking-the-status-of-your-rate-limit
+	limits := rateLimits{
+		Limit:     p.int("x-ratelimit-limit"),
+		Remaining: p.int("x-ratelimit-remaining"),
+		Used:      p.int("x-ratelimit-used"),
+		Resource:  p.string("x-ratelimit-resource"),
+	}
+
+	if limits.Limit == 0 &&
+		limits.Remaining == 0 &&
+		limits.Used == 0 {
+		// For some requests, github has no rate limit. In which case,
+		// it returns all 0s. We can just omit these.
+		return limits, false
+	}
+
+	// Reset is when the rate limit "used" will be reset to 0.
+	// If it's unix 0, then we do not know when it will reset.
+	// Change it to a zero time as that is easier to handle in golang.
+	unix := p.int("x-ratelimit-reset")
+	resetAt := time.Unix(int64(unix), 0)
+	if unix == 0 {
+		resetAt = time.Time{}
+	}
+	limits.Reset = resetAt
+
+	// Unauthorized requests have their own rate limit, so we should
+	// track them separately.
+	if resp.StatusCode == http.StatusUnauthorized {
+		limits.Resource += "-unauthorized"
+	}
+
+	// A 401 or 429 means too many requests. This might mess up the
+	// "resource" string because we could hit the unauthorized limit,
+	// and we do not want that to override the authorized one.
+	// However, in testing, it seems a 401 is always a 401, even if
+	// the limit is hit.
+
+	if len(p.errors) > 0 {
+		// If we are missing any headers, then do not try and guess
+		// what the rate limits are.
+		return limits, false
+	}
+	return limits, true
+}
+
+type headerParser struct {
+	errors map[string]error
+	header http.Header
+}
+
+func (p *headerParser) string(key string) string {
+	if p.errors == nil {
+		p.errors = make(map[string]error)
+	}
+
+	v := p.header.Get(key)
+	if v == "" {
+		p.errors[key] = fmt.Errorf("missing header %q", key)
+	}
+	return v
+}
+
+func (p *headerParser) int(key string) int {
+	v := p.string(key)
+	if v == "" {
+		return -1
+	}
+
+	i, err := strconv.Atoi(v)
+	if err != nil {
+		p.errors[key] = err
+	}
+	return i
+}
diff --git a/coderd/promoauth/oauth2.go b/coderd/promoauth/oauth2.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"net/http"
+	"time"
 
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
@@ -46,11 +47,25 @@ var _ OAuth2Config = (*Config)(nil)
 // Primarily to avoid any prometheus errors registering duplicate metrics.
 type Factory struct {
 	metrics *metrics
+	// optional replace now func
+	Now func() time.Time
 }
 
 // metrics is the reusable metrics for all oauth2 providers.
 type metrics struct {
 	externalRequestCount *prometheus.CounterVec
+
+	// if the oauth supports it, rate limit metrics.
+	// rateLimit is the defined limit per interval
+	rateLimit          *prometheus.GaugeVec
+	rateLimitRemaining *prometheus.GaugeVec
+	rateLimitUsed      *prometheus.GaugeVec
+	// rateLimitReset is unix time of the next interval (when the rate limit resets).
+	rateLimitReset *prometheus.GaugeVec
+	// rateLimitResetIn is the time in seconds until the rate limit resets.
+	// This is included because it is sometimes more helpful to know the limit
+	// will reset in 600seconds, rather than at 1704000000 unix time.
+	rateLimitResetIn *prometheus.GaugeVec
 }
 
 func NewFactory(registry prometheus.Registerer) *Factory {
@@ -68,6 +83,53 @@ func NewFactory(registry prometheus.Registerer) *Factory {
 				"source",
 				"status_code",
 			}),
+			rateLimit: factory.NewGaugeVec(prometheus.GaugeOpts{
+				Namespace: "coderd",
+				Subsystem: "oauth2",
+				Name:      "external_requests_rate_limit_total",
+				Help:      "The total number of allowed requests per interval.",
+			}, []string{
+				"name",
+				// Resource allows different rate limits for the same oauth2 provider.
+				// Some IDPs have different buckets for different rate limits.
+				"resource",
+			}),
+			rateLimitRemaining: factory.NewGaugeVec(prometheus.GaugeOpts{
+				Namespace: "coderd",
+				Subsystem: "oauth2",
+				Name:      "external_requests_rate_limit_remaining",
+				Help:      "The remaining number of allowed requests in this interval.",
+			}, []string{
+				"name",
+				"resource",
+			}),
+			rateLimitUsed: factory.NewGaugeVec(prometheus.GaugeOpts{
+				Namespace: "coderd",
+				Subsystem: "oauth2",
+				Name:      "external_requests_rate_limit_used",
+				Help:      "The number of requests made in this interval.",
+			}, []string{
+				"name",
+				"resource",
+			}),
+			rateLimitReset: factory.NewGaugeVec(prometheus.GaugeOpts{
+				Namespace: "coderd",
+				Subsystem: "oauth2",
+				Name:      "external_requests_rate_limit_next_reset_unix",
+				Help:      "Unix timestamp for when the next interval starts",
+			}, []string{
+				"name",
+				"resource",
+			}),
+			rateLimitResetIn: factory.NewGaugeVec(prometheus.GaugeOpts{
+				Namespace: "coderd",
+				Subsystem: "oauth2",
+				Name:      "external_requests_rate_limit_reset_in_seconds",
+				Help:      "Seconds until the next interval",
+			}, []string{
+				"name",
+				"resource",
+			}),
 		},
 	}
 }
@@ -80,13 +142,53 @@ func (f *Factory) New(name string, under OAuth2Config) *Config {
 	}
 }
 
+// NewGithub returns a new instrumented oauth2 config for github. It tracks
+// rate limits as well as just the external request counts.
+//
+//nolint:bodyclose
+func (f *Factory) NewGithub(name string, under OAuth2Config) *Config {
+	cfg := f.New(name, under)
+	cfg.interceptors = append(cfg.interceptors, func(resp *http.Response, err error) {
+		limits, ok := githubRateLimits(resp, err)
+		if !ok {
+			return
+		}
+		labels := prometheus.Labels{
+			"name":     cfg.name,
+			"resource": limits.Resource,
+		}
+		// Default to -1 for "do not know"
+		resetIn := float64(-1)
+		if !limits.Reset.IsZero() {
+			now := time.Now()
+			if f.Now != nil {
+				now = f.Now()
+			}
+			resetIn = limits.Reset.Sub(now).Seconds()
+			if resetIn < 0 {
+				// If it just reset, just make it 0.
+				resetIn = 0
+			}
+		}
+
+		f.metrics.rateLimit.With(labels).Set(float64(limits.Limit))
+		f.metrics.rateLimitRemaining.With(labels).Set(float64(limits.Remaining))
+		f.metrics.rateLimitUsed.With(labels).Set(float64(limits.Used))
+		f.metrics.rateLimitReset.With(labels).Set(float64(limits.Reset.Unix()))
+		f.metrics.rateLimitResetIn.With(labels).Set(resetIn)
+	})
+	return cfg
+}
+
 type Config struct {
 	// Name is a human friendly name to identify the oauth2 provider. This should be
 	// deterministic from restart to restart, as it is going to be used as a label in
 	// prometheus metrics.
 	name       string
 	underlying OAuth2Config
 	metrics    *metrics
+	// interceptors are called after every request made by the oauth2 client.
+	interceptors []func(resp *http.Response, err error)
 }
 
 func (c *Config) Do(ctx context.Context, source Oauth2Source, req *http.Request) (*http.Response, error) {
@@ -169,5 +271,10 @@ func (i *instrumentedTripper) RoundTrip(r *http.Request) (*http.Response, error)
 		"source":      string(i.source),
 		"status_code": fmt.Sprintf("%d", statusCode),
 	}).Inc()
+
+	// Handle any extra interceptors.
+	for _, interceptor := range i.c.interceptors {
+		interceptor(resp, err)
+	}
 	return resp, err
 }