package middleware import ( _ "embed" "net" "net/http" "testing" E "github.com/yusing/go-proxy/internal/error" "github.com/yusing/go-proxy/internal/utils" . "github.com/yusing/go-proxy/internal/utils/testing" ) //go:embed test_data/cidr_whitelist_test.yml var testCIDRWhitelistCompose []byte var deny, accept *Middleware func TestCIDRWhitelistValidation(t *testing.T) { t.Run("valid", func(t *testing.T) { _, err := NewCIDRWhitelist(OptionsRaw{ "allow": []string{"1.2.3.4/32"}, "message": "test-message", }) ExpectNoError(t, err) }) t.Run("missing allow", func(t *testing.T) { _, err := NewCIDRWhitelist(OptionsRaw{ "message": "test-message", }) ExpectError(t, utils.ErrValidationError, err) }) t.Run("invalid cidr", func(t *testing.T) { _, err := NewCIDRWhitelist(OptionsRaw{ "allow": []string{"1.2.3.4/123"}, "message": "test-message", }) ExpectErrorT[*net.ParseError](t, err) }) t.Run("invalid status code", func(t *testing.T) { _, err := NewCIDRWhitelist(OptionsRaw{ "allow": []string{"1.2.3.4/32"}, "status_code": 600, "message": "test-message", }) ExpectError(t, utils.ErrValidationError, err) }) } func TestCIDRWhitelist(t *testing.T) { errs := E.NewBuilder("") mids := BuildMiddlewaresFromYAML("", testCIDRWhitelistCompose, errs) ExpectNoError(t, errs.Error()) deny = mids["deny@file"] accept = mids["accept@file"] if deny == nil || accept == nil { panic("bug occurred") } t.Run("deny", func(t *testing.T) { t.Parallel() for range 10 { result, err := newMiddlewareTest(deny, nil) ExpectNoError(t, err) ExpectEqual(t, result.ResponseStatus, cidrWhitelistDefaults.StatusCode) ExpectEqual(t, string(result.Data), cidrWhitelistDefaults.Message) } }) t.Run("accept", func(t *testing.T) { t.Parallel() for range 10 { result, err := newMiddlewareTest(accept, nil) ExpectNoError(t, err) ExpectEqual(t, result.ResponseStatus, http.StatusOK) } }) }