Fix blocking the introspection + add unit tests.

graphql.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"flag"
 	"strconv"
 	"strings"
 
@@ -10,7 +11,7 @@ import (
 	libpack_monitoring "github.com/lukaszraczylo/graphql-monitoring-proxy/monitoring"
 )
 
-var retrospection_queries = []string{
+var introspection_queries = []string{
 	"__schema",
 	"__type",
 	"__typename",
@@ -34,13 +35,29 @@ var retrospection_queries = []string{
 }
 
 // Saving the introspection queries as a map O(1) operation instead of O(n) for a slice.
-var retrospectionQuerySet = func() map[string]struct{} {
-	rsqs := make(map[string]struct{}, len(retrospection_queries))
-	for _, query := range retrospection_queries {
-		rsqs[strings.ToLower(query)] = struct{}{}
-	}
-	return rsqs
-}()
+
+var introspectionQuerySet = map[string]struct{}{}
+var introspectionAllowedQueries = map[string]struct{}{}
+
+func prepareQueriesAndExemptions() {
+	introspectionQuerySet = map[string]struct{}{}
+	introspectionQuerySet = func() map[string]struct{} {
+		rsqs := make(map[string]struct{}, len(introspection_queries))
+		for _, query := range introspection_queries {
+			rsqs[strings.ToLower(query)] = struct{}{}
+		}
+		return rsqs
+	}()
+
+	introspectionAllowedQueries = map[string]struct{}{}
+	introspectionAllowedQueries = func() map[string]struct{} {
+		rsqs := make(map[string]struct{}, len(cfg.Security.IntrospectionAllowed))
+		for _, query := range cfg.Security.IntrospectionAllowed {
+			rsqs[strings.ToLower(query)] = struct{}{}
+		}
+		return rsqs
+	}()
+}
 
 func parseGraphQLQuery(c *fiber.Ctx) (operationType, operationName string, cacheRequest bool, cache_time int, should_block bool, should_ignore bool) {
 	should_ignore = true
@@ -48,21 +65,27 @@ func parseGraphQLQuery(c *fiber.Ctx) (operationType, operationName string, cache
 	err := json.Unmarshal(c.Body(), &m)
 	if err != nil {
 		cfg.Logger.Debug("Can't unmarshal the request", map[string]interface{}{"error": err.Error(), "body": string(c.Body())})
-		cfg.Monitoring.Increment(libpack_monitoring.MetricsSkipped, nil)
+		if flag.Lookup("test.v") == nil {
+			cfg.Monitoring.Increment(libpack_monitoring.MetricsSkipped, nil)
+		}
 		return
 	}
 	// get the query
 	query, ok := m["query"].(string)
 	if !ok {
 		cfg.Logger.Error("Can't find the query", map[string]interface{}{"query": query, "m_val": m})
-		cfg.Monitoring.Increment(libpack_monitoring.MetricsSkipped, nil)
+		if flag.Lookup("test.v") == nil {
+			cfg.Monitoring.Increment(libpack_monitoring.MetricsSkipped, nil)
+		}
 		return
 	}
 
 	p, err := parser.Parse(parser.ParseParams{Source: query})
 	if err != nil {
 		cfg.Logger.Error("Can't parse the query", map[string]interface{}{"query": query, "m_val": m})
-		cfg.Monitoring.Increment(libpack_monitoring.MetricsFailed, nil)
+		if flag.Lookup("test.v") == nil {
+			cfg.Monitoring.Increment(libpack_monitoring.MetricsFailed, nil)
+		}
 		return
 	}
 
@@ -71,19 +94,21 @@ func parseGraphQLQuery(c *fiber.Ctx) (operationType, operationName string, cache
 	for _, d := range p.Definitions {
 		if oper, ok := d.(*ast.OperationDefinition); ok {
 			operationType = oper.Operation
+
+			if oper.Name != nil {
+				operationName = oper.Name.Value
+			}
+
 			if strings.ToLower(operationType) == "mutation" && cfg.Server.ReadOnlyMode {
 				cfg.Logger.Warning("Mutation blocked", m)
-				cfg.Monitoring.Increment(libpack_monitoring.MetricsSkipped, nil)
+				if flag.Lookup("test.v") == nil {
+					cfg.Monitoring.Increment(libpack_monitoring.MetricsSkipped, nil)
+				}
 				c.Status(403).SendString("The server is in read-only mode")
 				should_block = true
 				return
 			}
 
-			if oper.Name != nil {
-				operationName = oper.Name.Value
-			} else {
-				operationName = "undefined"
-			}
 			for _, dir := range oper.Directives {
 				if dir.Name.Value == "cached" {
 					cacheRequest = true
@@ -91,38 +116,67 @@ func parseGraphQLQuery(c *fiber.Ctx) (operationType, operationName string, cache
 						if arg.Name.Value == "ttl" {
 							cache_time, err = strconv.Atoi(arg.Value.GetValue().(string))
 							if err != nil {
-								cfg.Logger.Error("Can't parse the ttl", map[string]interface{}{"ttl": arg.Value.GetValue().(string)})
-								cfg.Monitoring.Increment(libpack_monitoring.MetricsFailed, nil)
+								cfg.Logger.Error("Can't parse the ttl, using global", map[string]interface{}{"bad_ttl": arg.Value.GetValue().(string)})
+								if flag.Lookup("test.v") == nil {
+									cfg.Monitoring.Increment(libpack_monitoring.MetricsFailed, nil)
+								}
 								return
 							}
 						}
 					}
 				}
 			}
-			if cfg.Security.BlockIntrospection {
 
-				for _, s := range oper.SelectionSet.Selections {
-					for _, s2 := range s.GetSelectionSet().Selections {
-						if _, exists := retrospectionQuerySet[strings.ToLower(s2.(*ast.Field).Name.Value)]; exists {
-							if len(cfg.Security.IntrospectionAllowed) > 0 {
-								for _, introspectionQueryAllowed := range cfg.Security.IntrospectionAllowed {
-									if strings.EqualFold(strings.ToLower(introspectionQueryAllowed), strings.ToLower(s2.(*ast.Field).Name.Value)) {
-										cfg.Logger.Debug("Introspection query allowed, passing through", m)
-										return
-									}
-								}
-							}
-							cfg.Logger.Warning("Introspection query blocked", m)
-							cfg.Monitoring.Increment(libpack_monitoring.MetricsSkipped, nil)
-							c.Status(403).SendString("Introspection queries are not allowed")
-							should_block = true
-							return
-						}
-					}
+			if cfg.Security.BlockIntrospection {
+				should_block = checkSelections(c, oper.GetSelectionSet().Selections)
+				if should_block {
+					return
 				}
 			}
 		}
 	}
+	return
+}
 
+func checkSelections(c *fiber.Ctx, selections []ast.Selection) bool {
+	for _, s := range selections {
+		field, ok := s.(*ast.Field)
+		if !ok {
+			continue // or handle the case where the type assertion fails
+		}
+		shouldBlock := checkIfContainsIntrospection(c, field.Name.Value)
+		if shouldBlock {
+			return true
+		}
+		if field.SelectionSet != nil {
+			if checkSelections(c, field.GetSelectionSet().Selections) {
+				return true
+			}
+		}
+	}
+	return false
+}
+
+func checkIfContainsIntrospection(c *fiber.Ctx, whatever string) (should_block bool) {
+	whateverLower := strings.ToLower(whatever)
+	got_exemption := false
+	if _, exists := introspectionQuerySet[whateverLower]; exists {
+		if len(cfg.Security.IntrospectionAllowed) > 0 {
+			if _, allowed_exists := introspectionAllowedQueries[whateverLower]; allowed_exists {
+				cfg.Logger.Debug("Introspection query allowed, passing through", map[string]interface{}{"query": whatever})
+				got_exemption = true
+				should_block = false
+			}
+		}
+		if !got_exemption {
+			should_block = true
+		}
+	}
+	if should_block {
+		if flag.Lookup("test.v") == nil {
+			cfg.Monitoring.Increment(libpack_monitoring.MetricsSkipped, nil)
+		}
+		c.Status(403).SendString("Introspection queries are not allowed")
+	}
 	return
 }