使用自定义的http.ResponseWriter基于代理请求的响应编写cookie？

Question

我在这里原来的问题被标记为该问题的重复。我没有实现它的运气，并且怀疑我的问题被误解了，所以在我的问题结束时，我从一个更具体的问题开始着手。

我正在尝试基于中间代理中响应代理的响应标头设置cookie，该请求被反向代理。

这是工作流程：

• 用户请求http://example.com/foo/bar
• Go应用使用ReverseProxy将请求转发到http://baz.com
• baz.com设置响应头 X-FOO
• Go app通过设置MYAPPFOO带有X-FOO响应标头值的cookie来修改响应
• Cookie被写入用户的浏览器

有人建议，可以使用自定义http.ResponseWriter方式，但是在尝试并搜索更多信息后，尚不清楚如何实现此目的。

由于我无法掌握用例的自定义ResponseWriter的概念，因此我将发布代码，以更精确地演示我在遇到困难时试图做的事情：

package main

import (

    "github.com/gorilla/mux"
    "log"
    "net/http"
    "net/http/httputil"
    "net/url"
)

func setCookie(w http.ResponseWriter, name string, value string) {
    ...
    http.SetCookie(w, &cookie)
}

func handler(next http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {

        // setCookie() works here
        // but I cannot access w.Header().Get("X-FOO")

        next.ServeHTTP(w, r)

        // I can access w.Header().Get("X-FOO") here
        // but setCookie() does not cookie the user's browser

        // If I could do it all in one place, this is what I would do:
        if r.Method == "POST" && r.URL.String() == "/login" {
            foo := w.Header().Get("X-FOO")
            setCookie(w, "MYAPPFOO", foo)

        }
    })
}

func main() {

    r := mux.NewRouter()
    r.Use(handler)
    proxy := httputil.NewSingleHostReverseProxy("https://baz.example.com/")
    r.PathPrefix("/").Handler(proxy)
    log.Fatal(http.ListenAndServe(":9001", r))
}

附带说明一下，我能够按照ReverseProxy.ModifyResponse上一个问题的注释中的建议进行这项工作，但我真的很想用中间件来实现这一点，以保持从配置中动态创建代理的代码干净。（不在示例代码中）

Answer 1

从有关http.ResponseWriter方法的文档中：（
添加了重点）

• Header() http.Header：

除非修改后的标头是预告片，否则在 调用WriteHeader（或Write）后 更改标头映射 无效 。

• WriteHeader(statusCode int)：

WriteHeader 发送 带有提供的状态代码 的HTTP响应标头 。

• Write([]byte) (int, error)：

如果尚未 调用WriteHeader ，则Write 在写入数据之前会 调用WriteHeader（http.StatusOK） 。

这应该强调为什么，你不能设置在后一个cookie的原因next.ServeHTTP(w, r)调用，这是在通过电话执行的中间件链中的处理程序中的一个或者打电话WriteHeader或Write直接或间接的。

因此，要能够在调用 之后 设置cookie ，next.ServeHTTP(w, r)您需要确保中间件链调用WriteHeader或Write原始http.ResponseWriter实例中没有任何处理程序。一种方法是将原始实例包装在自定义http.ResponseWriter实现中，该实现将
推迟 响应的编写，直到完成cookie设置为止。

例如这样的事情：

type responsewriter struct {
    w    http.ResponseWriter
    buf  bytes.Buffer
    code int
}

func (rw *responsewriter) Header() http.Header {
    return rw.w.Header()
}

func (rw *responsewriter) WriteHeader(statusCode int) {
    rw.code = statusCode
}

func (rw *responsewriter) Write(data []byte) (int, error) {
    return rw.buf.Write(data)
}

func (rw *responsewriter) Done() (int64, error) {
    if rw.code > 0 {
        rw.w.WriteHeader(rw.code)
    }
    return io.Copy(rw.w, &rw.buf)
}

您可以在中间件中像这样使用它：

func handler(next http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        rw := &responsewriter{w: w}
        next.ServeHTTP(rw, r)

        if r.Method == "POST" && r.URL.String() == "/login" {
            foo := rw.Header().Get("X-FOO")
            setCookie(rw, "MYAPPFOO", foo)
        }

        if _, err := rw.Done(); err != nil {
            log.Println(err)
        }
    })
}