使用自定义 http.ResponseWriter 根据代理请求的响应写入 cookie？

我原来的问题here https://stackoverflow.com/questions/58776775/how-can-i-set-a-cookie-based-on-response-header-with-reverse-proxy?noredirect=1#comment103840087_58776775被标记为重复项this https://stackoverflow.com/questions/38334072/control-http-headers-from-outer-go-middleware问题。我没有运气实现它，并怀疑我的问题被误解了，所以在我的问题结束后，我开始提出一个更具体的问题。

我正在尝试根据反向代理请求中的中间件内的响应标头设置 cookie。

这是工作流程：

• 用户请求http://example.com/foo/bar http://example.com/foo/bar
• Go 应用程序使用 ReverseProxy 将该请求代理到http://baz.com http://baz.com
• baz.com 设置响应标头X-FOO
• Go 应用程序通过设置来修改响应MYAPPFOOcookie 的值X-FOO响应头
• 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正如我上一个问题的评论中所建议的，但我真的很想通过中间件来实现这一点，以保持从配置动态创建代理的代码干净。 （示例代码中没有）

从文档上http.ResponseWriter https://golang.org/pkg/net/http/#ResponseWriter方法：（强调已添加）

• Header() http.Header:

  更改标题图调用 WriteHeader （或 Write）后没有 影响除非修改后的标头是预告片。

• WriteHeader(statusCode int):

  写头发送 HTTP 响应标头与提供的状态 代码。

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

  如果 WriteHeader 尚未被调用，则 Write来电 WriteHeader(http.StatusOK)在写入数据之前。

这应该突出说明为什么您无法在之后设置 cookienext.ServeHTTP(w, r)调用，这是该调用执行的中间件链中的处理程序之一正在调用WriteHeader or Write直接或间接。

这样才能设置cookieafter the next.ServeHTTP(w, r)调用您需要确保中间件链中没有任何处理程序调用WriteHeader or Write在原来的http.ResponseWriter实例。执行此操作的一种方法是将原始实例包装在自定义实例中http.ResponseWriter实施将postpone直到您完成 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)
        }
    })
}