From 96f7a6c1489377db3285aa09075e4bf7c664f6d1 Mon Sep 17 00:00:00 2001 From: Scott Leggett Date: Tue, 16 Jan 2024 23:52:38 +0800 Subject: [PATCH] feat: rate limit keycloak API access --- cmd/ssh-portal-api/serve.go | 12 ++++++++---- cmd/ssh-token/serve.go | 15 +++++++++++---- go.mod | 2 +- internal/keycloak/client.go | 5 ++++- internal/keycloak/jwt_test.go | 2 +- internal/keycloak/useraccesstoken.go | 8 ++++++++ internal/keycloak/userrolesandgroups.go | 4 ++++ 7 files changed, 37 insertions(+), 11 deletions(-) diff --git a/cmd/ssh-portal-api/serve.go b/cmd/ssh-portal-api/serve.go index 0517e9af..feddb18d 100644 --- a/cmd/ssh-portal-api/serve.go +++ b/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) diff --git a/cmd/ssh-token/serve.go b/cmd/ssh-token/serve.go index 59d988ce..8bb29799 100644 --- a/cmd/ssh-token/serve.go +++ b/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) } diff --git a/go.mod b/go.mod index 8610236d..98c398f7 100644 --- 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 index c5a37a3e..1d38f55f 100644 --- 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,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 } diff --git a/internal/keycloak/jwt_test.go b/internal/keycloak/jwt_test.go index b57d7c80..af125d81 100644 --- a/internal/keycloak/jwt_test.go +++ b/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) } diff --git a/internal/keycloak/useraccesstoken.go b/internal/keycloak/useraccesstoken.go index 8299e68e..5e7bd1d9 100644 --- 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 index a17b83ae..68bc7f6d 100644 --- 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,