diff --git a/cache/memory/memory.go b/cache/memory/memory.go index 43b1a8b..3578312 100644 --- a/cache/memory/memory.go +++ b/cache/memory/memory.go @@ -158,7 +158,12 @@ func (c *Cache) Set(key string, value []byte, ttl time.Duration) { entry.MemorySize = int64(len(key) + len(entry.Value) + approxEntryOverhead) } - // Check if this is a new entry or an update + // Check if this is a new entry or an update. The exists-check, counter + // update and store must be atomic together: without the lock, two + // concurrent Sets of the same brand-new key both observe it as absent and + // both increment entryCount/memoryUsage, leaving the counters inflated + // for a single stored entry. + c.Lock() oldEntry, exists := c.entries.Load(key) if exists { // Update memory usage: subtract old entry size, add new entry size @@ -172,6 +177,7 @@ func (c *Cache) Set(key string, value []byte, ttl time.Duration) { // Add new entry's memory size to total atomic.AddInt64(&c.memoryUsage, entry.MemorySize) c.entries.Store(key, entry) + c.Unlock() } func (c *Cache) Get(key string) ([]byte, bool) { diff --git a/cache/memory/memory_additional_test.go b/cache/memory/memory_additional_test.go index e5085e6..495fc02 100644 --- a/cache/memory/memory_additional_test.go +++ b/cache/memory/memory_additional_test.go @@ -118,3 +118,31 @@ func TestMemoryCacheConcurrentGetOnExpiredEntry_CountersConsistent(t *testing.T) t.Fatalf("memoryUsage = %d after concurrent expired gets, want 0", got) } } + +func TestMemoryCacheConcurrentSetSameNewKey_CounterNotDoubled(t *testing.T) { + // Regression: two (or more) concurrent Sets of the same brand-new key must + // not both take the "new entry" branch. Previously the exists-check and + // counter increment were separate atomic ops, so concurrent Sets of a new + // key each incremented entryCount, inflating it for a single stored entry. + cache := NewWithSize(DefaultTestExpiration, DefaultMaxMemorySize, DefaultMaxCacheSize) + + const workers = 32 + var wg sync.WaitGroup + for range workers { + wg.Add(1) + go func() { + defer wg.Done() + cache.Set("same-key", []byte("value"), DefaultTestExpiration) + }() + } + wg.Wait() + + if got := cache.CountQueries(); got != 1 { + t.Fatalf("entryCount = %d after %d concurrent Sets of one new key, want 1", got, workers) + } + + // Sanity: the entry is actually stored and retrievable. + if _, ok := cache.Get("same-key"); !ok { + t.Fatal("key should be retrievable after Set") + } +} diff --git a/graphql.go b/graphql.go index b64090a..a1fa964 100644 --- a/graphql.go +++ b/graphql.go @@ -322,6 +322,17 @@ func parseGraphQLQuery(c fiber.Ctx) *parseGraphQLQueryResult { res.shouldIgnore = false res.operationName = "undefined" + // Index named fragments so FragmentSpread nodes can be resolved during the + // introspection check. A named spread (e.g. `...F`) only carries a name; + // its content lives in the matching top-level FragmentDefinition, which the + // OperationDefinition loop below otherwise skips. + fragments := make(map[string]*ast.FragmentDefinition) + for _, d := range p.Definitions { + if fd, ok := d.(*ast.FragmentDefinition); ok && fd.Name != nil { + fragments[fd.Name.Value] = fd + } + } + // Single pass over definitions: gather operation type, mutation flag, // operation name, and process directives / introspection checks together. // Mutations take priority for operationType regardless of order. @@ -368,7 +379,7 @@ func parseGraphQLQuery(c fiber.Ctx) *parseGraphQLQueryResult { processDirectives(oper, res) // Check for introspection queries if they're blocked - if cfg.Security.BlockIntrospection && checkSelections(c, oper.GetSelectionSet().Selections) { + if cfg.Security.BlockIntrospection && checkSelections(c, oper.GetSelectionSet().Selections, fragments) { _ = c.Status(403).SendString("Introspection queries are not allowed") res.shouldBlock = true return res @@ -417,8 +428,10 @@ func processDirectives(oper *ast.OperationDefinition, res *parseGraphQLQueryResu } } -// checkSelections recursively checks if any selection is an introspection query that should be blocked -func checkSelections(c fiber.Ctx, selections []ast.Selection) bool { +// checkSelections recursively checks if any selection is an introspection query that should be blocked. +// fragments indexes top-level named FragmentDefinition nodes so that a FragmentSpread (which only +// carries a name) can be resolved to its actual selection set. +func checkSelections(c fiber.Ctx, selections []ast.Selection, fragments map[string]*ast.FragmentDefinition) bool { if len(selections) == 0 { return false } @@ -450,7 +463,7 @@ func checkSelections(c fiber.Ctx, selections []ast.Selection) bool { // Check nested selections if present if sel.SelectionSet != nil && len(sel.GetSelectionSet().Selections) > 0 { - if checkSelections(c, sel.GetSelectionSet().Selections) { + if checkSelections(c, sel.GetSelectionSet().Selections, fragments) { return true } } @@ -458,10 +471,22 @@ func checkSelections(c fiber.Ctx, selections []ast.Selection) bool { case *ast.InlineFragment: // Check nested selections in fragments if sel.SelectionSet != nil && len(sel.GetSelectionSet().Selections) > 0 { - if checkSelections(c, sel.GetSelectionSet().Selections) { + if checkSelections(c, sel.GetSelectionSet().Selections, fragments) { return true } } + + case *ast.FragmentSpread: + // A named fragment spread carries no selection set itself; resolve it + // to its top-level FragmentDefinition so introspection inside a named + // fragment (e.g. `...F` where F selects __schema) cannot bypass the block. + if sel.Name != nil { + if frag, ok := fragments[sel.Name.Value]; ok && frag.SelectionSet != nil { + if checkSelections(c, frag.SelectionSet.Selections, fragments) { + return true + } + } + } } } @@ -500,11 +525,18 @@ func checkIfContainsIntrospection(c fiber.Ctx, query string) bool { } if p != nil { + // Index named fragments so FragmentSpread nodes resolve correctly. + fragments := make(map[string]*ast.FragmentDefinition) + for _, def := range p.Definitions { + if fd, ok := def.(*ast.FragmentDefinition); ok && fd.Name != nil { + fragments[fd.Name.Value] = fd + } + } // It's a complete query, check all selections for _, def := range p.Definitions { if op, ok := def.(*ast.OperationDefinition); ok { if op.SelectionSet != nil { - blocked = checkSelections(c, op.GetSelectionSet().Selections) + blocked = checkSelections(c, op.GetSelectionSet().Selections, fragments) break } } diff --git a/graphql_test.go b/graphql_test.go index 9336605..5bb25f1 100644 --- a/graphql_test.go +++ b/graphql_test.go @@ -476,6 +476,30 @@ func (suite *Tests) Test_DeepIntrospectionQueries() { allowed: []string{}, expected: true, }, + { + name: "named fragment spreading introspection is blocked", + query: "query Q { ...IntroFields } fragment IntroFields on Query { __schema { queryType { name } } }", + allowed: []string{}, + expected: true, + }, + { + name: "named fragment spreading schema with allowlist is blocked", + query: "query Q { ...IntroFields } fragment IntroFields on Query { __schema { queryType { name } } }", + allowed: []string{"__typename"}, + expected: true, + }, + { + name: "named fragment without introspection is allowed", + query: "query Q { ...Fields } fragment Fields on User { id name }", + allowed: []string{}, + expected: false, + }, + { + name: "nested named fragment chain with introspection is blocked", + query: "query Q { ...A } fragment A on Query { ...B } fragment B on Query { __schema { queryType { name } } }", + allowed: []string{}, + expected: true, + }, } for _, tt := range tests { @@ -592,9 +616,15 @@ func TestIntrospectionQueryHandling(t *testing.T) { // Check selections var blocked bool + fragments := make(map[string]*ast.FragmentDefinition) + for _, def := range p.Definitions { + if fd, ok := def.(*ast.FragmentDefinition); ok && fd.Name != nil { + fragments[fd.Name.Value] = fd + } + } for _, def := range p.Definitions { if op, ok := def.(*ast.OperationDefinition); ok { - blocked = checkSelections(ctx, op.GetSelectionSet().Selections) + blocked = checkSelections(ctx, op.GetSelectionSet().Selections, fragments) break } } diff --git a/proxy.go b/proxy.go index 614b0f2..e916ac5 100644 --- a/proxy.go +++ b/proxy.go @@ -370,12 +370,15 @@ func setupTracing(c fiber.Ctx) context.Context { // performProxyRequest executes the proxy request with retries, circuit breaker, and request coalescing func performProxyRequest(c fiber.Ctx, proxyURL string) error { - // Extract user context for cache key (needed for coalescing and circuit breaker fallback) - userID, userRole := extractUserInfo(c) - - // Calculate cache key - includes user context for security - // This key is used for both request coalescing and cache fallback - cacheKey := libpack_cache.CalculateHash(c, userID, userRole) + // Reuse the cache key already computed by handleCaching (stored on the + // request context) to avoid a second MD5 of the full request body per + // proxied request. Fall back to computing it when this function is + // reached without going through handleCaching (direct callers). + cacheKey, _ := c.Locals("query_cache_hash").(string) + if cacheKey == "" { + userID, userRole := extractUserInfo(c) + cacheKey = libpack_cache.CalculateHash(c, userID, userRole) + } // Check if request coalescing is enabled rc := GetRequestCoalescer() diff --git a/server.go b/server.go index 2ed6c80..3dcf895 100644 --- a/server.go +++ b/server.go @@ -356,6 +356,10 @@ func handleCaching(c fiber.Ctx, parsedResult *parseGraphQLQueryResult, userID, u // Calculate query hash for cache key - now includes user context for security calculatedQueryHash := libpack_cache.CalculateHash(c, userID, userRole) + // Share the precomputed hash with the proxy path (request coalescing) so it + // is not hashed a second time per request (same c, userID, userRole). + c.Locals("query_cache_hash", calculatedQueryHash) + // Set cache time from header or default if parsedResult.cacheTime == 0 { if cacheQuery := c.Get("X-Cache-Graphql-Query"); cacheQuery != "" {