From 784417bacc845a928670a0fc90c2712ac092fe26 Mon Sep 17 00:00:00 2001 From: Michael Quigley Date: Mon, 23 Jun 2025 15:08:43 -0400 Subject: [PATCH] better middleware management; return 429 when over limit (#968) --- controller/share.go | 83 +++++++++---------- rest_client_zrok/share/share_responses.go | 62 ++++++++++++++ rest_server_zrok/embedded_spec.go | 6 ++ .../operations/share/share_responses.go | 25 ++++++ sdk/python/src/docs/ShareApi.md | 1 + sdk/python/src/zrok_api/api/share_api.py | 3 + specs/zrok.yml | 2 + 7 files changed, 140 insertions(+), 42 deletions(-) diff --git a/controller/share.go b/controller/share.go index 7e67b69e..697556c6 100644 --- a/controller/share.go +++ b/controller/share.go @@ -34,17 +34,19 @@ func (h *shareHandler) Handle(params share.ShareParams, principal *rest_model_zr envId, err := h.validateEnvironment(params.Body.EnvZID, principal, trx) if err != nil { - return h.handleEnvironmentError(err) - } - - if err := h.checkLimits(envId, principal, params.Body.Reserved, params.Body.UniqueName != "", sdk.ShareMode(params.Body.ShareMode), sdk.BackendMode(params.Body.BackendMode), trx); err != nil { - logrus.Errorf("limits error: %v", err) + logrus.Errorf("error validating environment '%v' for '%v': %v", params.Body.EnvZID, principal.Email, err) return share.NewShareUnauthorized() } - accessGrantAcctIds, responder := h.processAccessGrants(params, principal, trx) - if responder != nil { - return responder + if err := h.checkLimits(envId, principal, params.Body.Reserved, params.Body.UniqueName != "", sdk.ShareMode(params.Body.ShareMode), sdk.BackendMode(params.Body.BackendMode), trx); err != nil { + logrus.Errorf("limits error for '%v': %v", principal.Email, err) + return share.NewShareTooManyRequests() + } + + accessGrantAcctIds, err := h.processAccessGrants(params, principal, trx) + if err != nil { + logrus.Errorf("error processing access grants: %v", err) + return share.NewShareInternalServerError() } edge, err := zrokEdgeSdk.Client(cfg.Ziti) @@ -53,25 +55,29 @@ func (h *shareHandler) Handle(params share.ShareParams, principal *rest_model_zr return share.NewShareInternalServerError() } - shrToken, responder := h.createShareToken(params.Body.Reserved, params.Body.UniqueName, trx) - if responder != nil { - return responder + shrToken, err := h.createShareToken(params.Body.Reserved, params.Body.UniqueName, trx) + if err != nil { + logrus.Errorf("error creating share token: %v", err) + return share.NewShareInternalServerError() } - shrZId, frontendEndpoints, responder := h.allocateResources(params, principal, edge, shrToken, trx) - if responder != nil { - return responder + shrZId, frontendEndpoints, err := h.allocateResources(params, principal, edge, shrToken, trx) + if err != nil { + logrus.Errorf("error allocating resources: %v", err) + return share.NewShareInternalServerError() } sshr := h.createShareRecord(shrZId, shrToken, params, frontendEndpoints) - sid, responder := h.saveShareAndGrants(sshr, envId, accessGrantAcctIds, trx) - if responder != nil { - return responder + sid, err := h.saveShareAndGrants(sshr, envId, accessGrantAcctIds, trx) + if err != nil { + logrus.Errorf("error saving share and grants: %v", err) + return share.NewShareInternalServerError() } - if responder := h.handleAuthSecrets(params, sid, sshr, trx); responder != nil { - return responder + if err := h.handleAuthSecrets(params, sid, sshr, trx); err != nil { + logrus.Errorf("error handling auth secrets: %v", err) + return share.NewShareInternalServerError() } if err := trx.Commit(); err != nil { @@ -104,21 +110,14 @@ func (h *shareHandler) validateEnvironment(envZId string, principal *rest_model_ return 0, errors.New("environment not found") } -func (h *shareHandler) handleEnvironmentError(err error) middleware.Responder { - if err.Error() == "environment not found" { - return share.NewShareUnauthorized() - } - return share.NewShareInternalServerError() -} - -func (h *shareHandler) processAccessGrants(params share.ShareParams, principal *rest_model_zrok.Principal, trx *sqlx.Tx) ([]int, middleware.Responder) { +func (h *shareHandler) processAccessGrants(params share.ShareParams, principal *rest_model_zrok.Principal, trx *sqlx.Tx) ([]int, error) { var accessGrantAcctIds []int if store.PermissionMode(params.Body.PermissionMode) == store.ClosedPermissionMode { for _, email := range params.Body.AccessGrants { acct, err := str.FindAccountWithEmail(email, trx) if err != nil { logrus.Errorf("unable to find account '%v' for share request from '%v'", email, principal.Email) - return nil, share.NewShareNotFound() + return nil, err } logrus.Debugf("found id '%d' for '%v'", acct.Id, acct.Email) accessGrantAcctIds = append(accessGrantAcctIds, acct.Id) @@ -127,35 +126,35 @@ func (h *shareHandler) processAccessGrants(params share.ShareParams, principal * return accessGrantAcctIds, nil } -func (h *shareHandler) createShareToken(reserved bool, uniqueName string, trx *sqlx.Tx) (string, middleware.Responder) { +func (h *shareHandler) createShareToken(reserved bool, uniqueName string, trx *sqlx.Tx) (string, error) { if !reserved || uniqueName == "" { token, err := createShareToken() if err != nil { logrus.Error(err) - return "", share.NewShareInternalServerError() + return "", err } return token, nil } if !util.IsValidUniqueName(uniqueName) { logrus.Errorf("invalid unique name '%v'", uniqueName) - return "", share.NewShareUnprocessableEntity() + return "", errors.New("invalid unique name") } shareExists, err := str.ShareWithTokenExists(uniqueName, trx) if err != nil { logrus.Errorf("error checking share for token collision: %v", err) - return "", share.NewUpdateShareInternalServerError() + return "", err } if shareExists { logrus.Errorf("token '%v' already exists; cannot create share", uniqueName) - return "", share.NewShareConflict() + return "", errors.New("token already exists") } return uniqueName, nil } -func (h *shareHandler) allocateResources(params share.ShareParams, principal *rest_model_zrok.Principal, edge *rest_management_api_client.ZitiEdgeManagement, shrToken string, trx *sqlx.Tx) (string, []string, middleware.Responder) { +func (h *shareHandler) allocateResources(params share.ShareParams, principal *rest_model_zrok.Principal, edge *rest_management_api_client.ZitiEdgeManagement, shrToken string, trx *sqlx.Tx) (string, []string, error) { var shrZId string var frontendEndpoints []string var err error @@ -167,12 +166,12 @@ func (h *shareHandler) allocateResources(params share.ShareParams, principal *re shrZId, frontendEndpoints, err = h.allocatePrivateResources(params, edge, shrToken) default: logrus.Errorf("unknown share mode '%v'", params.Body.ShareMode) - return "", nil, share.NewShareInternalServerError() + return "", nil, errors.New("unknown share mode") } if err != nil { logrus.Error(err) - return "", nil, share.NewShareInternalServerError() + return "", nil, err } return shrZId, frontendEndpoints, nil @@ -257,11 +256,11 @@ func (h *shareHandler) createShareRecord(shrZId string, shrToken string, params return sshr } -func (h *shareHandler) saveShareAndGrants(sshr *store.Share, envId int, accessGrantAcctIds []int, trx *sqlx.Tx) (int, middleware.Responder) { +func (h *shareHandler) saveShareAndGrants(sshr *store.Share, envId int, accessGrantAcctIds []int, trx *sqlx.Tx) (int, error) { sid, err := str.CreateShare(envId, sshr, trx) if err != nil { logrus.Errorf("error creating share record: %v", err) - return 0, share.NewShareInternalServerError() + return 0, err } if sshr.PermissionMode == store.ClosedPermissionMode { @@ -269,7 +268,7 @@ func (h *shareHandler) saveShareAndGrants(sshr *store.Share, envId int, accessGr _, err := str.CreateAccessGrant(sid, acctId, trx) if err != nil { logrus.Errorf("error creating access grant: %v", err) - return 0, share.NewShareInternalServerError() + return 0, err } } } @@ -277,7 +276,7 @@ func (h *shareHandler) saveShareAndGrants(sshr *store.Share, envId int, accessGr return sid, nil } -func (h *shareHandler) handleAuthSecrets(params share.ShareParams, sid int, sshr *store.Share, trx *sqlx.Tx) middleware.Responder { +func (h *shareHandler) handleAuthSecrets(params share.ShareParams, sid int, sshr *store.Share, trx *sqlx.Tx) error { if sshr.ShareMode == string(sdk.PublicShareMode) && params.Body.AuthScheme == string(sdk.Basic) { logrus.Infof("writing basic auth secrets for '%v'", sshr.Token) authUsersMap := make(map[string]string) @@ -287,7 +286,7 @@ func (h *shareHandler) handleAuthSecrets(params share.ShareParams, sid int, sshr authUsersMapJson, err := json.Marshal(authUsersMap) if err != nil { logrus.Errorf("error marshalling auth secrets for '%v': %v", sshr.Token, err) - return share.NewShareInternalServerError() + return err } secrets := store.Secrets{ ShareId: sid, @@ -298,7 +297,7 @@ func (h *shareHandler) handleAuthSecrets(params share.ShareParams, sid int, sshr } if err := str.CreateSecrets(secrets, trx); err != nil { logrus.Errorf("error creating secrets: %v", err) - return share.NewShareInternalServerError() + return err } logrus.Infof("wrote auth secrets for '%v'", sshr.Token) } diff --git a/rest_client_zrok/share/share_responses.go b/rest_client_zrok/share/share_responses.go index 8af0e374..46c170ac 100644 --- a/rest_client_zrok/share/share_responses.go +++ b/rest_client_zrok/share/share_responses.go @@ -53,6 +53,12 @@ func (o *ShareReader) ReadResponse(response runtime.ClientResponse, consumer run return nil, err } return nil, result + case 429: + result := NewShareTooManyRequests() + if err := result.readResponse(response, consumer, o.formats); err != nil { + return nil, err + } + return nil, result case 500: result := NewShareInternalServerError() if err := result.readResponse(response, consumer, o.formats); err != nil { @@ -356,6 +362,62 @@ func (o *ShareUnprocessableEntity) readResponse(response runtime.ClientResponse, return nil } +// NewShareTooManyRequests creates a ShareTooManyRequests with default headers values +func NewShareTooManyRequests() *ShareTooManyRequests { + return &ShareTooManyRequests{} +} + +/* +ShareTooManyRequests describes a response with status code 429, with default header values. + +over limit +*/ +type ShareTooManyRequests struct { +} + +// IsSuccess returns true when this share too many requests response has a 2xx status code +func (o *ShareTooManyRequests) IsSuccess() bool { + return false +} + +// IsRedirect returns true when this share too many requests response has a 3xx status code +func (o *ShareTooManyRequests) IsRedirect() bool { + return false +} + +// IsClientError returns true when this share too many requests response has a 4xx status code +func (o *ShareTooManyRequests) IsClientError() bool { + return true +} + +// IsServerError returns true when this share too many requests response has a 5xx status code +func (o *ShareTooManyRequests) IsServerError() bool { + return false +} + +// IsCode returns true when this share too many requests response a status code equal to that given +func (o *ShareTooManyRequests) IsCode(code int) bool { + return code == 429 +} + +// Code gets the status code for the share too many requests response +func (o *ShareTooManyRequests) Code() int { + return 429 +} + +func (o *ShareTooManyRequests) Error() string { + return fmt.Sprintf("[POST /share][%d] shareTooManyRequests ", 429) +} + +func (o *ShareTooManyRequests) String() string { + return fmt.Sprintf("[POST /share][%d] shareTooManyRequests ", 429) +} + +func (o *ShareTooManyRequests) readResponse(response runtime.ClientResponse, consumer runtime.Consumer, formats strfmt.Registry) error { + + return nil +} + // NewShareInternalServerError creates a ShareInternalServerError with default headers values func NewShareInternalServerError() *ShareInternalServerError { return &ShareInternalServerError{} diff --git a/rest_server_zrok/embedded_spec.go b/rest_server_zrok/embedded_spec.go index 0f15ad1d..cf402b88 100644 --- a/rest_server_zrok/embedded_spec.go +++ b/rest_server_zrok/embedded_spec.go @@ -2298,6 +2298,9 @@ func init() { "422": { "description": "unprocessable" }, + "429": { + "description": "over limit" + }, "500": { "description": "internal server error", "schema": { @@ -5194,6 +5197,9 @@ func init() { "422": { "description": "unprocessable" }, + "429": { + "description": "over limit" + }, "500": { "description": "internal server error", "schema": { diff --git a/rest_server_zrok/operations/share/share_responses.go b/rest_server_zrok/operations/share/share_responses.go index 9bf46588..aeb0e90f 100644 --- a/rest_server_zrok/operations/share/share_responses.go +++ b/rest_server_zrok/operations/share/share_responses.go @@ -158,6 +158,31 @@ func (o *ShareUnprocessableEntity) WriteResponse(rw http.ResponseWriter, produce rw.WriteHeader(422) } +// ShareTooManyRequestsCode is the HTTP code returned for type ShareTooManyRequests +const ShareTooManyRequestsCode int = 429 + +/* +ShareTooManyRequests over limit + +swagger:response shareTooManyRequests +*/ +type ShareTooManyRequests struct { +} + +// NewShareTooManyRequests creates ShareTooManyRequests with default headers values +func NewShareTooManyRequests() *ShareTooManyRequests { + + return &ShareTooManyRequests{} +} + +// WriteResponse to the client +func (o *ShareTooManyRequests) WriteResponse(rw http.ResponseWriter, producer runtime.Producer) { + + rw.Header().Del(runtime.HeaderContentType) //Remove Content-Type on empty responses + + rw.WriteHeader(429) +} + // ShareInternalServerErrorCode is the HTTP code returned for type ShareInternalServerError const ShareInternalServerErrorCode int = 500 diff --git a/sdk/python/src/docs/ShareApi.md b/sdk/python/src/docs/ShareApi.md index 2831f1d2..b3fc247f 100644 --- a/sdk/python/src/docs/ShareApi.md +++ b/sdk/python/src/docs/ShareApi.md @@ -166,6 +166,7 @@ Name | Type | Description | Notes **404** | not found | - | **409** | conflict | - | **422** | unprocessable | - | +**429** | over limit | - | **500** | internal server error | - | [[Back to top]](#) [[Back to API list]](../README.md#documentation-for-api-endpoints) [[Back to Model list]](../README.md#documentation-for-models) [[Back to README]](../README.md) diff --git a/sdk/python/src/zrok_api/api/share_api.py b/sdk/python/src/zrok_api/api/share_api.py index 626035ee..79311bca 100644 --- a/sdk/python/src/zrok_api/api/share_api.py +++ b/sdk/python/src/zrok_api/api/share_api.py @@ -382,6 +382,7 @@ class ShareApi: '404': None, '409': None, '422': None, + '429': None, '500': "str", } response_data = self.api_client.call_api( @@ -453,6 +454,7 @@ class ShareApi: '404': None, '409': None, '422': None, + '429': None, '500': "str", } response_data = self.api_client.call_api( @@ -524,6 +526,7 @@ class ShareApi: '404': None, '409': None, '422': None, + '429': None, '500': "str", } response_data = self.api_client.call_api( diff --git a/specs/zrok.yml b/specs/zrok.yml index 76ab90ef..7310c204 100644 --- a/specs/zrok.yml +++ b/specs/zrok.yml @@ -1497,6 +1497,8 @@ paths: description: conflict 422: description: unprocessable + 429: + description: over limit 500: description: internal server error schema: