diff --git a/coderd/coderdtest/oidctest/idp.go b/coderd/coderdtest/oidctest/idp.go index a185332e87..ff42e31997 100644 --- a/coderd/coderdtest/oidctest/idp.go +++ b/coderd/coderdtest/oidctest/idp.go @@ -244,6 +244,56 @@ func WithIssuer(issuer string) func(*FakeIDP) { } } +type With429Arguments struct { + AllPaths bool + TokenPath bool + AuthorizePath bool + KeysPath bool + UserInfoPath bool + DeviceAuth bool + DeviceVerify bool +} + +// With429 will emulate a 429 response for the selected paths. +func With429(params With429Arguments) func(*FakeIDP) { + return func(f *FakeIDP) { + f.middlewares = append(f.middlewares, func(next http.Handler) http.Handler { + return http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) { + if params.AllPaths { + http.Error(rw, "429, being manually blocked (all)", http.StatusTooManyRequests) + return + } + if params.TokenPath && strings.Contains(r.URL.Path, tokenPath) { + http.Error(rw, "429, being manually blocked (token)", http.StatusTooManyRequests) + return + } + if params.AuthorizePath && strings.Contains(r.URL.Path, authorizePath) { + http.Error(rw, "429, being manually blocked (authorize)", http.StatusTooManyRequests) + return + } + if params.KeysPath && strings.Contains(r.URL.Path, keysPath) { + http.Error(rw, "429, being manually blocked (keys)", http.StatusTooManyRequests) + return + } + if params.UserInfoPath && strings.Contains(r.URL.Path, userInfoPath) { + http.Error(rw, "429, being manually blocked (userinfo)", http.StatusTooManyRequests) + return + } + if params.DeviceAuth && strings.Contains(r.URL.Path, deviceAuth) { + http.Error(rw, "429, being manually blocked (device-auth)", http.StatusTooManyRequests) + return + } + if params.DeviceVerify && strings.Contains(r.URL.Path, deviceVerify) { + http.Error(rw, "429, being manually blocked (device-verify)", http.StatusTooManyRequests) + return + } + + next.ServeHTTP(rw, r) + }) + }) + } +} + const ( // nolint:gosec // It thinks this is a secret lol tokenPath = "/oauth2/token" diff --git a/scripts/testidp/main.go b/scripts/testidp/main.go index 82fc10c936..9feae443ab 100644 --- a/scripts/testidp/main.go +++ b/scripts/testidp/main.go @@ -6,6 +6,7 @@ import ( "log" "os" "os/signal" + "strings" "testing" "time" @@ -25,7 +26,8 @@ var ( clientSecret = flag.String("client-sec", "static-client-secret", "Client Secret, set empty to be random") deviceFlow = flag.Bool("device-flow", false, "Enable device flow") // By default, no regex means it will never match anything. So at least default to matching something. - extRegex = flag.String("ext-regex", `^(https?://)?example\.com(/.*)?$`, "External auth regex") + extRegex = flag.String("ext-regex", `^(https?://)?example\.com(/.*)?$`, "External auth regex") + tooManyRequests = flag.String("429", "", "Simulate too many requests for a given endpoint.") ) func main() { @@ -54,6 +56,31 @@ type withClientSecret struct { // RunIDP needs the testing.T because our oidctest package requires the // testing.T. func RunIDP() func(t *testing.T) { + tooManyRequestParams := oidctest.With429Arguments{} + if *tooManyRequests != "" { + for _, v := range strings.Split(*tooManyRequests, ",") { + v = strings.ToLower(strings.TrimSpace(v)) + switch v { + case "all": + tooManyRequestParams.AllPaths = true + case "auth": + tooManyRequestParams.AuthorizePath = true + case "token": + tooManyRequestParams.TokenPath = true + case "keys": + tooManyRequestParams.KeysPath = true + case "userinfo": + tooManyRequestParams.UserInfoPath = true + case "device": + tooManyRequestParams.DeviceAuth = true + case "device-verify": + tooManyRequestParams.DeviceVerify = true + default: + log.Printf("Unknown too-many-requests value: %s\nView the `testidp/main.go` for valid values.", v) + } + } + } + return func(t *testing.T) { idp := oidctest.NewFakeIDP(t, oidctest.WithServing(), @@ -63,6 +90,7 @@ func RunIDP() func(t *testing.T) { oidctest.WithStaticCredentials(*clientID, *clientSecret), oidctest.WithIssuer("http://localhost:4500"), oidctest.WithLogger(slog.Make(sloghuman.Sink(os.Stderr))), + oidctest.With429(tooManyRequestParams), ) id, sec := idp.AppCredentials() prov := idp.WellknownConfig()