From a2900e46f4531c6e1b9a4f151865798fdc596fcd Mon Sep 17 00:00:00 2001
From: Tw <tw19881113@gmail.com>
Date: Wed, 2 Nov 2016 12:08:02 +0800
Subject: [PATCH] header: only register deletion operation (#1212)

fix issue #1183

Signed-off-by: Tw <tw19881113@gmail.com>
---
 caddyhttp/header/header.go      | 25 ++++++++-----------------
 caddyhttp/header/header_test.go |  6 +++++-
 2 files changed, 13 insertions(+), 18 deletions(-)

diff --git a/caddyhttp/header/header.go b/caddyhttp/header/header.go
index dff925d4..7207c52e 100644
--- a/caddyhttp/header/header.go
+++ b/caddyhttp/header/header.go
@@ -32,9 +32,9 @@ func (h Headers) ServeHTTP(w http.ResponseWriter, r *http.Request) (int, error)
 				if strings.HasPrefix(header.Name, "-") {
 					rww.delHeader(strings.TrimLeft(header.Name, "-"))
 				} else if strings.HasPrefix(header.Name, "+") {
-					rww.addHeader(strings.TrimLeft(header.Name, "+"), replacer.Replace(header.Value))
+					rww.Header().Add(strings.TrimLeft(header.Name, "+"), replacer.Replace(header.Value))
 				} else {
-					rww.setHeader(header.Name, replacer.Replace(header.Value))
+					rww.Header().Set(header.Name, replacer.Replace(header.Value))
 				}
 			}
 		}
@@ -95,27 +95,18 @@ func (rww *responseWriterWrapper) WriteHeader(status int) {
 	rww.w.WriteHeader(status)
 }
 
-// addHeader registers a http.Header.Add operation
-func (rww *responseWriterWrapper) addHeader(key, value string) {
-	rww.ops = append(rww.ops, func(h http.Header) {
-		h.Add(key, value)
-	})
-}
-
-// delHeader registers a http.Header.Del operation
+// delHeader deletes the existing header according to the key
+// Also it will delete that header added later.
 func (rww *responseWriterWrapper) delHeader(key string) {
+	// remove the existing one if any
+	rww.Header().Del(key)
+
+	// register a future deletion
 	rww.ops = append(rww.ops, func(h http.Header) {
 		h.Del(key)
 	})
 }
 
-// setHeader registers a http.Header.Set operation
-func (rww *responseWriterWrapper) setHeader(key, value string) {
-	rww.ops = append(rww.ops, func(h http.Header) {
-		h.Set(key, value)
-	})
-}
-
 // Hijack implements http.Hijacker. It simply wraps the underlying
 // ResponseWriter's Hijack method if there is one, or returns an error.
 func (rww *responseWriterWrapper) Hijack() (net.Conn, *bufio.ReadWriter, error) {
diff --git a/caddyhttp/header/header_test.go b/caddyhttp/header/header_test.go
index 0e0aaa31..5c3b8304 100644
--- a/caddyhttp/header/header_test.go
+++ b/caddyhttp/header/header_test.go
@@ -25,6 +25,7 @@ func TestHeader(t *testing.T) {
 		{"/a", "Foo", "Bar"},
 		{"/a", "Bar", ""},
 		{"/a", "Baz", ""},
+		{"/a", "Server", ""},
 		{"/a", "ServerName", hostname},
 		{"/b", "Foo", ""},
 		{"/b", "Bar", "Removed in /a"},
@@ -32,7 +33,7 @@ func TestHeader(t *testing.T) {
 		he := Headers{
 			Next: httpserver.HandlerFunc(func(w http.ResponseWriter, r *http.Request) (int, error) {
 				w.Header().Set("Bar", "Removed in /a")
-				fmt.Fprint(w, "This is a test")
+				w.WriteHeader(http.StatusOK)
 				return 0, nil
 			}),
 			Rules: []Rule{
@@ -40,6 +41,7 @@ func TestHeader(t *testing.T) {
 					{Name: "Foo", Value: "Bar"},
 					{Name: "ServerName", Value: "{hostname}"},
 					{Name: "-Bar"},
+					{Name: "-Server"},
 				}},
 			},
 		}
@@ -50,6 +52,8 @@ func TestHeader(t *testing.T) {
 		}
 
 		rec := httptest.NewRecorder()
+		// preset header
+		rec.Header().Set("Server", "Caddy")
 
 		he.ServeHTTP(rec, req)