From 3bfebfd79e0d5e09ce3537075f816c9b408d7363 Mon Sep 17 00:00:00 2001
From: lilee <laelli@foxmail.com>
Date: Thu, 10 Dec 2015 15:31:53 +0800
Subject: [PATCH] add custom replacer

---
 middleware/replacer.go      | 17 +++++++++++++++++
 middleware/replacer_test.go | 38 +++++++++++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+)

diff --git a/middleware/replacer.go b/middleware/replacer.go
index 29c695b7..192cec49 100644
--- a/middleware/replacer.go
+++ b/middleware/replacer.go
@@ -6,6 +6,7 @@ import (
 	"path"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 )
 
@@ -86,6 +87,11 @@ func NewReplacer(r *http.Request, rr *responseRecorder, emptyValue string) Repla
 		rep.replacements[headerReplacer+header+"}"] = strings.Join(val, ",")
 	}
 
+	customReplacementsMutex.RLock()
+	for k, v := range customReplacements {
+		rep.replacements["{@"+k+"}"] = v
+	}
+	customReplacementsMutex.RUnlock()
 	return rep
 }
 
@@ -117,3 +123,14 @@ const (
 	timeFormat     = "02/Jan/2006:15:04:05 -0700"
 	headerReplacer = "{>"
 )
+
+var (
+	customReplacements      = map[string]string{}
+	customReplacementsMutex sync.RWMutex
+)
+
+func RegisterReplacement(name, value string) {
+	customReplacementsMutex.Lock()
+	defer customReplacementsMutex.Unlock()
+	customReplacements[name] = value
+}
diff --git a/middleware/replacer_test.go b/middleware/replacer_test.go
index 39d91a22..cbbe7a23 100644
--- a/middleware/replacer_test.go
+++ b/middleware/replacer_test.go
@@ -69,3 +69,41 @@ func TestReplace(t *testing.T) {
 	}
 
 }
+
+func TestCustomReplace(t *testing.T) {
+	w := httptest.NewRecorder()
+	recordRequest := NewResponseRecorder(w)
+	userJSON := `{"username": "dennis"}`
+
+	reader := strings.NewReader(userJSON) //Convert string to reader
+
+	request, err := http.NewRequest("POST", "http://caddyserver.com", reader) //Create request with JSON body
+	if err != nil {
+		t.Fatalf("Request Formation Failed \n")
+	}
+	for i := 0; i < 2; i++ {
+		go func() {
+			RegisterReplacement("customId", "foobar")
+			replaceValues := NewReplacer(request, recordRequest, "")
+			switch v := replaceValues.(type) {
+			case replacer:
+
+				if v.Replace("This host is {host}") != "This host is caddyserver.com" {
+					t.Errorf("Expected host replacement failed")
+				}
+				if v.Replace("This request method is {method}") != "This request method is POST" {
+					t.Errorf("Expected method  replacement failed")
+				}
+				if v.Replace("The response status is {status}") != "The response status is 200" {
+					t.Errorf("Expected status replacement failed")
+				}
+				if v.Replace("The custom replace is customId with value of {@customId}") != "The custom replace is customId with value of foobar" {
+					t.Errorf("Expected customId replacement failed")
+				}
+
+			default:
+				t.Fatalf("Return Value from New Replacer expected pass type assertion into a replacer type   \n")
+			}
+		}()
+	}
+}