Control HTTP headers from outer Go middleware

Question

Let's say I have a middleware in Go where I want to override any existing Server headers with my own value.

// Server attaches a Server header to the response.
func Server(h http.Handler, serverName string) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        w.Header().Set("Server", serverName)
        h.ServeHTTP(w, r)
    })
}

Then I add it to a response chain like this

http.Handle("/", authHandler)
http.ListenAndServe(":8000", Server(http.DefaultServeMux, "myapp"))

Unfortunately, if the authHandler or anything handled by DefaultServeMux calls w.Header().Add("Server", "foo") (like, say, httputil.ReverseProxy.ServeHTTP), I end up with two Server headers in the response.

$ http localhost:5962/v1/hello_world
HTTP/1.1 200 OK
Content-Length: 11
Content-Type: text/plain
Date: Tue, 12 Jul 2016 04:54:04 GMT
Server: inner-middleware
Server: myapp

What I really want is something like this:

// Server attaches a Server header to the response.
func Server(h http.Handler, serverName string) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        h.ServeHTTP(w, r)
        w.Header().Set("Server", serverName)
    })
}

However, this is disallowed by the semantics of ServeHTTP:

ServeHTTP should write reply headers and data to the ResponseWriter and then return. Returning signals that the request is finished; it is not valid to use the ResponseWriter or read from the Request.Body after or concurrently with the completion of the ServeHTTP call.

I have seen some fairly ugly hacks around this, for example the code in httputil.ReverseProxy.ServeHTTP, which copies headers and rewrites the response code, or using httptest.ResponseRecorder, which reads the entire body into a byte buffer.

I could also reverse the traditional middleware order and put my Server middleware last, somehow, or make it the inner-most middleware.

Is there any simple method I'm missing?

Answer 1

You can define a custom type which wraps a ResponseWriter and inserts the Server header just before all headers are written, at the cost of an additional layer of indirection. Here is an example:

type serverWriter struct {
    w           http.ResponseWriter
    name        string
    wroteHeader bool
}

func (s serverWriter) WriteHeader(code int) {
    if s.wroteHeader == false {
        s.w.Header().Set("Server", s.name)
        s.wroteHeader = true
    }
    s.w.WriteHeader(code)
}

func (s serverWriter) Write(b []byte) (int, error) {
    return s.w.Write(b)
}

func (s serverWriter) Header() http.Header {
    return s.w.Header()
}

// Server attaches a Server header to the response.
func Server(h http.Handler, serverName string) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        sw := serverWriter{
            w:           w,
            name:        serverName,
            wroteHeader: false,
        }
        h.ServeHTTP(sw, r)
    })
}

I wrote more about this here. https://kev.inburke.com/kevin/how-to-write-go-middleware/