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

Question

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

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

这是工作流程：

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

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