github.com/anuvu/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  }