如何在反向代理中基于后端响应头设置Cookie？

这个问题我之前也碰到过，核心原因很明确：HTTP响应头一旦发送给客户端，就再也无法修改。你原来的代码里，调用next.ServeHTTP(w, r)之后，反向代理已经把后端的响应头和内容直接发给用户了，这时候再调用setCookie()自然不会生效。

要实现你的需求，我们得绕开这个限制——通过包装http.ResponseWriter，先把后端的响应全部拦截下来，等我们完成Cookie的设置后，再把完整的响应发送给客户端。下面是具体的解决方案：

步骤1：定义自定义ResponseWriter拦截响应 #

我们需要创建一个结构体，实现http.ResponseWriter接口，用来捕获后端返回的响应头、状态码和响应体，而不是直接发送给客户端：

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

// responseCapture 包装原始ResponseWriter，捕获响应信息
type responseCapture struct {
    http.ResponseWriter
    statusCode  int
    headers     http.Header
    body        bytes.Buffer
    wroteHeader bool
}

// newResponseCapture 创建一个新的responseCapture实例
func newResponseCapture(w http.ResponseWriter) *responseCapture {
    return &responseCapture{
        ResponseWriter: w,
        headers:        make(http.Header),
        statusCode:     http.StatusOK, // 默认状态码
    }
}

// Header 返回捕获的响应头，让后端写入这里而不是直接发送给客户端
func (rc *responseCapture) Header() http.Header {
    return rc.headers
}

// WriteHeader 捕获状态码，不直接发送给客户端
func (rc *responseCapture) WriteHeader(code int) {
    if !rc.wroteHeader {
        rc.statusCode = code
        rc.wroteHeader = true
    }
}

// Write 捕获响应体，写入内部缓冲区
func (rc *responseCapture) Write(b []byte) (int, error) {
    if !rc.wroteHeader {
        rc.WriteHeader(http.StatusOK) // 如果后端没写状态码，默认200
    }
    return rc.body.Write(b)
}

步骤2：修改中间件逻辑 #

在中间件中使用我们的responseCapture，先让反向代理把响应写入捕获器，然后根据捕获到的X-FOO头设置Cookie，最后再把完整的响应发送给客户端：

func setCookie(w http.ResponseWriter, name string, value string) {
    cookie := &http.Cookie{
        Name:  name,
        Value: value,
        Path:  "/", // 根据你的需求调整路径
        // 可选：添加安全属性
        // HttpOnly: true,
        // Secure:   true,
        // SameSite: http.SameSiteLaxMode,
    }
    // 直接操作Header，后续会统一发送给客户端
    w.Header().Add("Set-Cookie", cookie.String())
}

func handler(next http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        // 创建响应捕获器，代替原始的w
        rc := newResponseCapture(w)
        
        // 让反向代理处理请求，响应会被写入捕获器而不是客户端
        next.ServeHTTP(rc, r)
        
        // 现在可以安全地获取后端返回的X-FOO头，同时设置Cookie
        if r.Method == http.MethodPost && r.URL.String() == "/login" {
            if foo := rc.headers.Get("X-FOO"); foo != "" {
                setCookie(rc, "MYAPPFOO", foo)
            }
        }
        
        // 把捕获到的所有响应头复制到原始ResponseWriter
        for key, values := range rc.headers {
            for _, value := range values {
                w.Header().Add(key, value)
            }
        }
        
        // 发送状态码和响应体给客户端
        w.WriteHeader(rc.statusCode)
        rc.body.WriteTo(w)
    })
}

func main() {
    r := mux.NewRouter()
    r.Use(handler)
    proxy := httputil.NewSingleHostReverseProxy(&url.URL{Scheme: "https", Host: "baz.example.com"})
    r.PathPrefix("/").Handler(proxy)
    log.Fatal(http.ListenAndServe(":9001", r))
}

为什么这个方案有效？ #

• 我们的responseCapture拦截了所有后端的响应操作：后端写入的响应头会存在rc.headers里，响应体存在rc.body里，状态码存在rc.statusCode里，不会直接发送给客户端。
• 在完成反向代理的处理后，我们可以自由读取X-FOO头，并添加Set-Cookie头到rc.headers中。
• 最后我们统一把所有响应头、状态码和响应体发送给客户端，这时候Set-Cookie头会和其他响应头一起被发送，浏览器就能正常接收并保存Cookie了。

内容的提问来源于stack exchange，提问作者Coder1