github.com/Tyktechnologies/tyk@v2.9.5+incompatible/gateway/mw_ip_blacklist_test.go (about)

     1  package gateway
     2  
     3  import (
     4  	"net/http"
     5  	"net/http/httptest"
     6  	"testing"
     7  )
     8  
     9  var testBlackListIPData = []struct {
    10  	remote, forwarded string
    11  	wantCode          int
    12  }{
    13  	{"127.0.0.1:80", "", http.StatusForbidden},         // remote exact match
    14  	{"127.0.0.2:80", "", http.StatusForbidden},         // remote CIDR match
    15  	{"10.0.0.1:80", "", http.StatusOK},                 // no match
    16  	{"10.0.0.1:80", "127.0.0.1", http.StatusForbidden}, // forwarded exact match
    17  	{"10.0.0.1:80", "127.0.0.2", http.StatusForbidden}, // forwarded CIDR match
    18  }
    19  
    20  func testPrepareIPBlacklistMiddleware() *APISpec {
    21  	return BuildAPI(func(spec *APISpec) {
    22  		spec.EnableIpBlacklisting = true
    23  		spec.BlacklistedIPs = []string{"127.0.0.1", "127.0.0.1/24"}
    24  	})[0]
    25  }
    26  
    27  func TestIPBlacklistMiddleware(t *testing.T) {
    28  	spec := testPrepareIPBlacklistMiddleware()
    29  
    30  	for ti, tc := range testBlackListIPData {
    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 := &IPBlackListMiddleware{}
    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 BenchmarkIPBlacklistMiddleware(b *testing.B) {
    50  	b.ReportAllocs()
    51  
    52  	spec := testPrepareIPBlacklistMiddleware()
    53  
    54  	mw := &IPBlackListMiddleware{}
    55  	mw.Spec = spec
    56  
    57  	rec := httptest.NewRecorder()
    58  	for i := 0; i < b.N; i++ {
    59  		for ti, tc := range testBlackListIPData {
    60  			req := TestReq(b, "GET", "/", nil)
    61  			req.RemoteAddr = tc.remote
    62  			if tc.forwarded != "" {
    63  				req.Header.Set("X-Forwarded-For", tc.forwarded)
    64  			}
    65  			_, code := mw.ProcessRequest(rec, req, nil)
    66  			if code != tc.wantCode {
    67  				b.Errorf("[%d] Response code %d should be %d\n%q %q", ti,
    68  					code, tc.wantCode, tc.remote, tc.forwarded)
    69  			}
    70  		}
    71  	}
    72  }