feat(max-tokens): Implement max-tokens rule, add documentation & test… (

#12) Implement max-tokens rule, add documentation & test… cases, update README.md Co-authored-by: ldebruijn <[email protected]>
ldebruijn · Dec 22, 2023 · 1879fee · 1879fee
1 parent 1455da7
commit 1879fee
Show file tree

Hide file tree

Showing 8 changed files with 262 additions and 5 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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.
+
+<!-- TOC -->
+
+## 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.
diff --git 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
@@ -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
@@ -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)
+			}
+		})
+	}
+}