github.com/dbernstein1/tyk@v2.9.0-beta9-dl-apic+incompatible/gateway/mw_ip_whitelist_test.go (about) 1 package gateway 2 3 import ( 4 "net/http" 5 "net/http/httptest" 6 "testing" 7 ) 8 9 var testWhiteListIPData = []struct { 10 remote, forwarded string 11 wantCode int 12 }{ 13 {"127.0.0.1:80", "", http.StatusOK}, // remote exact match 14 {"127.0.0.2:80", "", http.StatusOK}, // remote CIDR match 15 {"10.0.0.1:80", "", http.StatusForbidden}, // no match 16 {"10.0.0.1:80", "127.0.0.1", http.StatusOK}, // forwarded exact match 17 {"10.0.0.1:80", "127.0.0.2", http.StatusOK}, // forwarded CIDR match 18 } 19 20 func testPrepareIPMiddlewarePass() *APISpec { 21 return BuildAPI(func(spec *APISpec) { 22 spec.EnableIpWhiteListing = true 23 spec.AllowedIPs = []string{"127.0.0.1", "127.0.0.1/24"} 24 })[0] 25 } 26 27 func TestIPMiddlewarePass(t *testing.T) { 28 spec := testPrepareIPMiddlewarePass() 29 30 for ti, tc := range testWhiteListIPData { 31 rec := httptest.NewRecorder() 32 req := TestReq(t, "GET", "/", nil) 33 req.RemoteAddr = tc.remote 34 if tc.forwarded != "" { 35 req.Header.Set("X-Forwarded-For", tc.forwarded) 36 } 37 38 mw := &IPWhiteListMiddleware{} 39 mw.Spec = spec 40 _, code := mw.ProcessRequest(rec, req, nil) 41 42 if code != tc.wantCode { 43 t.Errorf("[%d] Response code %d should be %d\n%q %q", ti, 44 code, tc.wantCode, tc.remote, tc.forwarded) 45 } 46 } 47 } 48 49 func BenchmarkIPMiddlewarePass(b *testing.B) { 50 b.ReportAllocs() 51 52 spec := testPrepareIPMiddlewarePass() 53 mw := &IPWhiteListMiddleware{} 54 mw.Spec = spec 55 56 rec := httptest.NewRecorder() 57 for i := 0; i < b.N; i++ { 58 for ti, tc := range testWhiteListIPData { 59 req := TestReq(b, "GET", "/", nil) 60 req.RemoteAddr = tc.remote 61 if tc.forwarded != "" { 62 req.Header.Set("X-Forwarded-For", tc.forwarded) 63 } 64 65 _, code := mw.ProcessRequest(rec, req, nil) 66 67 if code != tc.wantCode { 68 b.Errorf("[%d] Response code %d should be %d\n%q %q", ti, 69 code, tc.wantCode, tc.remote, tc.forwarded) 70 } 71 } 72 } 73 }