Merge pull request #19 from sygmaprotocol/feat/control-access-per-token

feat: add multiple rate limited auth tokens

Author: MakMuftic

.gitignore

@@ -6,7 +6,3 @@
 
 .idea
 
-config.json
-config_*
-
-prometheus

README.md

@@ -107,15 +107,40 @@ Each JSON configuration file for the gateways can specify detailed settings for
 ```
 
 ## Authentication
-Authentication can be enabled using the `--auth` flag. The auth token should be set through environment variables `GATEWAY_PASSWORD`.
 
-Auth token needs to be the last entry in the RPC gateway URL. Example:
+Authentication can be enabled using the `--auth` flag. The authentication system uses a token-based approach with rate limiting.
 
-`https://sample/rpc-gateway/sepolia/a1b2c3d4e5f7`
+### Token Configuration
 
-### Running the Application
-To run the application with authentication:
+The token configuration should be provided through the `GATEWAY_TOKEN_MAP` environment variable. This variable should contain a JSON string representing a map of tokens to their corresponding information. Each token entry includes a name and the number of requests allowed per second.
 
+Example of `GATEWAY_TOKEN_MAP`:
+
+```json
+{
+  "token1": {"name": "User1", "numOfRequestPerSec": 10},
+  "token2": {"name": "User2", "numOfRequestPerSec": 20}
+}
 ```
-DEBUG=true GATEWAY_PASSWORD=my_auth_token go run . --config config.json --auth
-```
+
+### URL Format
+
+When authentication is enabled, the auth token needs to be the last entry in the RPC gateway URL. 
+
+Example:
+
+`https://sample/rpc-gateway/sepolia/token1`
+
+In this example, `token1` is the authentication token that must match one of the tokens defined in the `GATEWAY_TOKEN_MAP`.
+
+### Rate Limiting
+
+Each token has its own rate limit, defined by the `numOfRequestPerSec` value in the token configuration. If a client exceeds this limit, they will receive a 429 (Too Many Requests) status code.
+
+### Running the Application with Authentication
+
+To run the application with authentication:
+
+```bash
+export GATEWAY_TOKEN_MAP='{"token1":{"name":"User1","numOfRequestPerSec":10},"token2":{"name":"User2","numOfRequestPerSec":20}}'
+DEBUG=true go run . --config config.json --auth

config.json

@@ -0,0 +1,16 @@
+{
+  "metrics": {
+    "port": 9090
+  },
+  "port": 4000,
+  "gateways": [
+    {
+      "configFile": "/app/config_holesky.json",
+      "name": "Holesky gateway"
+    },
+    {
+      "configFile": "/app/config_sepolia.json",
+      "name": "Sepolia gateway"
+    }
+  ]
+}

config_holesky.json

@@ -0,0 +1,31 @@
+{
+  "name": "Holesky",
+  "proxy": {
+    "path": "holesky",
+    "upstreamTimeout": "1s"
+  },
+  "healthChecks": {
+    "interval": "20s",
+    "timeout": "1s",
+    "failureThreshold": 2,
+    "successThreshold": 1
+  },
+  "targets": [
+    {
+      "name": "ChainSafe",
+      "connection": {
+        "http": {
+          "url": "https://lodestar-holeskyrpc.chainsafe.io/"
+        }
+      }
+    },
+    {
+      "name": "Tenderly",
+      "connection": {
+        "http": {
+          "url": "https://holesky.gateway.tenderly.co"
+        }
+      }
+    }
+  ]
+}

config_sepolia.json

@@ -0,0 +1,31 @@
+{
+  "name": "Sepolia",
+  "proxy": {
+    "path": "sepolia",
+    "upstreamTimeout": "1s"
+  },
+  "healthChecks": {
+    "interval": "20s",
+    "timeout": "1s",
+    "failureThreshold": 2,
+    "successThreshold": 1
+  },
+  "targets": [
+    {
+      "name": "ChainSafe",
+      "connection": {
+        "http": {
+          "url": "https://lodestar-sepoliarpc.chainsafe.io"
+        }
+      }
+    },
+    {
+      "name": "Tenderly",
+      "connection": {
+        "http": {
+          "url": "https://sepolia.gateway.tenderly.co"
+        }
+      }
+    }
+  ]
+}

docker-compose.yml

@@ -0,0 +1,49 @@
+services:
+  rpc-gateway:
+    build:
+      context: .
+      dockerfile: Dockerfile
+    ports:
+      - 8080:4000  # Main port
+      - 9090:9090  # Metrics port
+    volumes:
+      - ./config.json:/app/config.json:ro
+      - ./config_sepolia.json:/app/config_sepolia.json:ro
+      - ./config_holesky.json:/app/config_holesky.json:ro
+    environment:
+      - GATEWAY_TOKEN_MAP={"token1":{"name":"token1","numOfRequestPerSec":10},"token2":{"name":"token2","numOfRequestPerSec":20}}
+    user: nobody
+    entrypoint: ["/app/rpc-gateway", "--config", "/app/config.json", "--auth"]
+    networks:
+      - app-network
+
+  prometheus:
+    image: prom/prometheus:v2.44.0
+    volumes:
+      - ./prometheus.yml:/etc/prometheus/prometheus.yml
+    command:
+      - '--config.file=/etc/prometheus/prometheus.yml'
+    ports:
+      - "9091:9090"  # Changed to 9091 on the host
+    networks:
+      - app-network
+
+  grafana:
+    image: grafana/grafana:latest
+    ports:
+      - 3000:3000
+    volumes:
+      - grafana-storage:/var/lib/grafana
+    environment:
+      - GF_SECURITY_ADMIN_PASSWORD=admin
+    depends_on:
+      - rpc-gateway
+    networks:
+      - app-network
+
+volumes:
+  grafana-storage:
+
+networks:
+  app-network:
+    driver: bridge

go.mod

@@ -44,6 +44,7 @@ require (
 	golang.org/x/mod v0.15.0 // indirect
 	golang.org/x/net v0.21.0 // indirect
 	golang.org/x/sys v0.17.0 // indirect
+	golang.org/x/time v0.7.0
 	golang.org/x/tools v0.18.0 // indirect
 	google.golang.org/protobuf v1.32.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect

go.sum

@@ -103,6 +103,8 @@ golang.org/x/sys v0.0.0-20190916202348-b4ddaad3f8a3/go.mod h1:h1NjWce9XRLGQEsW7w
 golang.org/x/sys v0.1.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.17.0 h1:25cE3gD+tdBA7lp7QfhuV+rJiE9YXTcS3VG1SqssI/Y=
 golang.org/x/sys v0.17.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
+golang.org/x/time v0.7.0 h1:ntUhktv3OPE6TgYxXWv9vKvUSJyIFJlyohwbkEwPrKQ=
+golang.org/x/time v0.7.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM=
 golang.org/x/tools v0.18.0 h1:k8NLag8AGHnn+PHbl7g43CtqZAwG60vZkLqgyZgIHgQ=
 golang.org/x/tools v0.18.0/go.mod h1:GL7B4CwcLLeo59yx/9UWWuNOW1n3VZ4f5axWfML7Lcg=
 google.golang.org/protobuf v1.32.0 h1:pPC6BG5ex8PDFnkbrGU3EixyhKcQ2aDuBS36lqK/C7I=

internal/auth/auth.go

@@ -1,21 +1,70 @@
 package auth
 
 import (
+	"context"
+	"fmt"
 	"net/http"
 	"strings"
+
+	"golang.org/x/time/rate"
 )
 
-func URLTokenAuth(token string) func(next http.Handler) http.Handler {
+type TokenInfo struct {
+	Name               string `json:"name"`
+	NumOfRequestPerSec int    `json:"numOfRequestPerSec"`
+}
+
+// ContextKeyType custom type for the context key.
+type ContextKeyType string
+
+const TokenInfoKey ContextKeyType = "tokeninfo"
+
+func URLTokenAuth(tokenToName map[string]TokenInfo) func(next http.Handler) http.Handler {
+	limiters := make(map[string]*rate.Limiter)
+	for token, info := range tokenToName {
+		limiters[token] = rate.NewLimiter(rate.Limit(info.NumOfRequestPerSec), info.NumOfRequestPerSec)
+		fmt.Printf("Configured limiter for %s, allowed %d requests per second\n",
+			info.Name, info.NumOfRequestPerSec,
+		)
+	}
+
 	return func(next http.Handler) http.Handler {
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			pathParts := strings.Split(r.URL.Path, "/")
-			if len(pathParts) < 2 || pathParts[len(pathParts)-1] != token {
+			if len(pathParts) < 2 {
+				w.WriteHeader(http.StatusUnauthorized)
+
+				return
+			}
+
+			token := pathParts[len(pathParts)-1]
+			tInfo, validToken := tokenToName[token]
+			if !validToken {
 				w.WriteHeader(http.StatusUnauthorized)
 
 				return
 			}
-			// Remove the token part from the path to forward the request to the next handler
+
+			limiter, exists := limiters[token]
+			if !exists {
+				w.WriteHeader(http.StatusInternalServerError)
+
+				return
+			}
+
+			if !limiter.Allow() {
+				w.WriteHeader(http.StatusTooManyRequests)
+
+				return
+			}
+
+			// Remove the token part from the path
 			r.URL.Path = strings.Join(pathParts[:len(pathParts)-1], "/")
+
+			// Add the user's name to the request context
+			ctx := context.WithValue(r.Context(), TokenInfoKey, tInfo)
+			r = r.WithContext(ctx)
+
 			next.ServeHTTP(w, r)
 		})
 	}

internal/auth/auth_test.go

@@ -4,11 +4,16 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"testing"
+	"time"
 )
 
 func TestURLTokenAuth(t *testing.T) {
 	validToken := "valid_token"
-	middleware := URLTokenAuth(validToken)
+	tokenInfo := TokenInfo{
+		Name:               "Test User",
+		NumOfRequestPerSec: 1, // Changed from 60 per minute to 1 per second
+	}
+	tokenMap := map[string]TokenInfo{validToken: tokenInfo}
 
 	tests := []struct {
 		name           string
@@ -21,7 +26,7 @@ func TestURLTokenAuth(t *testing.T) {
 			expectedStatus: http.StatusOK,
 		},
 		{
-			name:           "Valid token",
+			name:           "Valid token with long path",
 			url:            "/some/really/long/path/valid_token",
 			expectedStatus: http.StatusOK,
 		},
@@ -39,6 +44,7 @@ func TestURLTokenAuth(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
+			middleware := URLTokenAuth(tokenMap)
 			req, err := http.NewRequest("GET", tt.url, nil)
 			if err != nil {
 				t.Fatalf("could not create request: %v", err)
@@ -57,3 +63,50 @@ func TestURLTokenAuth(t *testing.T) {
 		})
 	}
 }
+
+func TestURLTokenAuthRateLimit(t *testing.T) {
+	validToken := "valid_token"
+	tokenInfo := TokenInfo{
+		Name:               "Test User",
+		NumOfRequestPerSec: 5, // Changed from 60 per minute to 1 per second
+	}
+	tokenMap := map[string]TokenInfo{validToken: tokenInfo}
+	middleware := URLTokenAuth(tokenMap)
+
+	handler := middleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+	}))
+
+	url := "/some/path/valid_token"
+
+	// Make requests up to the limit
+	for i := 0; i < tokenInfo.NumOfRequestPerSec; i++ {
+		req, _ := http.NewRequest("GET", url, nil)
+		rr := httptest.NewRecorder()
+		handler.ServeHTTP(rr, req)
+		if rr.Code != http.StatusOK {
+			t.Fatalf("Expected OK for request %d, got %d", i, rr.Code)
+		}
+	}
+
+	// This request should exceed the rate limit
+	req, _ := http.NewRequest("GET", url, nil)
+	rr := httptest.NewRecorder()
+	handler.ServeHTTP(rr, req)
+
+	if rr.Code != http.StatusTooManyRequests {
+		t.Errorf("Expected status %v for rate limit exceeded; got %v", http.StatusTooManyRequests, rr.Code)
+	}
+
+	// Wait for a second to allow the rate limiter to reset
+	time.Sleep(time.Second)
+
+	// This request should now succeed
+	req, _ = http.NewRequest("GET", url, nil)
+	rr = httptest.NewRecorder()
+	handler.ServeHTTP(rr, req)
+
+	if rr.Code != http.StatusOK {
+		t.Errorf("Expected status %v after rate limit reset; got %v", http.StatusOK, rr.Code)
+	}
+}