101 lines
2.4 KiB
Go
101 lines
2.4 KiB
Go
|
package middleware
|
||
|
|
|
||
|
|
import (
|
||
|
|
"io/ioutil"
|
||
|
|
"net/http"
|
||
|
|
"net/http/httptest"
|
||
|
|
"strings"
|
||
|
|
"testing"
|
||
|
|
|
||
|
|
"github.com/stretchr/testify/assert"
|
||
|
|
"github.com/stretchr/testify/require"
|
||
|
|
)
|
||
|
|
|
||
|
|
var realIPTestHandler = func(t *testing.T) http.HandlerFunc {
|
||
|
|
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
|
|
assert.Equal(t, r.RemoteAddr, "3.3.3.3")
|
||
|
|
|
||
|
|
_, err := w.Write([]byte("resp"))
|
||
|
|
require.NoError(t, err)
|
||
|
|
})
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestRealIPXRealIP(t *testing.T) {
|
||
|
|
server := httptest.NewServer(RealIP(realIPTestHandler(t)))
|
||
|
|
defer server.Close()
|
||
|
|
|
||
|
|
client := http.Client{}
|
||
|
|
|
||
|
|
req, err := http.NewRequest("GET", server.URL, nil)
|
||
|
|
require.Nil(t, err)
|
||
|
|
req.Header.Set("X-Real-IP", "3.3.3.3")
|
||
|
|
|
||
|
|
res, err := client.Do(req)
|
||
|
|
require.NoError(t, err)
|
||
|
|
defer res.Body.Close()
|
||
|
|
assert.Equal(t, http.StatusOK, res.StatusCode)
|
||
|
|
|
||
|
|
b, err := ioutil.ReadAll(res.Body)
|
||
|
|
assert.NoError(t, err)
|
||
|
|
assert.Equal(t, "resp", string(b))
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestRealIPXForwardedFor(t *testing.T) {
|
||
|
|
server := httptest.NewServer(RealIP(realIPTestHandler(t)))
|
||
|
|
defer server.Close()
|
||
|
|
|
||
|
|
client := http.Client{}
|
||
|
|
|
||
|
|
req, err := http.NewRequest("GET", server.URL, nil)
|
||
|
|
require.Nil(t, err)
|
||
|
|
req.Header.Set("X-Forwarded-For", "3.3.3.3")
|
||
|
|
|
||
|
|
res, err := client.Do(req)
|
||
|
|
require.NoError(t, err)
|
||
|
|
defer res.Body.Close()
|
||
|
|
assert.Equal(t, http.StatusOK, res.StatusCode)
|
||
|
|
|
||
|
|
req, err = http.NewRequest("GET", server.URL, nil)
|
||
|
|
require.Nil(t, err)
|
||
|
|
req.Header.Set("X-Forwarded-For", "3.3.3.3, 4.4.4.4, 5.5.5.5")
|
||
|
|
|
||
|
|
res, err = client.Do(req)
|
||
|
|
require.NoError(t, err)
|
||
|
|
defer res.Body.Close()
|
||
|
|
assert.Equal(t, http.StatusOK, res.StatusCode)
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestRealIPBothHeaders(t *testing.T) {
|
||
|
|
server := httptest.NewServer(RealIP(realIPTestHandler(t)))
|
||
|
|
defer server.Close()
|
||
|
|
|
||
|
|
client := http.Client{}
|
||
|
|
|
||
|
|
req, err := http.NewRequest("GET", server.URL, nil)
|
||
|
|
require.Nil(t, err)
|
||
|
|
req.Header.Set("X-Real-IP", "3.3.3.3")
|
||
|
|
req.Header.Set("X-Forwarded-For", "4.4.4.4")
|
||
|
|
|
||
|
|
res, err := client.Do(req)
|
||
|
|
require.NoError(t, err)
|
||
|
|
defer res.Body.Close()
|
||
|
|
assert.Equal(t, http.StatusOK, res.StatusCode)
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestRealIPNoHeaders(t *testing.T) {
|
||
|
|
handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
|
|
assert.Equal(t, strings.Split(r.RemoteAddr, ":")[0], "127.0.0.1")
|
||
|
|
|
||
|
|
_, err := w.Write([]byte("resp"))
|
||
|
|
require.NoError(t, err)
|
||
|
|
})
|
||
|
|
|
||
|
|
server := httptest.NewServer(RealIP(handler))
|
||
|
|
defer server.Close()
|
||
|
|
|
||
|
|
res, err := http.Get(server.URL)
|
||
|
|
require.NoError(t, err)
|
||
|
|
defer res.Body.Close()
|
||
|
|
assert.Equal(t, http.StatusOK, res.StatusCode)
|
||
|
|
}
|