gitee.com/ks-custle/core-gm@v0.0.0-20230922171213-b83bdd97b62c/handlers/proxy_headers_test.go (about)

     1  package handlers
     2  
     3  import (
     4  	http "gitee.com/ks-custle/core-gm/gmhttp"
     5  	"gitee.com/ks-custle/core-gm/gmhttp/httptest"
     6  	"testing"
     7  )
     8  
     9  type headerTable struct {
    10  	key      string // header key
    11  	val      string // header val
    12  	expected string // expected result
    13  }
    14  
    15  func TestGetIP(t *testing.T) {
    16  	headers := []headerTable{
    17  		{xForwardedFor, "8.8.8.8", "8.8.8.8"},                                   // Single address
    18  		{xForwardedFor, "8.8.8.8, 8.8.4.4", "8.8.8.8"},                          // Multiple
    19  		{xForwardedFor, "[2001:db8:cafe::17]:4711", "[2001:db8:cafe::17]:4711"}, // IPv6 address
    20  		{xForwardedFor, "", ""},                                                 // None
    21  		{xRealIP, "8.8.8.8", "8.8.8.8"},                                         // Single address
    22  		{xRealIP, "8.8.8.8, 8.8.4.4", "8.8.8.8, 8.8.4.4"},                       // Multiple
    23  		{xRealIP, "[2001:db8:cafe::17]:4711", "[2001:db8:cafe::17]:4711"},       // IPv6 address
    24  		{xRealIP, "", ""},                       // None
    25  		{forwarded, `for="_gazonk"`, "_gazonk"}, // Hostname
    26  		{forwarded, `For="[2001:db8:cafe::17]:4711`, `[2001:db8:cafe::17]:4711`},      // IPv6 address
    27  		{forwarded, `for=192.0.2.60;proto=http;by=203.0.113.43`, `192.0.2.60`},        // Multiple params
    28  		{forwarded, `for=192.0.2.43, for=198.51.100.17`, "192.0.2.43"},                // Multiple params
    29  		{forwarded, `for="workstation.local",for=198.51.100.17`, "workstation.local"}, // Hostname
    30  	}
    31  
    32  	for _, v := range headers {
    33  		req := &http.Request{
    34  			Header: http.Header{
    35  				v.key: []string{v.val},
    36  			}}
    37  		res := getIP(req)
    38  		if res != v.expected {
    39  			t.Fatalf("wrong header for %s: got %s want %s", v.key, res,
    40  				v.expected)
    41  		}
    42  	}
    43  }
    44  
    45  func TestGetScheme(t *testing.T) {
    46  	headers := []headerTable{
    47  		{xForwardedProto, "https", "https"},
    48  		{xForwardedProto, "http", "http"},
    49  		{xForwardedProto, "HTTP", "http"},
    50  		{xForwardedScheme, "https", "https"},
    51  		{xForwardedScheme, "http", "http"},
    52  		{xForwardedScheme, "HTTP", "http"},
    53  		{forwarded, `For="[2001:db8:cafe::17]:4711`, ""},                      // No proto
    54  		{forwarded, `for=192.0.2.43, for=198.51.100.17;proto=https`, "https"}, // Multiple params before proto
    55  		{forwarded, `for=172.32.10.15; proto=https;by=127.0.0.1`, "https"},    // Space before proto
    56  		{forwarded, `for=192.0.2.60;proto=http;by=203.0.113.43`, "http"},      // Multiple params
    57  	}
    58  
    59  	for _, v := range headers {
    60  		req := &http.Request{
    61  			Header: http.Header{
    62  				v.key: []string{v.val},
    63  			},
    64  		}
    65  		res := getScheme(req)
    66  		if res != v.expected {
    67  			t.Fatalf("wrong header for %s: got %s want %s", v.key, res,
    68  				v.expected)
    69  		}
    70  	}
    71  }
    72  
    73  // Test the middleware end-to-end
    74  func TestProxyHeaders(t *testing.T) {
    75  	rr := httptest.NewRecorder()
    76  	r := newRequest("GET", "/")
    77  
    78  	r.Header.Set(xForwardedFor, "8.8.8.8")
    79  	r.Header.Set(xForwardedProto, "https")
    80  	r.Header.Set(xForwardedHost, "google.com")
    81  	var (
    82  		addr  string
    83  		proto string
    84  		host  string
    85  	)
    86  	ProxyHeaders(http.HandlerFunc(
    87  		func(w http.ResponseWriter, r *http.Request) {
    88  			addr = r.RemoteAddr
    89  			proto = r.URL.Scheme
    90  			host = r.Host
    91  		})).ServeHTTP(rr, r)
    92  
    93  	if rr.Code != http.StatusOK {
    94  		t.Fatalf("bad status: got %d want %d", rr.Code, http.StatusOK)
    95  	}
    96  
    97  	if addr != r.Header.Get(xForwardedFor) {
    98  		t.Fatalf("wrong address: got %s want %s", addr,
    99  			r.Header.Get(xForwardedFor))
   100  	}
   101  
   102  	if proto != r.Header.Get(xForwardedProto) {
   103  		t.Fatalf("wrong address: got %s want %s", proto,
   104  			r.Header.Get(xForwardedProto))
   105  	}
   106  	if host != r.Header.Get(xForwardedHost) {
   107  		t.Fatalf("wrong address: got %s want %s", host,
   108  			r.Header.Get(xForwardedHost))
   109  	}
   110  
   111  }