Merge pull request #30 from NeuralTrust/x-selected-provider

fix the x-selected-provider header
NeuralTrust · Feb 11, 2025 · 6ebc63e · 6ebc63e
2 parents 0fc735e + ce9b021
commit 6ebc63e
Show file tree

Hide file tree

Showing 3 changed files with 159 additions and 161 deletions.
diff --git a/docker-compose.yaml b/docker-compose.yaml
@@ -2,12 +2,10 @@ version: '3'
 
 services:
   redis:
-    image: redis/redis-stack:latest
+    image: redis:7-alpine
     restart: always
     ports:
       - "6379:6379"
-    volumes:
-      - redis_data:/data
   redis-insight:
     image: redis/redisinsight:latest
     restart: always

diff --git a/examples/test_rate_limiter.sh b/examples/test_rate_limiter.sh
@@ -6,174 +6,169 @@ RED='\033[0;31m'
 NC='\033[0m'
 
 # Use environment variables or defaults
-ADMIN_URL=${ADMIN_URL:-"https://admin.neuraltrust.ai/api/v1"}
-PROXY_URL=${PROXY_URL:-"https://gateway.neuraltrust.app"}
-BASE_DOMAIN=${BASE_DOMAIN:-"neuraltrust.app"}
-SUBDOMAIN="gateway"
+ADMIN_URL=${ADMIN_URL:-"http://localhost:8080/api/v1"}
+PROXY_URL=${PROXY_URL:-"http://localhost:8081"}
+BASE_DOMAIN=${BASE_DOMAIN:-"example.com"}
+SUBDOMAIN="gateway-rate-limiter-$(date +%s)"
 
 echo -e "${GREEN}Testing Rate Limiter${NC}\n"
 
-# # 1. Create a gateway with rate limiting plugin
-# echo -e "${GREEN}1. Creating gateway with rate limiting plugin...${NC}"
-# GATEWAY_RESPONSE=$(curl -s -X POST "$ADMIN_URL/gateways" \
-#   -H "Content-Type: application/json" \
-#   -d '{
-#     "name": "Multi Rate Limited Gateway",
-#     "subdomain": "'$SUBDOMAIN'",
-#     "required_plugins": [
-#         {
-#             "name": "rate_limiter",
-#             "enabled": true,
-#             "stage": "pre_request",
-#             "priority": 1,
-#             "settings": {
-#                 "limits": {
-#                     "global": {
-#                         "limit": 15,
-#                         "window": "1m"
-#                     },
-#                     "per_ip": {
-#                         "limit": 5,
-#                         "window": "1m"
-#                     },
-#                     "per_user": {
-#                         "limit": 5,
-#                         "window": "1m"
-#                     }
-#                 },
-#                 "actions": {
-#                     "type": "reject",
-#                     "retry_after": "60"
-#                 }
-#             }
-#         }
-#     ]
-# }')
-# echo $GATEWAY_RESPONSE | jq
-# # Extract gateway details
-# GATEWAY_ID=$(echo $GATEWAY_RESPONSE | jq -r '.id')
-# SUBDOMAIN=$(echo $GATEWAY_RESPONSE | jq -r '.subdomain')
-
-# if [ "$GATEWAY_ID" == "null" ] || [ -z "$GATEWAY_ID" ]; then
-#     echo -e "${RED}Failed to create gateway. Response: $GATEWAY_RESPONSE${NC}"
-#     exit 1
-# fi
-
-# echo "Gateway created with ID: $GATEWAY_ID"
-
-# # Create API key
-# echo -e "\n${GREEN}2. Creating API key...${NC}"
-# API_KEY_RESPONSE=$(curl -s -X POST "$ADMIN_URL/gateways/$GATEWAY_ID/keys" \
-#   -H "Content-Type: application/json" \
-#   -d '{
-#     "name": "Test Key",
-#     "expires_at": "2026-01-01T00:00:00Z"
-# }')
-
-# API_KEY=$(echo $API_KEY_RESPONSE | jq -r '.key')
-
-# if [ "$API_KEY" == "null" ] || [ -z "$API_KEY" ]; then
-#     echo -e "${RED}Failed to create API key. Response: $API_KEY_RESPONSE${NC}"
-#     exit 1
-# fi
-
-# echo "API Key created: $API_KEY"
-
-# # Create upstream
-# echo -e "\n${GREEN}3. Creating upstream...${NC}"
-# UPSTREAM_RESPONSE=$(curl -s -X POST "$ADMIN_URL/gateways/$GATEWAY_ID/upstreams" \
-#   -H "Content-Type: application/json" \
-#   -d '{
-#     "name": "httpbin-upstream-'$(date +%s)'",
-#     "algorithm": "round-robin",
-#     "targets": [{
-#         "host": "httpbin.org",
-#         "port": 443,
-#         "protocol": "https",
-#         "weight": 100,
-#         "priority": 1
-#     }],
-#     "health_checks": {
-#         "passive": true,
-#         "threshold": 3,
-#         "interval": 60
-#     }
-# }')
-
-# UPSTREAM_ID=$(echo $UPSTREAM_RESPONSE | jq -r '.id')
-
-# if [ "$UPSTREAM_ID" == "null" ] || [ -z "$UPSTREAM_ID" ]; then
-#     echo -e "${RED}Failed to create upstream. Response: $UPSTREAM_RESPONSE${NC}"
-#     exit 1
-# fi
-
-# echo "Upstream created with ID: $UPSTREAM_ID"
-
-# # Create service
-# echo -e "\n${GREEN}4. Creating service...${NC}"
-# SERVICE_RESPONSE=$(curl -s -X POST "$ADMIN_URL/gateways/$GATEWAY_ID/services" \
-#   -H "Content-Type: application/json" \
-#   -d '{
-#     "name": "httpbin-service-'$(date +%s)'",
-#     "type": "upstream",
-#     "description": "HTTPBin test service",
-#     "upstream_id": "'$UPSTREAM_ID'"
-# }')
-
-# SERVICE_ID=$(echo $SERVICE_RESPONSE | jq -r '.id')
-
-# if [ "$SERVICE_ID" == "null" ] || [ -z "$SERVICE_ID" ]; then
-#     echo -e "${RED}Failed to create service. Response: $SERVICE_RESPONSE${NC}"
-#     exit 1
-# fi
-
-# echo "Service created with ID: $SERVICE_ID"
-
-# # Create rules for different paths
-# echo -e "\n${GREEN}5. Creating rules for different paths...${NC}"
-# RULE1_RESPONSE=$(curl -s -X POST "$ADMIN_URL/gateways/$GATEWAY_ID/rules" \
-#   -H "Content-Type: application/json" \
-#   -d '{
-#     "path": "/path1",
-#     "service_id": "'$SERVICE_ID'",
-#     "methods": ["GET"],
-#     "strip_path": true,
-#     "active": true
-# }')
-
-# RULE2_RESPONSE=$(curl -s -X POST "$ADMIN_URL/gateways/$GATEWAY_ID/rules" \
-#   -H "Content-Type: application/json" \
-#   -d '{
-#     "path": "/path2",
-#     "service_id": "'$SERVICE_ID'",
-#     "methods": ["GET"],
-#     "strip_path": true,
-#     "active": true
-# }')
-
-# # Wait for configuration to propagate
-# sleep 2
-
-# # Test different rate limit types
-# echo -e "\n${GREEN}6. Testing different rate limit types...${NC}"
+# 1. Create a gateway with rate limiting plugin
+echo -e "${GREEN}1. Creating gateway with rate limiting plugin...${NC}"
+GATEWAY_RESPONSE=$(curl -s -X POST "$ADMIN_URL/gateways" \
+  -H "Content-Type: application/json" \
+  -d '{
+    "name": "Multi Rate Limited Gateway",
+    "subdomain": "'$SUBDOMAIN'",
+    "required_plugins": [
+        {
+            "name": "rate_limiter",
+            "enabled": true,
+            "stage": "pre_request",
+            "priority": 1,
+            "settings": {
+                "limits": {
+                    "global": {
+                        "limit": 15,
+                        "window": "1m"
+                    },
+                    "per_ip": {
+                        "limit": 5,
+                        "window": "1m"
+                    },
+                    "per_user": {
+                        "limit": 5,
+                        "window": "1m"
+                    }
+                },
+                "actions": {
+                    "type": "reject",
+                    "retry_after": "60"
+                }
+            }
+        }
+    ]
+}')
+echo $GATEWAY_RESPONSE | jq
+# Extract gateway details
+GATEWAY_ID=$(echo $GATEWAY_RESPONSE | jq -r '.id')
+SUBDOMAIN=$(echo $GATEWAY_RESPONSE | jq -r '.subdomain')
+
+if [ "$GATEWAY_ID" == "null" ] || [ -z "$GATEWAY_ID" ]; then
+    echo -e "${RED}Failed to create gateway. Response: $GATEWAY_RESPONSE${NC}"
+    exit 1
+fi
+
+echo "Gateway created with ID: $GATEWAY_ID"
+
+# Create API key
+echo -e "\n${GREEN}2. Creating API key...${NC}"
+API_KEY_RESPONSE=$(curl -s -X POST "$ADMIN_URL/gateways/$GATEWAY_ID/keys" \
+  -H "Content-Type: application/json" \
+  -d '{
+    "name": "Test Key",
+    "expires_at": "2026-01-01T00:00:00Z"
+}')
+
+API_KEY=$(echo $API_KEY_RESPONSE | jq -r '.key')
+
+if [ "$API_KEY" == "null" ] || [ -z "$API_KEY" ]; then
+    echo -e "${RED}Failed to create API key. Response: $API_KEY_RESPONSE${NC}"
+    exit 1
+fi
+
+echo "API Key created: $API_KEY"
+
+# Create upstream
+echo -e "\n${GREEN}3. Creating upstream...${NC}"
+UPSTREAM_RESPONSE=$(curl -s -X POST "$ADMIN_URL/gateways/$GATEWAY_ID/upstreams" \
+  -H "Content-Type: application/json" \
+  -d '{
+    "name": "httpbin-upstream-'$(date +%s)'",
+    "algorithm": "round-robin",
+    "targets": [{
+        "host": "httpbin.org",
+        "port": 443,
+        "protocol": "https",
+        "weight": 100,
+        "priority": 1
+    }],
+    "health_checks": {
+        "passive": true,
+        "threshold": 3,
+        "interval": 60
+    }
+}')
+
+UPSTREAM_ID=$(echo $UPSTREAM_RESPONSE | jq -r '.id')
+
+if [ "$UPSTREAM_ID" == "null" ] || [ -z "$UPSTREAM_ID" ]; then
+    echo -e "${RED}Failed to create upstream. Response: $UPSTREAM_RESPONSE${NC}"
+    exit 1
+fi
+
+echo "Upstream created with ID: $UPSTREAM_ID"
+
+# Create service
+echo -e "\n${GREEN}4. Creating service...${NC}"
+SERVICE_RESPONSE=$(curl -s -X POST "$ADMIN_URL/gateways/$GATEWAY_ID/services" \
+  -H "Content-Type: application/json" \
+  -d '{
+    "name": "httpbin-service-'$(date +%s)'",
+    "type": "upstream",
+    "description": "HTTPBin test service",
+    "upstream_id": "'$UPSTREAM_ID'"
+}')
+
+SERVICE_ID=$(echo $SERVICE_RESPONSE | jq -r '.id')
+
+if [ "$SERVICE_ID" == "null" ] || [ -z "$SERVICE_ID" ]; then
+    echo -e "${RED}Failed to create service. Response: $SERVICE_RESPONSE${NC}"
+    exit 1
+fi
+
+echo "Service created with ID: $SERVICE_ID"
+
+# Create rules for different paths
+echo -e "\n${GREEN}5. Creating rules for different paths...${NC}"
+RULE1_RESPONSE=$(curl -s -X POST "$ADMIN_URL/gateways/$GATEWAY_ID/rules" \
+  -H "Content-Type: application/json" \
+  -d '{
+    "path": "/path1",
+    "service_id": "'$SERVICE_ID'",
+    "methods": ["GET"],
+    "strip_path": true,
+    "active": true
+}')
+
+RULE2_RESPONSE=$(curl -s -X POST "$ADMIN_URL/gateways/$GATEWAY_ID/rules" \
+  -H "Content-Type: application/json" \
+  -d '{
+    "path": "/path2",
+    "service_id": "'$SERVICE_ID'",
+    "methods": ["GET"],
+    "strip_path": true,
+    "active": true
+}')
+
+# Wait for configuration to propagate
+sleep 2
+
+# Test different rate limit types
+echo -e "\n${GREEN}6. Testing different rate limit types...${NC}"
 
 # Test IP-based rate limit
 echo -e "\n${GREEN}6.1 Testing IP-based rate limit (limit: 5/min)...${NC}"
-for i in {1..1}; do
+for i in {1..10}; do
     start_time=$(date +%s.%N)
 
-    RESPONSE=$(curl -s "$PROXY_URL/path2" \
+    RESPONSE=$(curl -s -w "\n%{http_code}\n%{time_total}" "$PROXY_URL/path2" \
         -H "Host: $SUBDOMAIN.$BASE_DOMAIN" \
-        -H "X-API-Key: 9x9PD71fWwqVwldtnfl41--HqOnYWlbbsKJ92-aM2sU=")
+        -H "X-API-Key: $API_KEY")
 
-    duration=$(echo "$RESPONSE" | tail -n1)
     HTTP_CODE=$(echo "$RESPONSE" | tail -n2 | head -n1)
     BODY=$(echo "$RESPONSE" | head -n1)
-    end_time=$(date +%s.%N)
-    total_time=$(echo "$end_time - $start_time" | bc)
 
-    echo -e "Request duration (curl): ${duration}s"
-    echo -e "Total script duration: ${total_time}s"
 
     if [ "$HTTP_CODE" == "200" ]; then
         echo -e "${GREEN}IP-based Request $i: Success${NC}"

diff --git a/pkg/server/proxy_server.go b/pkg/server/proxy_server.go
@@ -457,7 +457,6 @@ func (s *ProxyServer) HandleForward(c *gin.Context) {
 		c.JSON(http.StatusBadGateway, gin.H{"error": "Failed to forward request"})
 		return
 	}
-
 	// Record upstream latency if available
 	if metrics.Config.EnableUpstreamLatency {
 		upstreamLatency := float64(time.Since(startTime).Milliseconds())
@@ -865,6 +864,9 @@ func (s *ProxyServer) ForwardRequest(req *types.RequestContext, rule *types.Forw
 			}).Debug("Attempting request")
 
 			response, err := s.doForwardRequest(req, rule, target, service.Type, lb)
+			s.logger.WithFields(logrus.Fields{
+				"response": response,
+			}).Debug("Forwarded request")
 			if err == nil {
 				lb.ReportSuccess(target)
 				return response, nil
@@ -995,6 +997,9 @@ func (s *ProxyServer) doForwardRequest(req *types.RequestContext, rule *types.Fo
 	}
 
 	// Set provider in response header
+	s.logger.WithFields(logrus.Fields{
+		"provider": target.Provider,
+	}).Debug("Selected provider")
 	httpResp.Header.Set("X-Selected-Provider", target.Provider)
 
 	// Handle response status
@@ -1097,7 +1102,7 @@ func (s *ProxyServer) handleStreamingResponse(req *types.RequestContext, target
 				}
 			}
 		}
-
+		w.Header().Add("X-Selected-Provider", target.Provider)
 		w.WriteHeader(resp.StatusCode)
 
 		reader := bufio.NewReader(resp.Body)
@@ -1163,7 +1168,7 @@ func (s *ProxyServer) handleStreamingResponse(req *types.RequestContext, target
 
 	return &types.ResponseContext{
 		StatusCode: resp.StatusCode,
-		Headers:    make(map[string][]string),
+		Headers:    resp.Header,
 		Streaming:  true,
 		Metadata:   req.Metadata, // Include the metadata with token usage
 	}, nil