From 1879fee9ae2d04b0461ff5657abe01ced4e6e925 Mon Sep 17 00:00:00 2001 From: Lars de Bruijn <9264036+ldebruijn@users.noreply.github.com> Date: Fri, 22 Dec 2023 01:49:14 +0100 Subject: [PATCH] =?UTF-8?q?feat(max-tokens):=20Implement=20max-tokens=20ru?= =?UTF-8?q?le,=20add=20documentation=20&=20test=E2=80=A6=20(#12)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Implement max-tokens rule, add documentation & test… cases, update README.md Co-authored-by: ldebruijn --- README.md | 2 +- cmd/main.go | 22 +++++-- cmd/main_test.go | 45 +++++++++++++ docs/configuration.md | 8 +++ docs/max_tokens.md | 37 +++++++++++ internal/app/config/config.go | 2 + internal/business/tokens/tokens.go | 66 +++++++++++++++++++ internal/business/tokens/tokens_test.go | 85 +++++++++++++++++++++++++ 8 files changed, 262 insertions(+), 5 deletions(-) create mode 100644 docs/max_tokens.md create mode 100644 internal/business/tokens/tokens.go create mode 100644 internal/business/tokens/tokens_test.go diff --git a/README.md b/README.md index 79bcf33..d345cd0 100644 --- a/README.md +++ b/README.md @@ -12,9 +12,9 @@ It is dead-simple yet highly customizable security sidecar compatible with any H * [Persisted Operations](docs/persisted_operations.md) * [Block Field Suggestions](docs/block_field_suggestions.md) * [Max Aliases](docs/max_aliases.md) +* [Max Tokens](docs/max_tokens.md) * _Max Depth (coming soon)_ * _Max Directives (coming soon)_ -* _Max Tokens (coming soon)_ * _Cost Limit (coming soon)_ Curious why you need these features? Check out this [Excellent talk on GraphQL security](https://www.youtube.com/watch?v=hyB2UKsEkqA&list=PLP1igyLx8foE9SlDLI1Vtlshcon5r1jMJ) on YouTube. diff --git a/cmd/main.go b/cmd/main.go index b8ff3a6..a105946 100644 --- a/cmd/main.go +++ b/cmd/main.go @@ -16,9 +16,11 @@ import ( "github.com/ldebruijn/go-graphql-armor/internal/business/proxy" "github.com/ldebruijn/go-graphql-armor/internal/business/readiness" "github.com/ldebruijn/go-graphql-armor/internal/business/schema" + "github.com/ldebruijn/go-graphql-armor/internal/business/tokens" "github.com/prometheus/client_golang/prometheus" "github.com/prometheus/client_golang/prometheus/promhttp" "github.com/vektah/gqlparser/v2/ast" + "github.com/vektah/gqlparser/v2/gqlerror" "github.com/vektah/gqlparser/v2/parser" "github.com/vektah/gqlparser/v2/validator" log2 "log" @@ -164,7 +166,8 @@ func middleware(log *slog.Logger, cfg *config.Config, po *persisted_operations.P httpInstrumentation := HttpInstrumentation() aliases.NewMaxAliasesRule(cfg.MaxAliases) - vr := ValidationRules(schema) + tks := tokens.MaxTokens(cfg.Token) + vr := ValidationRules(schema, tks) fn := func(next http.Handler) http.Handler { return rec(httpInstrumentation(po.Execute(vr(next)))) @@ -187,7 +190,7 @@ func HttpInstrumentation() func(next http.Handler) http.Handler { } } -func ValidationRules(schema *schema.Provider) func(next http.Handler) http.Handler { +func ValidationRules(schema *schema.Provider, tks *tokens.MaxTokensRule) func(next http.Handler) http.Handler { return func(next http.Handler) http.Handler { fn := func(w http.ResponseWriter, r *http.Request) { payload, err := gql.ParseRequestPayload(r) @@ -196,10 +199,21 @@ func ValidationRules(schema *schema.Provider) func(next http.Handler) http.Handl return } - var query, _ = parser.ParseQuery(&ast.Source{ + operationSource := &ast.Source{ Name: payload.OperationName, Input: payload.Query, - }) + } + err = tks.Validate(operationSource) + if err != nil { + response := map[string]interface{}{ + "data": nil, + "errors": gqlerror.List{gqlerror.Wrap(err)}, + } + _ = json.NewEncoder(w).Encode(response) + return + } + + var query, _ = parser.ParseQuery(operationSource) errs := validator.Validate(schema.Get(), query) diff --git a/cmd/main_test.go b/cmd/main_test.go index 1aeae52..3204158 100644 --- a/cmd/main_test.go +++ b/cmd/main_test.go @@ -402,6 +402,51 @@ type Product { assert.NotContains(t, string(actual), "\"errors\":") }, }, + { + name: "throws error when exceeding max tokens", + args: args{ + request: func() *http.Request { + body := map[string]interface{}{ + "query": "query Foo($id: ID!) { product(id: $id) { id name } }", + } + + bts, _ := json.Marshal(body) + r := httptest.NewRequest("POST", "/graphql", bytes.NewBuffer(bts)) + return r + }(), + schema: ` +extend type Query { + product(id: ID!): Product +} + +type Product { + id: ID! + name: String +} +`, + cfgOverrides: func(cfg *config.Config) *config.Config { + cfg.Token.Enabled = true + cfg.Token.Max = 1 + return cfg + }, + mockResponse: map[string]interface{}{ + "data": map[string]interface{}{ + "product": map[string]interface{}{ + "id": "1", + "name": "name", + }, + }, + }, + mockStatusCode: http.StatusOK, + }, + want: func(t *testing.T, response *http.Response) { + assert.Equal(t, http.StatusOK, response.StatusCode) + actual, err := io.ReadAll(response.Body) + assert.NoError(t, err) + _ = actual + assert.Contains(t, string(actual), "operation has exceeded maximum tokens. found [22], max [1]") + }, + }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { diff --git a/docs/configuration.md b/docs/configuration.md index 31e3cd1..61e8b04 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -65,6 +65,14 @@ max_aliases: block_field_suggestions: enabled: "true" mask: [redacted] + +max_tokens: + # Enable the feature + enable: "true" + # The maximum number of allowed tokens within a single request. + max: 10000 + # Reject the request when the rule fails. Disable this to allow the request regardless of token count. + reject_on_failure: "true" ``` For a more in-depth view of each option visit the accompanying documentation page. diff --git a/docs/max_tokens.md b/docs/max_tokens.md new file mode 100644 index 0000000..f536feb --- /dev/null +++ b/docs/max_tokens.md @@ -0,0 +1,37 @@ +# Max Tokens + +Restricting the maximum number of tokens in an operation helps prevent excessively large operations reaching your landscape. +This can be useful to prevent DDoS attacks, Heap Overflows or Server overload. + + + +## Configuration + +You can configure `go-graphql-armor` to limit the maximum number of tokens allowed on an operation. + +```yaml +max_tokens: + # Enable the feature + enable: "true" + # The maximum number of allowed tokens within a single request. + max: 10000 + # Reject the request when the rule fails. Disable this to allow the request regardless of token count. + reject_on_failure: "true" +``` + +## Metrics + +This rule produces metrics to help you gain insights into the behavior of the rule. + +``` +go_graphql_armor_max_tokens_results{result} +``` + + +| `result` | Description | +|---------|--------------------------------------------------------------------------------------------------------------| +| `allowed` | The rule condition succeeded | +| `rejected` | The rule condition failed and the request was rejected | +| `failed` | The rule condition failed but the request was not rejected. This happens when `reject_on_failure` is `false` | + +No metrics are produced when the rule is disabled. \ No newline at end of file diff --git a/internal/app/config/config.go b/internal/app/config/config.go index 9c43167..626d4ee 100644 --- a/internal/app/config/config.go +++ b/internal/app/config/config.go @@ -10,6 +10,7 @@ import ( "github.com/ldebruijn/go-graphql-armor/internal/business/persisted_operations" "github.com/ldebruijn/go-graphql-armor/internal/business/proxy" "github.com/ldebruijn/go-graphql-armor/internal/business/schema" + "github.com/ldebruijn/go-graphql-armor/internal/business/tokens" "os" "time" ) @@ -26,6 +27,7 @@ type Config struct { //DebugHost string `conf:"default:0.0.0.0:4000"` } Schema schema.Config `yaml:"schema"` + Token tokens.Config `yaml:"token"` Target proxy.Config `yaml:"target"` PersistedOperations persisted_operations.Config `yaml:"persisted_operations"` BlockFieldSuggestions block_field_suggestions.Config `yaml:"block_field_suggestions"` diff --git a/internal/business/tokens/tokens.go b/internal/business/tokens/tokens.go new file mode 100644 index 0000000..edd6084 --- /dev/null +++ b/internal/business/tokens/tokens.go @@ -0,0 +1,66 @@ +package tokens + +import ( + "fmt" + "github.com/prometheus/client_golang/prometheus" + "github.com/vektah/gqlparser/v2/ast" + "github.com/vektah/gqlparser/v2/lexer" +) + +var resultCounter = prometheus.NewCounterVec(prometheus.CounterOpts{ + Namespace: "go_graphql_armor", + Subsystem: "max_tokens", + Name: "results", + Help: "The results of the max tokens rule", +}, + []string{"result"}, +) + +type Config struct { + Enabled bool `conf:"default:true" yaml:"enabled"` + Max int `conf:"default:10000" yaml:"max"` + RejectOnFailure bool `conf:"default:true" yaml:"reject-on-failure"` +} + +type MaxTokensRule struct { + cfg Config +} + +func MaxTokens(cfg Config) *MaxTokensRule { + return &MaxTokensRule{ + cfg: cfg, + } +} + +func (t *MaxTokensRule) Validate(source *ast.Source) error { + if !t.cfg.Enabled { + return nil + } + + lex := lexer.New(source) + count := 0 + + for { + tok, err := lex.ReadToken() + + if err != nil { + return err + } + + if tok.Kind == lexer.EOF { + break + } + + count++ + } + + if count > t.cfg.Max { + if t.cfg.RejectOnFailure { + resultCounter.WithLabelValues("rejected").Inc() + return fmt.Errorf("operation has exceeded maximum tokens. found [%d], max [%d]", count, t.cfg.Max) + } + resultCounter.WithLabelValues("failed").Inc() + } + resultCounter.WithLabelValues("allowed").Inc() + return nil +} diff --git a/internal/business/tokens/tokens_test.go b/internal/business/tokens/tokens_test.go new file mode 100644 index 0000000..5de0a3f --- /dev/null +++ b/internal/business/tokens/tokens_test.go @@ -0,0 +1,85 @@ +package tokens + +import ( + "github.com/stretchr/testify/assert" + "github.com/vektah/gqlparser/v2/ast" + "testing" +) + +func TestMaxTokens(t *testing.T) { + type args struct { + cfg Config + operation string + } + tests := []struct { + name string + args args + wantErr bool + }{ + { + name: "rule disabled does nothing", + args: args{ + cfg: Config{ + Enabled: false, + Max: 0, + RejectOnFailure: true, + }, + operation: "query { foo }", + }, + wantErr: false, + }, + { + name: "yields error when tokens exceed max", + args: args{ + cfg: Config{ + Enabled: true, + Max: 1, + RejectOnFailure: true, + }, + operation: "query { foo }", + }, + wantErr: true, + }, + { + name: "yields no error when tokens less than max", + args: args{ + cfg: Config{ + Enabled: true, + Max: 1000000, + RejectOnFailure: true, + }, + operation: "query { foo }", + }, + wantErr: false, + }, + { + name: "yields no error when tokens exceed max but failure on rejections is false", + args: args{ + cfg: Config{ + Enabled: true, + Max: 1, + RejectOnFailure: false, + }, + operation: "query { foo }", + }, + wantErr: false, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + rule := MaxTokens(tt.args.cfg) + + source := &ast.Source{ + Input: tt.args.operation, + } + + err := rule.Validate(source) + + if tt.wantErr { + assert.Error(t, err) + } else { + assert.NoError(t, err) + } + }) + } +}