General code optimisations. (#16)

* General code optimisations.

server.go

@@ -3,6 +3,7 @@ package main
 import (
 	"fmt"
 	"strconv"
+	"sync"
 	"time"
 
 	"github.com/goccy/go-json"
@@ -16,13 +17,24 @@ import (
 	libpack_monitoring "github.com/lukaszraczylo/graphql-monitoring-proxy/monitoring"
 )
 
-// StartHTTPProxy starts the HTTP and points it to the GraphQL server.
+const (
+	healthCheckQueryStr = `{ __typename }`
+)
+
+var (
+	ctxPool = sync.Pool{
+		New: func() interface{} {
+			return new(fiber.Ctx)
+		},
+	}
+)
+
 func StartHTTPProxy() {
 	cfg.Logger.Debug(&libpack_logger.LogMessage{
 		Message: "Starting the HTTP proxy",
-		Pairs:   nil,
 	})
-	server := fiber.New(fiber.Config{
+
+	serverConfig := fiber.Config{
 		DisableStartupMessage: true,
 		AppName:               fmt.Sprintf("GraphQL Monitoring Proxy - %s v%s", libpack_config.PKG_NAME, libpack_config.PKG_VERSION),
 		IdleTimeout:           time.Duration(cfg.Client.ClientTimeout) * time.Second * 2,
@@ -30,13 +42,14 @@ func StartHTTPProxy() {
 		WriteTimeout:          time.Duration(cfg.Client.ClientTimeout) * time.Second * 2,
 		JSONEncoder:           json.Marshal,
 		JSONDecoder:           json.Unmarshal,
-	})
+	}
+
+	server := fiber.New(serverConfig)
 
 	server.Use(cors.New(cors.Config{
 		AllowOrigins: "*",
 	}))
 
-	// add middleware to check if the request is a GraphQL query
 	server.Use(AddRequestUUID)
 
 	server.Get("/healthz", healthCheck)
@@ -49,11 +62,11 @@ func StartHTTPProxy() {
 		Message: "GraphQL proxy started",
 		Pairs:   map[string]interface{}{"port": cfg.Server.PortGraphQL},
 	})
-	err := server.Listen(fmt.Sprintf(":%d", cfg.Server.PortGraphQL))
-	if err != nil {
+
+	if err := server.Listen(fmt.Sprintf(":%d", cfg.Server.PortGraphQL)); err != nil {
 		cfg.Logger.Critical(&libpack_logger.LogMessage{
 			Message: "Can't start the service",
-			Pairs:   map[string]interface{}{"port": cfg.Server.PortGraphQL},
+			Pairs:   map[string]interface{}{"port": cfg.Server.PortGraphQL, "error": err.Error()},
 		})
 	}
 }
@@ -71,7 +84,8 @@ func checkAllowedURLs(c *fiber.Ctx) bool {
 	if len(allowedUrls) == 0 {
 		return true
 	}
-	_, ok := allowedUrls[c.Path()]
+	path := c.OriginalURL()
+	_, ok := allowedUrls[path]
 	return ok
 }
 
@@ -81,93 +95,71 @@ func healthCheck(c *fiber.Ctx) error {
 			Message: "Health check enabled",
 			Pairs:   map[string]interface{}{"url": cfg.Server.HealthcheckGraphQL},
 		})
-		query := `{ __typename }`
-		_, err := cfg.Client.GQLClient.Query(query, nil, nil)
+
+		_, err := cfg.Client.GQLClient.Query(healthCheckQueryStr, nil, nil)
 		if err != nil {
 			cfg.Logger.Error(&libpack_logger.LogMessage{
 				Message: "Can't reach the GraphQL server",
 				Pairs:   map[string]interface{}{"error": err.Error()},
 			})
 			cfg.Monitoring.Increment(libpack_monitoring.MetricsFailed, nil)
-			c.Status(500).SendString("Can't reach the GraphQL server with {__typename} query")
-			return err
+			return c.Status(500).SendString("Can't reach the GraphQL server with {__typename} query")
 		}
 	}
+
 	cfg.Logger.Debug(&libpack_logger.LogMessage{
 		Message: "Health check returning OK",
-		Pairs:   nil,
 	})
-	c.Status(200).SendString("Health check OK")
-	return nil
+	return c.Status(200).SendString("Health check OK")
 }
 
 func processGraphQLRequest(c *fiber.Ctx) error {
 	startTime := time.Now()
 
-	// Initialize variables with default values
 	extractedUserID := "-"
 	extractedRoleName := "-"
-	var queryCacheHash string
 
-	authorization := c.Request().Header.Peek("Authorization")
-	if authorization != nil && (len(cfg.Client.JWTUserClaimPath) > 0 || len(cfg.Client.JWTRoleClaimPath) > 0) {
-		extractedUserID, extractedRoleName = extractClaimsFromJWTHeader(string(authorization))
+	if authorization := c.Get("Authorization"); authorization != "" && (len(cfg.Client.JWTUserClaimPath) > 0 || len(cfg.Client.JWTRoleClaimPath) > 0) {
+		extractedUserID, extractedRoleName = extractClaimsFromJWTHeader(authorization)
 	}
 
 	if checkIfUserIsBanned(c, extractedUserID) {
-		c.Status(403).SendString("User is banned")
-		return nil
+		return c.Status(403).SendString("User is banned")
 	}
 
-	if len(cfg.Client.RoleFromHeader) > 0 {
-		extractedRoleName = string(c.Request().Header.Peek(cfg.Client.RoleFromHeader))
-		if extractedRoleName == "" {
-			extractedRoleName = "-"
+	if cfg.Client.RoleFromHeader != "" {
+		if role := c.Get(cfg.Client.RoleFromHeader); role != "" {
+			extractedRoleName = role
 		}
 	}
 
-	// Implementing rate limiting if enabled
 	if cfg.Client.RoleRateLimit {
 		cfg.Logger.Debug(&libpack_logger.LogMessage{
 			Message: "Rate limiting enabled",
 			Pairs:   map[string]interface{}{"user_id": extractedUserID, "role_name": extractedRoleName},
 		})
 		if !rateLimitedRequest(extractedUserID, extractedRoleName) {
-			c.Status(429).SendString("Rate limit exceeded, try again later")
-			return nil
+			return c.Status(429).SendString("Rate limit exceeded, try again later")
 		}
 	}
 
 	parsedResult := parseGraphQLQuery(c)
 	if parsedResult.shouldBlock {
-		c.Status(403).SendString("Request blocked")
-		return nil
+		return c.Status(403).SendString("Request blocked")
 	}
 
 	if parsedResult.shouldIgnore {
 		cfg.Logger.Debug(&libpack_logger.LogMessage{
 			Message: "Request passed as-is - probably not a GraphQL",
-			Pairs:   nil,
 		})
 		return proxyTheRequest(c, parsedResult.activeEndpoint)
 	}
 
 	calculatedQueryHash := libpack_cache.CalculateHash(c)
 
-	if parsedResult.cacheTime > 0 {
-		cfg.Logger.Debug(&libpack_logger.LogMessage{
-			Message: "Cache time set via query",
-			Pairs:   map[string]interface{}{"cacheTime": parsedResult.cacheTime},
-		})
-	} else {
-		// If not set via query, try setting via header
-		cacheQuery := c.Request().Header.Peek("X-Cache-Graphql-Query")
-		if cacheQuery != nil {
-			parsedResult.cacheTime, _ = strconv.Atoi(string(cacheQuery))
-			cfg.Logger.Debug(&libpack_logger.LogMessage{
-				Message: "Cache time set via header",
-				Pairs:   map[string]interface{}{"cacheTime": parsedResult.cacheTime},
-			})
+	if parsedResult.cacheTime == 0 {
+		if cacheQuery := c.Get("X-Cache-Graphql-Query"); cacheQuery != "" {
+			parsedResult.cacheTime, _ = strconv.Atoi(cacheQuery)
 		} else {
 			parsedResult.cacheTime = cfg.Cache.CacheTTL
 		}
@@ -183,82 +175,67 @@ func processGraphQLRequest(c *fiber.Ctx) error {
 		libpack_cache.CacheDelete(calculatedQueryHash)
 	}
 
-	// Handling Cache Logic
 	if parsedResult.cacheRequest || cfg.Cache.CacheEnable || cfg.Cache.CacheRedisEnable {
 		cfg.Logger.Debug(&libpack_logger.LogMessage{
 			Message: "Cache enabled",
 			Pairs:   map[string]interface{}{"via_query": parsedResult.cacheRequest, "via_env": cfg.Cache.CacheEnable},
 		})
-		queryCacheHash = calculatedQueryHash
 
-		if cachedResponse := libpack_cache.CacheLookup(queryCacheHash); cachedResponse != nil {
+		if cachedResponse := libpack_cache.CacheLookup(calculatedQueryHash); cachedResponse != nil {
 			cfg.Monitoring.Increment(libpack_monitoring.MetricsCacheHit, nil)
 			cfg.Logger.Debug(&libpack_logger.LogMessage{
 				Message: "Cache hit",
-				Pairs:   map[string]interface{}{"hash": queryCacheHash, "user_id": extractedUserID, "request_uuid": c.Locals("request_uuid")},
+				Pairs:   map[string]interface{}{"hash": calculatedQueryHash, "user_id": extractedUserID, "request_uuid": c.Locals("request_uuid")},
 			})
-			c.Request().Header.Add("X-Cache-Hit", "true")
-			err := c.Send(cachedResponse)
-			if err != nil {
-				cfg.Logger.Error(&libpack_logger.LogMessage{
-					Message: "Can't send the cached response",
-					Pairs:   map[string]interface{}{"error": err.Error()},
-				})
-				cfg.Monitoring.Increment(libpack_monitoring.MetricsFailed, nil)
-				c.Status(500).SendString("Can't send the cached response - try again later")
-			}
+			c.Set("X-Cache-Hit", "true")
 			wasCached = true
-		} else {
-			cfg.Monitoring.Increment(libpack_monitoring.MetricsCacheMiss, nil)
-			cfg.Logger.Debug(&libpack_logger.LogMessage{
-				Message: "Cache miss",
-				Pairs:   map[string]interface{}{"hash": queryCacheHash, "user_id": extractedUserID, "request_uuid": c.Locals("request_uuid")},
-			})
-			proxyAndCacheTheRequest(c, queryCacheHash, parsedResult.cacheTime, parsedResult.activeEndpoint)
+			return c.Send(cachedResponse)
+		}
+
+		cfg.Monitoring.Increment(libpack_monitoring.MetricsCacheMiss, nil)
+		cfg.Logger.Debug(&libpack_logger.LogMessage{
+			Message: "Cache miss",
+			Pairs:   map[string]interface{}{"hash": calculatedQueryHash, "user_id": extractedUserID, "request_uuid": c.Locals("request_uuid")},
+		})
+		if err := proxyAndCacheTheRequest(c, calculatedQueryHash, parsedResult.cacheTime, parsedResult.activeEndpoint); err != nil {
+			return err
 		}
 	} else {
-		err := proxyTheRequest(c, parsedResult.activeEndpoint)
-		if err != nil {
+		if err := proxyTheRequest(c, parsedResult.activeEndpoint); err != nil {
 			cfg.Logger.Error(&libpack_logger.LogMessage{
 				Message: "Can't proxy the request",
 				Pairs:   map[string]interface{}{"error": err.Error()},
 			})
 			cfg.Monitoring.Increment(libpack_monitoring.MetricsFailed, nil)
-			c.Status(500).SendString("Can't proxy the request - try again later")
-			return nil
+			return c.Status(500).SendString("Can't proxy the request - try again later")
 		}
 	}
 
-	timeTaken := time.Since(startTime)
-
-	// Logging & Monitoring
-	logAndMonitorRequest(c, extractedUserID, parsedResult.operationType, parsedResult.operationName, wasCached, timeTaken, startTime)
+	logAndMonitorRequest(c, extractedUserID, parsedResult.operationType, parsedResult.operationName, wasCached, time.Since(startTime), startTime)
 
 	return nil
 }
 
-// Additional helper function to avoid code repetition
-func proxyAndCacheTheRequest(c *fiber.Ctx, queryCacheHash string, cacheTime int, currentEndpoint string) {
-	err := proxyTheRequest(c, currentEndpoint)
-	if err != nil {
+func proxyAndCacheTheRequest(c *fiber.Ctx, queryCacheHash string, cacheTime int, currentEndpoint string) error {
+	if err := proxyTheRequest(c, currentEndpoint); err != nil {
 		cfg.Logger.Error(&libpack_logger.LogMessage{
 			Message: "Can't proxy the request",
 			Pairs:   map[string]interface{}{"error": err.Error()},
 		})
 		cfg.Monitoring.Increment(libpack_monitoring.MetricsFailed, nil)
-		c.Status(500).SendString("Can't proxy the request - try again later")
-		return
+		return c.Status(500).SendString("Can't proxy the request - try again later")
 	}
+
 	libpack_cache.CacheStoreWithTTL(queryCacheHash, c.Response().Body(), time.Duration(cacheTime)*time.Second)
 	cfg.Monitoring.Increment(libpack_monitoring.MetricsQueriesCached, nil)
-	c.Send(c.Response().Body())
+	return c.Send(c.Response().Body())
 }
 
 func logAndMonitorRequest(c *fiber.Ctx, userID, opType, opName string, wasCached bool, duration time.Duration, startTime time.Time) {
 	labels := map[string]string{
 		"op_type": opType,
 		"op_name": opName,
-		"cached":  fmt.Sprintf("%t", wasCached),
+		"cached":  strconv.FormatBool(wasCached),
 		"user_id": userID,
 	}
 
@@ -267,7 +244,7 @@ func logAndMonitorRequest(c *fiber.Ctx, userID, opType, opName string, wasCached
 			Message: "Request processed",
 			Pairs: map[string]interface{}{
 				"ip":           c.IP(),
-				"fwd-ip":       string(c.Request().Header.Peek("X-Forwarded-For")),
+				"fwd-ip":       c.Get("X-Forwarded-For"),
 				"user_id":      userID,
 				"op_type":      opType,
 				"op_name":      opName,