Merge pull request #354 from uselagoon/keycloak-rate-limit

feat: rate limit keycloak API access

cmd/ssh-portal-api/serve.go

@@ -25,6 +25,7 @@ type ServeCmd struct {
 	KeycloakBaseURL      string `kong:"required,env='KEYCLOAK_BASE_URL',help='Keycloak Base URL'"`
 	KeycloakClientID     string `kong:"default='service-api',env='KEYCLOAK_SERVICE_API_CLIENT_ID',help='Keycloak OAuth2 Client ID'"`
 	KeycloakClientSecret string `kong:"required,env='KEYCLOAK_SERVICE_API_CLIENT_SECRET',help='Keycloak OAuth2 Client Secret'"`
+	KeycloakRateLimit    int    `kong:"default=10,env='KEYCLOAK_RATE_LIMIT',help='Keycloak API Rate Limit (requests/second)'"`
 	NATSURL              string `kong:"required,env='NATS_URL',help='NATS server URL (nats://... or tls://...)'"`
 }
 
@@ -53,13 +54,16 @@ func (cmd *ServeCmd) Run(log *slog.Logger) error {
 	dbConf.User = cmd.APIDBUsername
 	l, err := lagoondb.NewClient(ctx, dbConf.FormatDSN())
 	if err != nil {
-		return fmt.Errorf("couldn't init lagoon DBClient: %v", err)
+		return fmt.Errorf("couldn't init lagoondb client: %v", err)
 	}
 	// init keycloak client
-	k, err := keycloak.NewClient(ctx, log, cmd.KeycloakBaseURL,
-		cmd.KeycloakClientID, cmd.KeycloakClientSecret)
+	k, err := keycloak.NewClient(ctx, log,
+		cmd.KeycloakBaseURL,
+		cmd.KeycloakClientID,
+		cmd.KeycloakClientSecret,
+		cmd.KeycloakRateLimit)
 	if err != nil {
-		return fmt.Errorf("couldn't init keycloak Client: %v", err)
+		return fmt.Errorf("couldn't init keycloak client: %v", err)
 	}
 	// start serving NATS requests
 	return sshportalapi.ServeNATS(ctx, stop, log, p, l, k, cmd.NATSURL)

cmd/ssh-token/serve.go

@@ -29,6 +29,7 @@ type ServeCmd struct {
 	KeycloakBaseURL                string `kong:"required,env='KEYCLOAK_BASE_URL',help='Keycloak Base URL'"`
 	KeycloakPermissionClientID     string `kong:"default='service-api',env='KEYCLOAK_SERVICE_API_CLIENT_ID',help='Keycloak service-api OAuth2 Client ID'"`
 	KeycloakPermissionClientSecret string `kong:"env='KEYCLOAK_SERVICE_API_CLIENT_SECRET',help='Keycloak service-api OAuth2 Client Secret'"`
+	KeycloakRateLimit              int    `kong:"default=10,env='KEYCLOAK_RATE_LIMIT',help='Keycloak API Rate Limit (requests/second)'"`
 	KeycloakTokenClientID          string `kong:"default='auth-server',env='KEYCLOAK_AUTH_SERVER_CLIENT_ID',help='Keycloak auth-server OAuth2 Client ID'"`
 	KeycloakTokenClientSecret      string `kong:"required,env='KEYCLOAK_AUTH_SERVER_CLIENT_SECRET',help='Keycloak auth-server OAuth2 Client Secret'"`
 	SSHServerPort                  uint   `kong:"default='2222',env='SSH_SERVER_PORT',help='Port the SSH server will listen on for SSH client connections'"`
@@ -62,14 +63,20 @@ func (cmd *ServeCmd) Run(log *slog.Logger) error {
 		return fmt.Errorf("couldn't init lagoonDB client: %v", err)
 	}
 	// init token / auth-server keycloak client
-	keycloakToken, err := keycloak.NewClient(ctx, log, cmd.KeycloakBaseURL,
-		cmd.KeycloakTokenClientID, cmd.KeycloakTokenClientSecret)
+	keycloakToken, err := keycloak.NewClient(ctx, log,
+		cmd.KeycloakBaseURL,
+		cmd.KeycloakTokenClientID,
+		cmd.KeycloakTokenClientSecret,
+		cmd.KeycloakRateLimit)
 	if err != nil {
 		return fmt.Errorf("couldn't init keycloak token client: %v", err)
 	}
 	// init permission / service-api keycloak client
-	keycloakPermission, err := keycloak.NewClient(ctx, log, cmd.KeycloakBaseURL,
-		cmd.KeycloakPermissionClientID, cmd.KeycloakPermissionClientSecret)
+	keycloakPermission, err := keycloak.NewClient(ctx, log,
+		cmd.KeycloakBaseURL,
+		cmd.KeycloakPermissionClientID,
+		cmd.KeycloakPermissionClientSecret,
+		cmd.KeycloakRateLimit)
 	if err != nil {
 		return fmt.Errorf("couldn't init keycloak permission client: %v", err)
 	}

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

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,11 +26,12 @@ 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.
 func NewClient(ctx context.Context, log *slog.Logger, keycloakURL, clientID,
-	clientSecret string) (*Client, error) {
+	clientSecret string, rateLimit int) (*Client, error) {
 	// discover OIDC config
 	issuerURL, err := url.Parse(keycloakURL)
 	if err != nil {
@@ -53,5 +55,6 @@ func NewClient(ctx context.Context, log *slog.Logger, keycloakURL, clientID,
 		jwks:         jwks,
 		log:          log,
 		oidcConfig:   oidcConfig,
+		limiter:      rate.NewLimiter(rate.Limit(rateLimit), rateLimit),
 	}, nil
 }

internal/keycloak/jwt_test.go

@@ -261,7 +261,7 @@ func TestValidateTokenClaims(t *testing.T) {
 			// NOTE: client secret is empty because it isn't used in this test, but
 			// client ID is checked against azp in the token.
 			k, err := keycloak.NewClient(context.Background(), log, ts.URL,
-				"auth-server", "")
+				"auth-server", "", 10)
 			if err != nil {
 				tt.Fatal(err)
 			}

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 {

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,