feat: rate limit keycloak API access

uselagoon · Jan 16, 2024 · 0be19ad · 0be19ad
1 parent 321a508
commit 0be19ad
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 1 deletion.
diff --git a/go.mod b/go.mod
@@ -20,6 +20,7 @@ require (
 	golang.org/x/exp v0.0.0-20230817173708-d852ddb80c63
 	golang.org/x/oauth2 v0.16.0
 	golang.org/x/sync v0.6.0
+	golang.org/x/time v0.3.0
 	k8s.io/api v0.29.0
 	k8s.io/apimachinery v0.29.0
 	k8s.io/client-go v0.29.0
@@ -71,7 +72,6 @@ require (
 	golang.org/x/sys v0.16.0 // indirect
 	golang.org/x/term v0.16.0 // indirect
 	golang.org/x/text v0.14.0 // indirect
-	golang.org/x/time v0.3.0 // indirect
 	google.golang.org/appengine v1.6.7 // indirect
 	google.golang.org/protobuf v1.31.0 // indirect
 	gopkg.in/inf.v0 v0.9.1 // indirect

diff --git a/internal/keycloak/client.go b/internal/keycloak/client.go
@@ -14,6 +14,7 @@ import (
 	"github.com/MicahParks/keyfunc/v2"
 	oidcClient "github.com/zitadel/oidc/v3/pkg/client"
 	"github.com/zitadel/oidc/v3/pkg/oidc"
+	"golang.org/x/time/rate"
 )
 
 const pkgName = "github.com/uselagoon/ssh-portal/internal/keycloak"
@@ -25,6 +26,7 @@ type Client struct {
 	jwks         *keyfunc.JWKS
 	log          *slog.Logger
 	oidcConfig   *oidc.DiscoveryConfiguration
+	limiter      *rate.Limiter
 }
 
 // NewClient creates a new keycloak client for the lagoon realm.
@@ -53,5 +55,6 @@ func NewClient(ctx context.Context, log *slog.Logger, keycloakURL, clientID,
 		jwks:         jwks,
 		log:          log,
 		oidcConfig:   oidcConfig,
+		limiter:      rate.NewLimiter(10, 10), // 10 requests per second
 	}, nil
 }
diff --git a/internal/keycloak/useraccesstoken.go b/internal/keycloak/useraccesstoken.go
@@ -54,6 +54,10 @@ func (c *Client) UserAccessTokenResponse(ctx context.Context,
 	// set up tracing
 	ctx, span := otel.Tracer(pkgName).Start(ctx, "UserAccessToken")
 	defer span.End()
+	// rate limit keycloak API access
+	if err := c.limiter.Wait(ctx); err != nil {
+		return "", fmt.Errorf("couldn't wait for limiter: %v", err)
+	}
 	// get user token
 	userToken, err := c.getUserToken(ctx, userUUID)
 	if err != nil {
@@ -74,6 +78,10 @@ func (c *Client) UserAccessToken(ctx context.Context,
 	// set up tracing
 	ctx, span := otel.Tracer(pkgName).Start(ctx, "UserAccessToken")
 	defer span.End()
+	// rate limit keycloak API access
+	if err := c.limiter.Wait(ctx); err != nil {
+		return "", fmt.Errorf("couldn't wait for limiter: %v", err)
+	}
 	// get user token
 	userToken, err := c.getUserToken(ctx, userUUID)
 	if err != nil {

diff --git a/internal/keycloak/userrolesandgroups.go b/internal/keycloak/userrolesandgroups.go
@@ -19,6 +19,10 @@ func (c *Client) UserRolesAndGroups(ctx context.Context,
 	// set up tracing
 	ctx, span := otel.Tracer(pkgName).Start(ctx, "UserRolesAndGroups")
 	defer span.End()
+	// rate limit keycloak API access
+	if err := c.limiter.Wait(ctx); err != nil {
+		return nil, nil, nil, fmt.Errorf("couldn't wait for limiter: %v", err)
+	}
 	// get user token
 	userConfig := oauth2.Config{
 		ClientID:     c.clientID,