chore: add faking 429 responses from fake idp (coder#12365)

Emyrk · web-flow · commit cbcf4ef2c4d4 · 2024-02-29T09:45:53.000-06:00
Required to trigger error condition in fe. See pull (coder#12367)
diff --git 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
@@ -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()
