Merge pull request #7 from Thijmen/feature/add-original-header

Add Header with original URL

Author: Thijmen

traefikremovequeryparametersbyregex.go

@@ -5,7 +5,6 @@ package traefik_remove_query_parameters_by_regex
 import (
 	"context"
 	"errors"
-	"log"
 	"net/http"
 	"regexp"
 )
@@ -18,9 +17,10 @@ const (
 
 // Config is the configuration for this plugin.
 type Config struct {
-	Type               modificationType `json:"type"`
-	AllowedValuesRegex string           `json:"allowedValuesRegex"`
-	ExceptURIRegex     string           `json:"exceptUriRegex"`
+Type               				modificationType `json:"type"`
+	AllowedValuesRegex 			string           `json:"allowedValuesRegex"`
+	ExceptURIRegex     			string           `json:"exceptUriRegex"`
+	AddOriginalHostnameHeader	bool			 `json:"addOriginalHostnameHeader"`
 }
 
 // CreateConfig creates a new configuration for this plugin.
@@ -77,6 +77,8 @@ func New(ctx context.Context, next http.Handler, config *Config, name string) (h
 func (q *QueryParameterRemover) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 	qry := req.URL.Query()
 
+	originalQuery := req.URL.String()
+
 	switch q.config.Type {
 	case deleteExceptType:
 
@@ -92,13 +94,19 @@ func (q *QueryParameterRemover) ServeHTTP(rw http.ResponseWriter, req *http.Requ
 
 		regex := regexp.MustCompile(q.config.AllowedValuesRegex)
 
+		addOriginalHeader := false
+
 		for param := range req.URL.Query() {
 			if !regex.MatchString(param) {
 				qry.Del(param)
 				req.URL.RawQuery = qry.Encode()
-				log.Printf("Removed parameter: %s \n", param)
+				addOriginalHeader = true
 			}
 		}
+
+		if q.config.AddOriginalHostnameHeader && addOriginalHeader {
+			req.Header.Add("Plugin-Original-Uri", originalQuery)
+		}
 	}
 
 	req.URL.RawQuery = qry.Encode()

traefikremovequeryparametersbyregex_test.go

@@ -60,10 +60,21 @@ func TestDeleteQueryParamDoesNotWorkWithRegexWithADash(t *testing.T) {
 	expected := "x-live-preview=1"
 	previous := "test=1&x-live-preview=1"
 
-
 	assertQueryModificationHelper(t, cfg, previous, expected, "")
 }
 
+func TestItSendsTheHeader(t *testing.T) {
+	cfg := traefik_remove_query_parameters_by_regex.CreateConfig()
+	cfg.Type = "deleteexcept"
+	cfg.AllowedValuesRegex = "(testing|x-live-preview)"
+	cfg.ExceptURIRegex = "(qontrol)"
+	cfg.AddOriginalHostnameHeader = true
+	headerValue := "http://localhost?test=1&x-live-preview=1"
+	previous := "test=1&x-live-preview=1"
+
+	assertHeaderValue(t, cfg, previous, headerValue)
+}
+
 func TestErrorInvalidType(t *testing.T) {
 	cfg := traefik_remove_query_parameters_by_regex.CreateConfig()
 	cfg.Type = "bla"
@@ -117,3 +128,20 @@ func assertQueryModificationHelper(t *testing.T, cfg *traefik_remove_query_param
 		t.Errorf("Expected %s, got %s", expected, req.URL.Query().Encode())
 	}
 }
+
+func assertHeaderValue(t *testing.T, cfg *traefik_remove_query_parameters_by_regex.Config, previous, expectedHeaderValue string) {
+	t.Helper()
+	handler, recorder, req, err := createReqAndRecorder(cfg)
+	if err != nil {
+		t.Fatal(err)
+		return
+	}
+	req.URL.RawQuery = previous
+	handler.ServeHTTP(recorder, req)
+
+	header := req.Header.Get("Plugin-Original-Uri")
+
+	if header != expectedHeaderValue {
+		t.Errorf("Expected %s, got %s", expectedHeaderValue, header)
+	}
+}