ratelimit_test.go

package main

import (
	"fmt"
	"os"
	"path/filepath"
	"time"

	"github.com/goccy/go-json"
	goratecounter "github.com/lukaszraczylo/go-ratecounter"
	libpack_logger "github.com/lukaszraczylo/graphql-monitoring-proxy/logging"
)

func (suite *Tests) Test_loadRatelimitConfig() {
	// Setup
	cfg = &config{}
	parseConfig()
	cfg.Logger = libpack_logger.New()

	// Create a temporary test ratelimit.json file
	tempDir := os.TempDir()
	testConfigPath := filepath.Join(tempDir, "test_ratelimit.json")

	testConfig := struct {
		RateLimit map[string]RateLimitConfig `json:"ratelimit"`
	}{
		RateLimit: map[string]RateLimitConfig{
			"admin": {
				Interval: 1 * time.Second,
				Req:      100,
			},
			"user": {
				Interval: 1 * time.Second,
				Req:      10,
			},
		},
	}

	configData, err := json.Marshal(testConfig)
	suite.NoError(err)

	err = os.WriteFile(testConfigPath, configData, 0o644)
	suite.NoError(err)
	defer func() { _ = os.Remove(testConfigPath) }()

	// Test loading config from custom path
	suite.Run("load from custom path", func() {
		// Clear existing rate limits
		rateLimitMu.Lock()
		rateLimits = make(map[string]RateLimitConfig)
		rateLimitMu.Unlock()

		err := loadConfigFromPath(testConfigPath)
		suite.NoError(err)

		// Verify rate limits were loaded
		rateLimitMu.RLock()
		defer rateLimitMu.RUnlock()

		suite.Equal(2, len(rateLimits))
		suite.Contains(rateLimits, "admin")
		suite.Contains(rateLimits, "user")
		suite.Equal(100, rateLimits["admin"].Req)
		suite.Equal(10, rateLimits["user"].Req)
		suite.NotNil(rateLimits["admin"].RateCounterTicker)
		suite.NotNil(rateLimits["user"].RateCounterTicker)
	})

	// Test loading config from non-existent path
	suite.Run("load from non-existent path", func() {
		err := loadConfigFromPath("/non/existent/path.json")
		suite.Error(err)
	})

	// Test loading config with invalid JSON
	suite.Run("load invalid JSON", func() {
		invalidPath := filepath.Join(tempDir, "invalid_ratelimit.json")
		err := os.WriteFile(invalidPath, []byte("{invalid json}"), 0o644)
		suite.NoError(err)
		defer func() { _ = os.Remove(invalidPath) }()

		err = loadConfigFromPath(invalidPath)
		suite.Error(err)
	})

	// Test with a temporary ratelimit.json file in the current directory
	suite.Run("load from current directory", func() {
		// Create a temporary ratelimit.json in current directory
		currentDirPath := "./ratelimit.json"
		err := os.WriteFile(currentDirPath, configData, 0o644)
		suite.NoError(err)
		defer func() { _ = os.Remove(currentDirPath) }()

		// Clear existing rate limits
		rateLimitMu.Lock()
		rateLimits = make(map[string]RateLimitConfig)
		rateLimitMu.Unlock()

		// This should find the file in the current directory
		err = loadRatelimitConfig()
		suite.NoError(err)

		// Verify rate limits were loaded
		rateLimitMu.RLock()
		defer rateLimitMu.RUnlock()

		suite.Equal(2, len(rateLimits))
	})

	// Test with all files missing
	suite.Run("all files missing", func() {
		// Save the original load function and restore it when done
		originalLoadFunc := loadConfigFunc
		defer func() {
			loadConfigFunc = originalLoadFunc
		}()

		// Replace with a mock function that always returns "file does not exist" error
		loadConfigFunc = func(string) error {
			return fmt.Errorf("file does not exist")
		}

		// Clear existing rate limits
		rateLimitMu.Lock()
		rateLimits = make(map[string]RateLimitConfig)
		rateLimitMu.Unlock()

		// This should fail as our mock returns errors for all paths
		err = loadRatelimitConfig()
		suite.Error(err)

		// The error should be a RateLimitConfigError
		configErr, ok := err.(*RateLimitConfigError)
		suite.True(ok, "Expected *RateLimitConfigError but got %T", err)

		// All path errors should contain our mock error message
		for _, errMsg := range configErr.PathErrors {
			suite.Equal("file does not exist", errMsg)
		}
	})
}

func (suite *Tests) Test_rateLimitedRequest() {
	// Setup
	cfg = &config{}
	parseConfig()
	cfg.Logger = libpack_logger.New()

	// Create test rate limits
	rateLimitMu.Lock()
	rateLimits = make(map[string]RateLimitConfig)

	// Admin role with high limit
	adminCounter := goratecounter.NewRateCounter().WithConfig(goratecounter.RateCounterConfig{
		Interval: 1 * time.Second,
	})
	rateLimits["admin"] = RateLimitConfig{
		RateCounterTicker: adminCounter,
		Interval:          1 * time.Second,
		Req:               100,
	}

	// User role with low limit
	userCounter := goratecounter.NewRateCounter().WithConfig(goratecounter.RateCounterConfig{
		Interval: 1 * time.Second,
	})
	rateLimits["user"] = RateLimitConfig{
		RateCounterTicker: userCounter,
		Interval:          1 * time.Second,
		Req:               2, // Set very low for testing
	}
	rateLimitMu.Unlock()

	// Test non-existent role - should be denied for security
	suite.Run("non-existent role", func() {
		allowed := rateLimitedRequest("test-user-1", "non-existent-role")
		suite.False(allowed, "Unknown roles should be denied for security")
	})

	// Test admin role (high limit)
	suite.Run("admin role within limit", func() {
		allowed := rateLimitedRequest("admin-user", "admin")
		suite.True(allowed, "Admin should be within rate limit")
	})

	// Test user role (low limit)
	suite.Run("user role within limit", func() {
		// First request should be allowed
		allowed := rateLimitedRequest("regular-user", "user")
		suite.True(allowed, "First request should be within rate limit")

		// Second request should be allowed
		allowed = rateLimitedRequest("regular-user", "user")
		suite.True(allowed, "Second request should be within rate limit")

		// Third request should exceed limit
		allowed = rateLimitedRequest("regular-user", "user")
		suite.False(allowed, "Third request should exceed rate limit")
	})
}

func (suite *Tests) Test_RateLimitConfig_UnmarshalJSON() {
	// Test unmarshaling of string-based intervals
	suite.Run("unmarshal string intervals", func() {
		// Test JSON with string-based intervals
		jsonString := `{
			"ratelimit": {
				"admin": {
					"req": 100,
					"interval": "second"
				},
				"guest": {
					"req": 5,
					"interval": "minute"
				},
				"user": {
					"req": 1000,
					"interval": "hour"
				},
				"service": {
					"req": 10000,
					"interval": "day"
				},
				"custom": {
					"req": 50,
					"interval": "5s"
				}
			}
		}`

		var config struct {
			RateLimit map[string]RateLimitConfig `json:"ratelimit"`
		}

		err := json.Unmarshal([]byte(jsonString), &config)
		suite.NoError(err)

		// Verify correct parsing of intervals
		suite.Equal(time.Second, config.RateLimit["admin"].Interval)
		suite.Equal(time.Minute, config.RateLimit["guest"].Interval)
		suite.Equal(time.Hour, config.RateLimit["user"].Interval)
		suite.Equal(24*time.Hour, config.RateLimit["service"].Interval)
		suite.Equal(5*time.Second, config.RateLimit["custom"].Interval)

		// Verify req values
		suite.Equal(100, config.RateLimit["admin"].Req)
		suite.Equal(5, config.RateLimit["guest"].Req)
	})

	// Test unmarshaling of invalid interval formats
	suite.Run("unmarshal invalid intervals", func() {
		// Test with an invalid interval format
		jsonString := `{
			"req": 100,
			"interval": "invalid_format"
		}`

		var config RateLimitConfig
		err := json.Unmarshal([]byte(jsonString), &config)
		suite.Error(err)
		suite.Contains(err.Error(), "invalid duration format")
	})

	// Test unmarshaling of numeric intervals
	suite.Run("unmarshal numeric intervals", func() {
		// Test with a numeric interval (seconds)
		jsonString := `{
			"req": 100,
			"interval": 60
		}`

		var config RateLimitConfig
		err := json.Unmarshal([]byte(jsonString), &config)
		suite.NoError(err)
		suite.Equal(60*time.Second, config.Interval)
	})
}