test: add test suite

Author: acouvreur

go.mod

@@ -1,3 +1,11 @@
 module github.com/acouvreur/traefik-modsecurity-plugin
 
 go 1.17
+
+require github.com/stretchr/testify v1.7.0
+
+require (
+	github.com/davecgh/go-spew v1.1.0 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
+	gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c // indirect
+)

go.sum

@@ -0,0 +1,11 @@
+github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8=
+github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
+github.com/stretchr/testify v1.7.0 h1:nwc3DEeHmmLAfoZucVR881uASk0Mfjw8xYJ99tb5CcY=
+github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c h1:dUUwHk2QECo/6vqA44rthZ8ie2QXMNeKRTHCNY2nXvo=
+gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

modsecurity.go

@@ -85,16 +85,7 @@ func (a *Modsecurity) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 	defer resp.Body.Close()
 
 	if resp.StatusCode >= 400 {
-		// copy headers
-		for k, vv := range resp.Header {
-			for _, v := range vv {
-				rw.Header().Set(k, v)
-			}
-		}
-		// copy status
-		rw.WriteHeader(resp.StatusCode)
-		// copy body
-		io.Copy(rw, resp.Body)
+		forwardResponse(resp, rw)
 		return
 	}
 
@@ -109,3 +100,16 @@ func isWebsocket(req *http.Request) bool {
 	}
 	return false
 }
+
+func forwardResponse(resp *http.Response, rw http.ResponseWriter) {
+	// copy headers
+	for k, vv := range resp.Header {
+		for _, v := range vv {
+			rw.Header().Set(k, v)
+		}
+	}
+	// copy status
+	rw.WriteHeader(resp.StatusCode)
+	// copy body
+	io.Copy(rw, resp.Body)
+}

modsecurity_test.go

@@ -0,0 +1,118 @@
+// Package traefik_modsecurity_plugin a modsecurity plugin.
+package traefik_modsecurity_plugin
+
+import (
+	"bytes"
+	"io"
+	"log"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestModsecurity_ServeHTTP(t *testing.T) {
+
+	req, err := http.NewRequest(http.MethodGet, "http://proxy.com/test", bytes.NewBuffer([]byte("Request")))
+
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	type response struct {
+		Body       string
+		StatusCode int
+	}
+
+	serviceResponse := response{
+		StatusCode: 200,
+		Body:       "Response from service",
+	}
+
+	tests := []struct {
+		name            string
+		request         http.Request
+		wafResponse     response
+		serviceResponse response
+		expectBody      string
+		expectStatus    int
+	}{
+		{
+			name:    "Forward request when WAF found no threats",
+			request: *req,
+			wafResponse: response{
+				StatusCode: 200,
+				Body:       "Response from waf",
+			},
+			serviceResponse: serviceResponse,
+			expectBody:      "Response from service",
+			expectStatus:    200,
+		},
+		{
+			name:    "Intercepts request when WAF found threats",
+			request: *req,
+			wafResponse: response{
+				StatusCode: 403,
+				Body:       "Response from waf",
+			},
+			serviceResponse: serviceResponse,
+			expectBody:      "Response from waf",
+			expectStatus:    403,
+		},
+		{
+			name: "Does not forward Websockets",
+			request: http.Request{
+				Body: http.NoBody,
+				Header: http.Header{
+					"Upgrade": []string{"Websocket"},
+				},
+			},
+			wafResponse: response{
+				StatusCode: 200,
+				Body:       "Response from waf",
+			},
+			serviceResponse: serviceResponse,
+			expectBody:      "Response from service",
+			expectStatus:    200,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+
+			modsecurityMockServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				resp := http.Response{
+					Body:       io.NopCloser(bytes.NewReader([]byte(tt.wafResponse.Body))),
+					StatusCode: tt.wafResponse.StatusCode,
+					Header:     http.Header{},
+				}
+				forwardResponse(&resp, w)
+			}))
+
+			httpServiceHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				resp := http.Response{
+					Body:       io.NopCloser(bytes.NewReader([]byte(tt.serviceResponse.Body))),
+					StatusCode: tt.serviceResponse.StatusCode,
+					Header:     http.Header{},
+				}
+				forwardResponse(&resp, w)
+			})
+
+			middleware := &Modsecurity{
+				next:           httpServiceHandler,
+				modSecurityUrl: modsecurityMockServer.URL,
+				name:           "modsecurity-middleware",
+			}
+
+			rw := httptest.NewRecorder()
+
+			middleware.ServeHTTP(rw, &tt.request)
+
+			resp := rw.Result()
+			body, _ := io.ReadAll(resp.Body)
+
+			assert.Equal(t, tt.expectBody, string(body))
+			assert.Equal(t, tt.expectStatus, resp.StatusCode)
+		})
+	}
+}