middleware/real_ip_test.go

package middleware

import (
	"io"
	"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, http.NoBody)
	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 := io.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, http.NoBody)
	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, http.NoBody)
	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, http.NoBody)
	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)
}

func TestRealIPContext(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")
		assert.Equal(t, r.Context().Value(RealIPCtxKey), "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)
}