coder/coderd/promoauth/oauth2_test.go

package promoauth_test

import (
	"context"
	"fmt"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"testing"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"golang.org/x/oauth2"
	"golang.org/x/sync/singleflight"

	"github.com/coder/coder/v2/coderd/coderdtest/oidctest"
	"github.com/coder/coder/v2/coderd/coderdtest/promhelp"
	"github.com/coder/coder/v2/coderd/externalauth"
	"github.com/coder/coder/v2/coderd/promoauth"
	"github.com/coder/coder/v2/testutil"
)

func TestInstrument(t *testing.T) {
	t.Parallel()

	ctx := testutil.Context(t, testutil.WaitShort)
	idp := oidctest.NewFakeIDP(t, oidctest.WithServing())
	reg := prometheus.NewRegistry()
	t.Cleanup(func() {
		if t.Failed() {
			t.Log(promhelp.RegistryDump(reg))
		}
	})

	const id = "test"
	labels := prometheus.Labels{
		"name":        id,
		"status_code": "200",
	}
	const metricname = "coderd_oauth2_external_requests_total"
	count := func(source string) int {
		labels["source"] = source
		return promhelp.CounterValue(t, reg, "coderd_oauth2_external_requests_total", labels)
	}

	factory := promoauth.NewFactory(reg)

	cfg := externalauth.Config{
		InstrumentedOAuth2Config: factory.New(id, idp.OIDCConfig(t, []string{})),
		ID:                       "test",
		ValidateURL:              must[*url.URL](t)(idp.IssuerURL().Parse("/oauth2/userinfo")).String(),
		RefreshGroup:             new(singleflight.Group),
	}

	// 0 Requests before we start
	require.Nil(t, promhelp.MetricValue(t, reg, metricname, labels), "no metrics at start")

	noClientCtx := ctx
	// This should never be done, but promoauth should not break the default client
	// even if this happens. So intentionally do this to verify nothing breaks.
	ctx = context.WithValue(ctx, oauth2.HTTPClient, http.DefaultClient)
	// Exchange should trigger a request
	code := idp.CreateAuthCode(t, "foo")
	_, err := cfg.Exchange(ctx, code)
	require.NoError(t, err)
	require.Equal(t, count("Exchange"), 1)

	// Do an exchange without a default http client as well to verify original
	// transport is not broken.
	code = idp.CreateAuthCode(t, "bar")
	token, err := cfg.Exchange(noClientCtx, code)
	require.NoError(t, err)
	require.Equal(t, count("Exchange"), 2)

	// Force a refresh
	token.Expiry = time.Now().Add(time.Hour * -1)
	src := cfg.TokenSource(ctx, token)
	refreshed, err := src.Token()
	require.NoError(t, err)
	require.NotEqual(t, token.AccessToken, refreshed.AccessToken, "token refreshed")
	require.Equal(t, count("TokenSource"), 1)

	// Try a validate
	valid, _, err := cfg.ValidateToken(ctx, refreshed)
	require.NoError(t, err)
	require.True(t, valid)
	require.Equal(t, count("ValidateToken"), 1)

	// Verify the default client was not broken. This check is added because we
	// extend the http.DefaultTransport. If a `.Clone()` is not done, this can be
	// mis-used. It is cheap to run this quick check.
	snapshot := promhelp.RegistryDump(reg)
	req, err := http.NewRequestWithContext(ctx, http.MethodGet,
		must[*url.URL](t)(idp.IssuerURL().Parse("/.well-known/openid-configuration")).String(), nil)
	require.NoError(t, err)

	client := &http.Client{}
	resp, err := client.Do(req)
	require.NoError(t, err)
	_ = resp.Body.Close()

	require.NoError(t, promhelp.Compare(reg, snapshot), "http default client corrupted")
}

func TestGithubRateLimits(t *testing.T) {
	t.Parallel()

	now := time.Now()
	cases := []struct {
		Name            string
		NoHeaders       bool
		Omit            []string
		ExpectNoMetrics bool
		Limit           int
		Remaining       int
		Used            int
		Reset           time.Time

		at time.Time
	}{
		{
			Name:            "NoHeaders",
			NoHeaders:       true,
			ExpectNoMetrics: true,
		},
		{
			Name:            "ZeroHeaders",
			ExpectNoMetrics: true,
		},
		{
			Name:      "OverLimit",
			Limit:     100,
			Remaining: 0,
			Used:      500,
			Reset:     now.Add(time.Hour),
			at:        now,
		},
		{
			Name:      "UnderLimit",
			Limit:     100,
			Remaining: 0,
			Used:      500,
			Reset:     now.Add(time.Hour),
			at:        now,
		},
		{
			Name:            "Partial",
			Omit:            []string{"x-ratelimit-remaining"},
			ExpectNoMetrics: true,
			Limit:           100,
			Remaining:       0,
			Used:            500,
			Reset:           now.Add(time.Hour),
			at:              now,
		},
	}

	for _, c := range cases {
		t.Run(c.Name, func(t *testing.T) {
			t.Parallel()

			reg := prometheus.NewRegistry()
			idp := oidctest.NewFakeIDP(t, oidctest.WithMiddlewares(
				func(next http.Handler) http.Handler {
					return http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
						if !c.NoHeaders {
							rw.Header().Set("x-ratelimit-limit", fmt.Sprintf("%d", c.Limit))
							rw.Header().Set("x-ratelimit-remaining", fmt.Sprintf("%d", c.Remaining))
							rw.Header().Set("x-ratelimit-used", fmt.Sprintf("%d", c.Used))
							rw.Header().Set("x-ratelimit-resource", "core")
							rw.Header().Set("x-ratelimit-reset", fmt.Sprintf("%d", c.Reset.Unix()))
							for _, omit := range c.Omit {
								rw.Header().Del(omit)
							}
						}

						// Ignore the well known, required for setup
						if !strings.Contains(r.URL.String(), ".well-known") {
							// GitHub rate limits only are instrumented for unauthenticated calls. So emulate
							// that here. We cannot actually use invalid credentials because the fake IDP
							// will throw a test error, as it only expects things to work. And we want to use
							// a real idp to emulate the right http calls to be instrumented.
							rw.WriteHeader(http.StatusUnauthorized)
							return
						}
						next.ServeHTTP(rw, r)
					})
				}))

			factory := promoauth.NewFactory(reg)
			if !c.at.IsZero() {
				factory.Now = func() time.Time {
					return c.at
				}
			}

			cfg := factory.NewGithub("test", idp.OIDCConfig(t, []string{}))

			// Do a single oauth2 call
			ctx := testutil.Context(t, testutil.WaitShort)
			ctx = context.WithValue(ctx, oauth2.HTTPClient, idp.HTTPClient(nil))
			_, err := cfg.Exchange(ctx, "invalid")
			require.Error(t, err, "expected unauthorized exchange")

			// Verify
			labels := prometheus.Labels{
				"name":     "test",
				"resource": "core-unauthorized",
			}
			pass := true
			if !c.ExpectNoMetrics {
				pass = pass && assert.Equal(t, promhelp.GaugeValue(t, reg, "coderd_oauth2_external_requests_rate_limit", labels), c.Limit, "limit")
				pass = pass && assert.Equal(t, promhelp.GaugeValue(t, reg, "coderd_oauth2_external_requests_rate_limit_remaining", labels), c.Remaining, "remaining")
				pass = pass && assert.Equal(t, promhelp.GaugeValue(t, reg, "coderd_oauth2_external_requests_rate_limit_used", labels), c.Used, "used")
				if !c.at.IsZero() {
					until := c.Reset.Sub(c.at)
					// Float accuracy is not great, so we allow a delta of 2
					pass = pass && assert.InDelta(t, promhelp.GaugeValue(t, reg, "coderd_oauth2_external_requests_rate_limit_reset_in_seconds", labels), int(until.Seconds()), 2, "reset in")
				}
			} else {
				pass = pass && assert.Nil(t, promhelp.MetricValue(t, reg, "coderd_oauth2_external_requests_rate_limit", labels), "not exists")
			}

			// Helpful debugging
			if !pass {
				t.Log(promhelp.RegistryDump(reg))
			}
		})
	}
}

func must[V any](t *testing.T) func(v V, err error) V {
	return func(v V, err error) V {
		t.Helper()
		require.NoError(t, err)
		return v
	}
}

func TestExternalRequestRateLimitedMetric(t *testing.T) {
	t.Parallel()

	const metricName = "coderd_oauth2_external_requests_rate_limited_total"
	labels := func(status int) prometheus.Labels {
		return prometheus.Labels{
			"name":        "test",
			"source":      string(promoauth.SourceValidateToken),
			"status_code": fmt.Sprintf("%d", status),
		}
	}

	reg := prometheus.NewRegistry()
	cfg := promoauth.NewFactory(reg).New("test", &oauth2.Config{})
	ctx := testutil.Context(t, testutil.WaitShort)

	do := func(t *testing.T, status int, headers map[string]string) {
		t.Helper()
		srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
			for k, v := range headers {
				w.Header().Set(k, v)
			}
			w.WriteHeader(status)
		}))
		t.Cleanup(srv.Close)
		req, err := http.NewRequestWithContext(ctx, http.MethodGet, srv.URL, nil)
		require.NoError(t, err)
		resp, err := cfg.Do(ctx, promoauth.SourceValidateToken, req)
		require.NoError(t, err)
		_ = resp.Body.Close()
	}

	// A 429 is counted as rate-limited under its status code.
	do(t, http.StatusTooManyRequests, nil)
	assert.Equal(t, 1, promhelp.CounterValue(t, reg, metricName, labels(http.StatusTooManyRequests)))

	// A 403 carrying rate-limit headers is counted under its own status code.
	do(t, http.StatusForbidden, map[string]string{"X-RateLimit-Remaining": "0"})
	assert.Equal(t, 1, promhelp.CounterValue(t, reg, metricName, labels(http.StatusForbidden)))

	// A 200 and a plain 403 (a genuine revocation) are not counted.
	do(t, http.StatusOK, nil)
	do(t, http.StatusForbidden, nil)
	assert.Equal(t, 1, promhelp.CounterValue(t, reg, metricName, labels(http.StatusTooManyRequests)))
	assert.Equal(t, 1, promhelp.CounterValue(t, reg, metricName, labels(http.StatusForbidden)))
	assert.Nil(t, promhelp.MetricValue(t, reg, metricName, labels(http.StatusOK)))
}