How Can Go Middleware Reliably Control HTTP Headers Without Duplication?

Controlling HTTP Headers from Outer Go Middleware

HTTP middleware in Go provides a convenient way to intercept and modify HTTP requests and responses. However, controlling headers from an outer middleware can be challenging, as it requires overriding existing headers without introducing duplicates.

Consider the following server middleware that sets the "Server" header:

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

When added to the response chain, this middleware successfully sets the "Server" header. However, if any other handler in the chain also sets the "Server" header, duplicate headers will result in the response.

The challenge arises because ServeHTTP explicitly prohibits writing to the ResponseWriter after the request is finished. One approach is to create a custom ResponseWriter that intercepts header writes and inserts the "Server" header before the first write.

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

By using a custom ResponseWriter, we can ensure that the "Server" header is added only once, regardless of the behavior of other handlers. This approach introduces an additional layer of indirection, but it maintains the desired functionality.

The above is the detailed content of How Can Go Middleware Reliably Control HTTP Headers Without Duplication?. For more information, please follow other related articles on the PHP Chinese website!