From 5f76a1f40df81b73fbc63c3622cb436058434112 Mon Sep 17 00:00:00 2001
From: Drew DeVault <sir@cmpwn.com>
Date: Fri, 20 Aug 2021 15:22:33 +0200
Subject: [PATCH] Enforce complexity limit on webhook payload

---
 server/server.go    | 14 +++++++-------
 webhooks/context.go | 11 ++++++++++-
 2 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/server/server.go b/server/server.go
index 7bb382226096ea11d1e903f26beeb7c89b17c394..936afcea3d8cd0c861bd29b1b42362ead9c4b187 100644
--- a/server/server.go
+++ b/server/server.go
@@ -56,6 +56,8 @@ type Server struct {
 	service string
 	queues  []*work.Queue
 	email   *work.Queue
+
+	MaxComplexity int
 }
 
 // Creates a new common server context for a SourceHut GraphQL daemon.
@@ -81,22 +83,19 @@ func (server *Server) WithSchema(
 	schema graphql.ExecutableSchema, scopes []string) *Server {
 	server.Schema = schema
 
-	var (
-		complexity int
-		err        error
-	)
+	var err error
 	if limit, ok := server.conf.Get(
 		server.service+"::api", "max-complexity"); ok {
-		complexity, err = strconv.Atoi(limit)
+		server.MaxComplexity, err = strconv.Atoi(limit)
 		if err != nil {
 			panic(err)
 		}
 	} else {
-		complexity = 250
+		server.MaxComplexity = 250
 	}
 
 	srv := handler.GraphQL(schema,
-		handler.ComplexityLimit(complexity),
+		handler.ComplexityLimit(server.MaxComplexity),
 		handler.RecoverFunc(EmailRecover),
 		handler.UploadMaxSize(1073741824)) // 1 GiB (TODO: configurable?)
 
@@ -234,6 +233,7 @@ func (server *Server) WithQueues(queues ...*work.Queue) *Server {
 	ctx = database.Context(ctx, server.db)
 	ctx = redis.Context(ctx, server.redis)
 	ctx = email.Context(ctx, server.email)
+	ctx = context.WithValue(ctx, serverCtxKey, server)
 
 	server.queues = append(server.queues, queues...)
 	for _, queue := range queues {
diff --git a/webhooks/context.go b/webhooks/context.go
index 029b683ce983470533f8f5f50046b211852fce23..082f7d2da8efca92934bd638939e7f1c339c0592 100644
--- a/webhooks/context.go
+++ b/webhooks/context.go
@@ -7,6 +7,7 @@ import (
 	"errors"
 	"fmt"
 
+	"github.com/99designs/gqlgen/complexity"
 	"github.com/99designs/gqlgen/graphql"
 	"github.com/99designs/gqlgen/graphql/executor"
 	"github.com/google/uuid"
@@ -64,7 +65,6 @@ func (webhook *WebhookContext) Exec(ctx context.Context,
 		return nil, err
 	}
 
-	// TODO: Set complexity limit
 	exec := executor.New(schema)
 	params := graphql.RawParams{
 		Query: sub.Query,
@@ -80,6 +80,15 @@ func (webhook *WebhookContext) Exec(ctx context.Context,
 	}
 	rc.RecoverFunc = server.EmailRecover
 
+	op := rc.Doc.Operations.ForName(rc.OperationName)
+	complexity := complexity.Calculate(schema, op, rc.Variables)
+	srv := server.ForContext(ctx)
+	if complexity > srv.MaxComplexity {
+		// TODO: This doesn't bubble up to the user well
+		return nil, fmt.Errorf("operation has complexity %d, which exceeds the maximum of %d",
+			complexity, srv.MaxComplexity)
+	}
+
 	var resp graphql.ResponseHandler
 	ctx = graphql.WithOperationContext(ctx, rc)
 	resp, ctx = exec.DispatchOperation(ctx, rc)