从外部 Go 中间件控制 HTTP 标头

Question

假设我在 Go 中有一个中间件，我想用Server我自己的值覆盖任何现有的标头。

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

然后我将它添加到这样的响应链中

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

不幸的是，如果 authHandler 或由 DefaultServeMux 调用处理的任何内容w.Header().Add("Server", "foo")（例如httputil.ReverseProxy.ServeHTTP），我最终会在响应中得到两个服务器标头。

$ 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

我真正想要的是这样的：

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

然而，ServeHTTP 的语义不允许这样做：

ServeHTTP 应该将回复标头和数据写入 ResponseWriter，然后返回。返回信号表明请求已完成；在 ServeHTTP 调用完成之后或同时，使用 ResponseWriter 或从 Request.Body 读取是无效的。

我见过一些围绕此问题的相当丑陋的黑客行为，例如httputil.ReverseProxy.ServeHTTP中的代码，它复制标头并重写响应代码，或者使用 httptest.ResponseRecorder 将整个正文读取到字节缓冲区中。

我还可以颠倒传统的中间件顺序，以某种方式将服务器中间件放在最后，或者使其成为最里面的中间件。

我缺少什么简单的方法吗？

Answer 1

您可以定义一个自定义类型，它包装 ResponseWriter 并在写入所有标头之前插入服务器标头，但需要额外的间接层。这是一个例子：

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

我在这里写了更多相关内容。https://kev.inburke.com/kevin/how-to-write-go-middleware/