github.com/in4it/ecs-deploy@v0.0.42-0.20240508120354-ed77ff16df25/ipfilter/main_test.go (about)

     1  package ipfilter
     2  
     3  import (
     4  	"net/http"
     5  	"net/http/httptest"
     6  	"testing"
     7  
     8  	"github.com/gin-gonic/gin"
     9  	"github.com/stretchr/testify/assert"
    10  )
    11  
    12  func setupMiddleware(whitelist string) *gin.Engine {
    13  	r := gin.Default()
    14  	r.Use(IPWhiteList(whitelist))
    15  	r.GET("/ping", func(c *gin.Context) {
    16  		c.String(200, "pong")
    17  	})
    18  	return r
    19  }
    20  
    21  func TestPingRoute(t *testing.T) {
    22  	router := setupMiddleware("0.0.0.0/0")
    23  
    24  	w := httptest.NewRecorder()
    25  	req, _ := http.NewRequest("GET", "/ping", nil)
    26  	req.Header.Add("X-Forwarded-For", "127.0.0.1")
    27  	router.ServeHTTP(w, req)
    28  
    29  	assert.Equal(t, 200, w.Code)
    30  	assert.Equal(t, "pong", w.Body.String())
    31  }
    32  
    33  func TestProcessingInput(t *testing.T) {
    34  	router := setupMiddleware("10.10.10.0/24,20.20.20.0/24,30.30.30.0/24,0.0.0.0/0")
    35  	w := httptest.NewRecorder()
    36  	req, _ := http.NewRequest("GET", "/ping", nil)
    37  	req.Header.Add("X-Forwarded-For", "20.20.20.5")
    38  	router.ServeHTTP(w, req)
    39  
    40  	assert.Equal(t, 200, w.Code)
    41  	assert.Equal(t, "pong", w.Body.String())
    42  }
    43  
    44  func TestDenyRoute(t *testing.T) {
    45  	router := setupMiddleware("10.10.10.0/24")
    46  
    47  	w := httptest.NewRecorder()
    48  	req, _ := http.NewRequest("GET", "/ping", nil)
    49  	req.Header.Add("X-Forwarded-For", "127.0.0.1")
    50  	router.ServeHTTP(w, req)
    51  
    52  	assert.Equal(t, 403, w.Code)
    53  	assert.Equal(t, "{\"message\":\"Permission denied\",\"status\":403}", w.Body.String())
    54  }
    55  
    56  func TestBadClientIPRoute(t *testing.T) {
    57  	router := setupMiddleware("10.10.10.0/24")
    58  
    59  	w := httptest.NewRecorder()
    60  	req, _ := http.NewRequest("GET", "/ping", nil)
    61  	req.Header.Add("X-Forwarded-For", "10.10.10.1badinput")
    62  	router.ServeHTTP(w, req)
    63  
    64  	assert.Equal(t, 403, w.Code)
    65  	assert.Equal(t, "{\"message\":\"Permission denied\",\"status\":403}", w.Body.String())
    66  }
    67  
    68  func TestBadInput(t *testing.T) {
    69  	router := setupMiddleware("0.0.0.0/0badinput")
    70  	w := httptest.NewRecorder()
    71  	req, _ := http.NewRequest("GET", "/ping", nil)
    72  	req.Header.Add("X-Forwarded-For", "127.0.0.1")
    73  	router.ServeHTTP(w, req)
    74  
    75  	assert.Equal(t, 403, w.Code)
    76  	assert.Equal(t, "{\"message\":\"Permission denied\",\"status\":403}", w.Body.String())
    77  }