Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix: updating ratelimits for identities #2067

Merged
merged 1 commit into from
Sep 5, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion apps/agent/integration/cluster/docker/ratelimits_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ import (
"time"

"github.com/stretchr/testify/require"
"github.com/unkeyed/unkey/apps/agent/gen/openapi"
"github.com/unkeyed/unkey/apps/agent/pkg/openapi"

"github.com/unkeyed/unkey/apps/agent/pkg/uid"

Expand Down
2 changes: 1 addition & 1 deletion apps/agent/pkg/api/errors/internal_server_error.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,8 @@ import (
"net/http"

"github.com/Southclaws/fault/fmsg"
"github.com/unkeyed/unkey/apps/agent/gen/openapi"
"github.com/unkeyed/unkey/apps/agent/pkg/api/ctxutil"
"github.com/unkeyed/unkey/apps/agent/pkg/openapi"
)

// HandleError takes in any unforseen error and returns a BaseError to be sent to the client
Expand Down
2 changes: 1 addition & 1 deletion apps/agent/pkg/api/errors/validation_error.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,8 @@ import (
"net/http"

"github.com/Southclaws/fault/fmsg"
"github.com/unkeyed/unkey/apps/agent/gen/openapi"
"github.com/unkeyed/unkey/apps/agent/pkg/api/ctxutil"
"github.com/unkeyed/unkey/apps/agent/pkg/openapi"
)

func HandleValidationError(ctx context.Context, err error) openapi.ValidationError {
Expand Down
2 changes: 2 additions & 0 deletions apps/agent/pkg/api/register_routes.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ package api
import (
"github.com/unkeyed/unkey/apps/agent/pkg/api/routes"
notFound "github.com/unkeyed/unkey/apps/agent/pkg/api/routes/not_found"
openapi "github.com/unkeyed/unkey/apps/agent/pkg/api/routes/openapi"
v1Liveness "github.com/unkeyed/unkey/apps/agent/pkg/api/routes/v1_liveness"
v1RatelimitCommitLease "github.com/unkeyed/unkey/apps/agent/pkg/api/routes/v1_ratelimit_commitLease"
v1RatelimitMultiRatelimit "github.com/unkeyed/unkey/apps/agent/pkg/api/routes/v1_ratelimit_multiRatelimit"
Expand All @@ -27,6 +28,7 @@ func (s *Server) RegisterRoutes() {
staticBearerAuth := newBearerAuthMiddleware(s.authToken)

v1Liveness.New(svc).Register(s.mux)
openapi.New(svc).Register(s.mux)

v1RatelimitCommitLease.New(svc).
WithMiddleware(staticBearerAuth).
Expand Down
2 changes: 1 addition & 1 deletion apps/agent/pkg/api/routes/not_found/handler.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,9 +3,9 @@ package notFound
import (
"net/http"

"github.com/unkeyed/unkey/apps/agent/gen/openapi"
"github.com/unkeyed/unkey/apps/agent/pkg/api/ctxutil"
"github.com/unkeyed/unkey/apps/agent/pkg/api/routes"
"github.com/unkeyed/unkey/apps/agent/pkg/openapi"
)

// This is a hack, because / matches everything, so we need to make sure this is the last route
Expand Down
21 changes: 21 additions & 0 deletions apps/agent/pkg/api/routes/openapi/handler.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,21 @@
package openapi

import (
"net/http"

"github.com/unkeyed/unkey/apps/agent/pkg/api/routes"
"github.com/unkeyed/unkey/apps/agent/pkg/openapi"
)

func New(svc routes.Services) *routes.Route {

return routes.NewRoute("GET", "/openapi.json",
func(w http.ResponseWriter, r *http.Request) {

w.WriteHeader(200)
w.Header().Set("Content-Type", "application/json")
w.Write(openapi.Spec)

},
)
}
Comment on lines +10 to +21
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Add error handling for response writing.

The function writes the OpenAPI specification to the response but does not handle potential errors from the w.Write method. It's recommended to handle this error to ensure robustness, especially in production environments.

- w.Write(openapi.Spec)
+ if _, err := w.Write(openapi.Spec); err != nil {
+     http.Error(w, "Failed to write OpenAPI spec", http.StatusInternalServerError)
+     return
+ }
Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
func New(svc routes.Services) *routes.Route {
return routes.NewRoute("GET", "/openapi.json",
func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(200)
w.Header().Set("Content-Type", "application/json")
w.Write(openapi.Spec)
},
)
}
func New(svc routes.Services) *routes.Route {
return routes.NewRoute("GET", "/openapi.json",
func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(200)
w.Header().Set("Content-Type", "application/json")
if _, err := w.Write(openapi.Spec); err != nil {
http.Error(w, "Failed to write OpenAPI spec", http.StatusInternalServerError)
return
}
},
)
}

Refactor for improved readability and maintainability.

Consider refactoring the handler function to separate concerns and improve readability. This can be achieved by defining the handler logic in a separate function, which can then be tested more easily.

+ func handleOpenAPISpec(w http.ResponseWriter, r *http.Request) {
+     w.WriteHeader(http.StatusOK)
+     w.Header().Set("Content-Type", "application/json")
+     if _, err := w.Write(openapi.Spec); err != nil {
+         http.Error(w, "Failed to write OpenAPI spec", http.StatusInternalServerError)
+         return
+     }
+ }
+
  func New(svc routes.Services) *routes.Route {
      return routes.NewRoute("GET", "/openapi.json", handleOpenAPISpec)
  }
Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
func New(svc routes.Services) *routes.Route {
return routes.NewRoute("GET", "/openapi.json",
func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(200)
w.Header().Set("Content-Type", "application/json")
w.Write(openapi.Spec)
},
)
}
func handleOpenAPISpec(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusOK)
w.Header().Set("Content-Type", "application/json")
if _, err := w.Write(openapi.Spec); err != nil {
http.Error(w, "Failed to write OpenAPI spec", http.StatusInternalServerError)
return
}
}
func New(svc routes.Services) *routes.Route {
return routes.NewRoute("GET", "/openapi.json", handleOpenAPISpec)
}

2 changes: 1 addition & 1 deletion apps/agent/pkg/api/routes/sender.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,9 +5,9 @@ import (
"encoding/json"
"net/http"

"github.com/unkeyed/unkey/apps/agent/gen/openapi"
"github.com/unkeyed/unkey/apps/agent/pkg/api/ctxutil"
"github.com/unkeyed/unkey/apps/agent/pkg/logging"
"github.com/unkeyed/unkey/apps/agent/pkg/openapi"
)

type Sender interface {
Expand Down
2 changes: 1 addition & 1 deletion apps/agent/pkg/api/routes/v1_liveness/handler.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,8 @@ package v1Liveness
import (
"net/http"

"github.com/unkeyed/unkey/apps/agent/gen/openapi"
"github.com/unkeyed/unkey/apps/agent/pkg/api/routes"
"github.com/unkeyed/unkey/apps/agent/pkg/openapi"
)

func New(svc routes.Services) *routes.Route {
Expand Down
2 changes: 1 addition & 1 deletion apps/agent/pkg/api/routes/v1_liveness/handler_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,9 +4,9 @@ import (
"testing"

"github.com/stretchr/testify/require"
"github.com/unkeyed/unkey/apps/agent/gen/openapi"
v1Liveness "github.com/unkeyed/unkey/apps/agent/pkg/api/routes/v1_liveness"
"github.com/unkeyed/unkey/apps/agent/pkg/api/testutil"
"github.com/unkeyed/unkey/apps/agent/pkg/openapi"
)

func TestLiveness(t *testing.T) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -6,10 +6,10 @@ import (
"github.com/Southclaws/fault"
"github.com/Southclaws/fault/fmsg"
"github.com/btcsuite/btcutil/base58"
"github.com/unkeyed/unkey/apps/agent/gen/openapi"
ratelimitv1 "github.com/unkeyed/unkey/apps/agent/gen/proto/ratelimit/v1"
"github.com/unkeyed/unkey/apps/agent/pkg/api/errors"
"github.com/unkeyed/unkey/apps/agent/pkg/api/routes"
"github.com/unkeyed/unkey/apps/agent/pkg/openapi"
"google.golang.org/protobuf/proto"
)

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -5,10 +5,10 @@ import (
"time"

"github.com/stretchr/testify/require"
"github.com/unkeyed/unkey/apps/agent/gen/openapi"
v1RatelimitCommitLease "github.com/unkeyed/unkey/apps/agent/pkg/api/routes/v1_ratelimit_commitLease"
v1RatelimitRatelimit "github.com/unkeyed/unkey/apps/agent/pkg/api/routes/v1_ratelimit_ratelimit"
"github.com/unkeyed/unkey/apps/agent/pkg/api/testutil"
"github.com/unkeyed/unkey/apps/agent/pkg/openapi"
"github.com/unkeyed/unkey/apps/agent/pkg/uid"
"github.com/unkeyed/unkey/apps/agent/pkg/util"
)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -3,10 +3,10 @@ package v1RatelimitMultiRatelimit
import (
"net/http"

"github.com/unkeyed/unkey/apps/agent/gen/openapi"
ratelimitv1 "github.com/unkeyed/unkey/apps/agent/gen/proto/ratelimit/v1"
"github.com/unkeyed/unkey/apps/agent/pkg/api/errors"
"github.com/unkeyed/unkey/apps/agent/pkg/api/routes"
"github.com/unkeyed/unkey/apps/agent/pkg/openapi"
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Optimize the handling of ratelimits and response construction.

Consider the following optimizations:

  • Pre-allocate the ratelimits array with the exact required capacity to avoid potential reallocations.
  • Use a similar approach for the res.Ratelimits array to ensure efficient memory usage.

Apply this diff to optimize the handling of the ratelimits array and the response construction:

- ratelimits := make([]*ratelimitv1.RatelimitRequest, len(req.Ratelimits))
+ ratelimits := make([]*ratelimitv1.RatelimitRequest, 0, len(req.Ratelimits))
- res.Ratelimits = make([]openapi.SingleRatelimitResponse, len(res.Ratelimits))
+ res.Ratelimits = make([]openapi.SingleRatelimitResponse, 0, len(svcRes.Ratelimits))

Committable suggestion was skipped due to low confidence.

)

func New(svc routes.Services) *routes.Route {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -4,10 +4,10 @@ import (
"net/http"

"github.com/btcsuite/btcutil/base58"
"github.com/unkeyed/unkey/apps/agent/gen/openapi"
ratelimitv1 "github.com/unkeyed/unkey/apps/agent/gen/proto/ratelimit/v1"
"github.com/unkeyed/unkey/apps/agent/pkg/api/errors"
"github.com/unkeyed/unkey/apps/agent/pkg/api/routes"
"github.com/unkeyed/unkey/apps/agent/pkg/openapi"
"github.com/unkeyed/unkey/apps/agent/pkg/util"
"google.golang.org/protobuf/proto"
)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -5,9 +5,9 @@ import (
"time"

"github.com/stretchr/testify/require"
"github.com/unkeyed/unkey/apps/agent/gen/openapi"
v1RatelimitRatelimit "github.com/unkeyed/unkey/apps/agent/pkg/api/routes/v1_ratelimit_ratelimit"
"github.com/unkeyed/unkey/apps/agent/pkg/api/testutil"
"github.com/unkeyed/unkey/apps/agent/pkg/openapi"
"github.com/unkeyed/unkey/apps/agent/pkg/uid"
)

Expand Down
2 changes: 1 addition & 1 deletion apps/agent/pkg/api/routes/v1_vault_decrypt/handler.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,10 +5,10 @@ import (

"github.com/Southclaws/fault"
"github.com/Southclaws/fault/fmsg"
"github.com/unkeyed/unkey/apps/agent/gen/openapi"
vaultv1 "github.com/unkeyed/unkey/apps/agent/gen/proto/vault/v1"
"github.com/unkeyed/unkey/apps/agent/pkg/api/errors"
"github.com/unkeyed/unkey/apps/agent/pkg/api/routes"
"github.com/unkeyed/unkey/apps/agent/pkg/openapi"
)

func New(svc routes.Services) *routes.Route {
Expand Down
2 changes: 1 addition & 1 deletion apps/agent/pkg/api/routes/v1_vault_encrypt/handler.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,10 +3,10 @@ package v1VaultEncrypt
import (
"net/http"

"github.com/unkeyed/unkey/apps/agent/gen/openapi"
vaultv1 "github.com/unkeyed/unkey/apps/agent/gen/proto/vault/v1"
"github.com/unkeyed/unkey/apps/agent/pkg/api/errors"
"github.com/unkeyed/unkey/apps/agent/pkg/api/routes"
"github.com/unkeyed/unkey/apps/agent/pkg/openapi"
)

func New(svc routes.Services) *routes.Route {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -5,10 +5,10 @@ import (

"github.com/Southclaws/fault"
"github.com/Southclaws/fault/fmsg"
"github.com/unkeyed/unkey/apps/agent/gen/openapi"
vaultv1 "github.com/unkeyed/unkey/apps/agent/gen/proto/vault/v1"
"github.com/unkeyed/unkey/apps/agent/pkg/api/errors"
"github.com/unkeyed/unkey/apps/agent/pkg/api/routes"
"github.com/unkeyed/unkey/apps/agent/pkg/openapi"
)

type Request = openapi.V1EncryptBulkRequestBody
Expand Down
2 changes: 1 addition & 1 deletion apps/agent/pkg/api/server.go
Original file line number Diff line number Diff line change
Expand Up @@ -83,7 +83,7 @@ func New(config Config) (*Server, error) {
// // validationMiddleware,
// )
// s.app.Use(tracingMiddleware)
v, err := validation.New("./pkg/openapi/openapi.json")
v, err := validation.New()
if err != nil {
return nil, err
}
Expand Down
2 changes: 1 addition & 1 deletion apps/agent/pkg/api/testutil/harness.go
Original file line number Diff line number Diff line change
Expand Up @@ -81,7 +81,7 @@ func (h *Harness) Register(route *routes.Route) {

func (h *Harness) SetupRoute(constructor func(svc routes.Services) *routes.Route) *routes.Route {

validator, err := validation.New("./pkg/openapi/openapi.json")
validator, err := validation.New()
require.NoError(h.t, err)
route := constructor(routes.Services{
Logger: h.logger,
Expand Down
12 changes: 4 additions & 8 deletions apps/agent/pkg/api/validation/validator.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,14 +6,13 @@ import (
"fmt"
"io"
"net/http"
"os"

"github.com/Southclaws/fault"
"github.com/Southclaws/fault/fmsg"
"github.com/pb33f/libopenapi"
validator "github.com/pb33f/libopenapi-validator"
"github.com/unkeyed/unkey/apps/agent/gen/openapi"
"github.com/unkeyed/unkey/apps/agent/pkg/api/ctxutil"
"github.com/unkeyed/unkey/apps/agent/pkg/openapi"
"github.com/unkeyed/unkey/apps/agent/pkg/util"
)

Expand All @@ -25,12 +24,9 @@ type Validator struct {
validator validator.Validator
}

func New(specPath string) (*Validator, error) {
b, err := os.ReadFile(specPath)
if err != nil {
return nil, fault.Wrap(err, fmsg.With("failed to read spec file"))
}
document, err := libopenapi.NewDocument(b)
func New() (*Validator, error) {

document, err := libopenapi.NewDocument(openapi.Spec)
if err != nil {
return nil, fault.Wrap(err, fmsg.With("failed to create OpenAPI document"))
}
Expand Down
2 changes: 1 addition & 1 deletion apps/agent/pkg/openapi/config.yaml
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
package: openapi
output: ./gen/openapi/gen.go
output: ./pkg/openapi/gen.go
generate:
models: true
output-options:
Expand Down
File renamed without changes.
11 changes: 11 additions & 0 deletions apps/agent/pkg/openapi/spec.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,11 @@
package openapi

import (
_ "embed"
)

// Spec is the OpenAPI specification for the service
// It's loaded from our openapi file and embedded into the binary
//
//go:embed openapi.json
var Spec []byte
2 changes: 1 addition & 1 deletion apps/agent/pkg/tinybird/tinybird.go
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ import (
"net/http"
"strings"

"github.com/unkeyed/unkey/apps/agent/gen/openapi"
"github.com/unkeyed/unkey/apps/agent/pkg/openapi"
)

type Client struct {
Expand Down
13 changes: 4 additions & 9 deletions apps/agent/schema.json
Original file line number Diff line number Diff line change
Expand Up @@ -13,19 +13,13 @@
"clickhouse": {
"type": "object",
"properties": {
"addr": {
"url": {
"type": "string",
"minLength": 1
},
"password": {
"type": "string"
},
"username": {
"type": "string"
}
},
"additionalProperties": false,
"required": ["addr", "username", "password"]
"required": ["url"]
},
"cluster": {
"type": "object",
Expand Down Expand Up @@ -282,7 +276,8 @@
"required": ["s3Bucket", "s3Url", "s3AccessKeyId", "s3AccessKeySecret", "masterKeys"]
}
},
"additionalProperties": false
"additionalProperties": false,
"required": ["vault"]
},
"tracing": {
"type": "object",
Expand Down
2 changes: 1 addition & 1 deletion apps/agent/services/eventrouter/service.go
Original file line number Diff line number Diff line change
Expand Up @@ -8,13 +8,13 @@ import (
"net/http"
"time"

"github.com/unkeyed/unkey/apps/agent/gen/openapi"
"github.com/unkeyed/unkey/apps/agent/pkg/auth"
"github.com/unkeyed/unkey/apps/agent/pkg/batch"
"github.com/unkeyed/unkey/apps/agent/pkg/clickhouse"
"github.com/unkeyed/unkey/apps/agent/pkg/clickhouse/schema"
"github.com/unkeyed/unkey/apps/agent/pkg/logging"
"github.com/unkeyed/unkey/apps/agent/pkg/metrics"
"github.com/unkeyed/unkey/apps/agent/pkg/openapi"
"github.com/unkeyed/unkey/apps/agent/pkg/prometheus"
"github.com/unkeyed/unkey/apps/agent/pkg/tinybird"
"github.com/unkeyed/unkey/apps/agent/pkg/tracing"
Expand Down
Loading
Loading