header: only register deletion operation (#1212)

fix issue #1183

Signed-off-by: Tw <tw19881113@gmail.com>
This commit is contained in:
Tw 2016-11-02 12:08:02 +08:00 committed by Matt Holt
parent 08c17c7c31
commit a2900e46f4
2 changed files with 13 additions and 18 deletions

View File

@ -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) {

View File

@ -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)