如何在Go反向代理中间件中用自定义http.ResponseWriter设置Cookie？

我明白你的需求——不想在每个代理实例里写ModifyResponse逻辑，而是通过统一的中间件来处理，这样动态生成代理的代码能保持简洁。你之前遇到的瓶颈核心问题是：当next.ServeHTTP(w, r)执行完成后，标准http.ResponseWriter的响应头已经被反向代理提交给客户端了，这时候再修改Header设置Cookie，客户端根本收不到。

要解决这个问题，我们需要自定义一个http.ResponseWriter包装器，提前拦截并保存反向代理的响应头，等我们添加完Cookie后，再把完整的响应头发送给客户端。

完整实现代码 #

package main

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

// responseCapture 包装原生ResponseWriter，用于拦截并保存响应头
type responseCapture struct {
	http.ResponseWriter
	statusCode int
	header     http.Header
}

func newResponseCapture(w http.ResponseWriter) *responseCapture {
	return &responseCapture{
		ResponseWriter: w,
		statusCode:     http.StatusOK,
		header:         make(http.Header),
	}
}

// Header 返回我们自定义的header，用于拦截响应头的写入
func (c *responseCapture) Header() http.Header {
	return c.header
}

// WriteHeader 保存状态码，但暂不发送给客户端
func (c *responseCapture) WriteHeader(code int) {
	c.statusCode = code
}

// Write 先处理响应头（包括我们添加的Cookie），再写入响应体
func (c *responseCapture) Write(b []byte) (int, error) {
	if c.statusCode != 0 {
		// 将自定义header中的所有内容复制到原生ResponseWriter的header中
		for k, v := range c.header {
			c.ResponseWriter.Header()[k] = v
		}
		// 发送状态码和响应头
		c.ResponseWriter.WriteHeader(c.statusCode)
		c.statusCode = 0 // 标记响应头已发送，避免重复执行
	}
	// 写入响应体
	return c.ResponseWriter.Write(b)
}

func handler(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// 创建响应捕获器
		capture := newResponseCapture(w)
		// 让代理处理请求，响应头会被捕获到capture中
		next.ServeHTTP(capture, r)

		// 按照你的业务逻辑判断是否需要设置Cookie
		if r.Method == http.MethodPost && r.URL.Path == "/login" {
			if fooVal := capture.Header().Get("X-FOO"); fooVal != "" {
				cookie := &http.Cookie{
					Name:     "MYAPPFOO",
					Value:    fooVal,
					Path:     "/",
					HttpOnly: true, // 推荐开启，提高安全性
					Secure:   r.TLS != nil, // 仅在HTTPS请求下设置Secure属性
					SameSite: http.SameSiteLaxMode,
				}
				// 将Cookie添加到捕获的响应头中
				http.SetCookie(capture, cookie)
			}
		}

		// 处理反向代理没有调用Write()的情况（比如只返回状态码没有响应体）
		if capture.statusCode != 0 {
			for k, v := range capture.header {
				w.Header()[k] = v
			}
			w.WriteHeader(capture.statusCode)
		}
	})
}

func main() {
	r := mux.NewRouter()
	r.Use(handler)

	targetURL, _ := url.Parse("https://baz.example.com")
	proxy := httputil.NewSingleHostReverseProxy(targetURL)
	r.PathPrefix("/").Handler(proxy)

	log.Fatal(http.ListenAndServe(":9001", r))
}

实现原理说明 #

1. 自定义ResponseWriter拦截响应：responseCapture替代了原生的ResponseWriter，所有响应头的写入都会先存在它的header属性里，不会直接发送给客户端。
2. 延迟发送响应头：只有当反向代理调用Write()（写入响应体）或者我们手动触发时，才会把包含新Cookie的完整响应头发送给客户端。
3. 中间件后续处理：在next.ServeHTTP执行完成后，我们可以安全地从capture.Header()获取X-FOO的值，添加Cookie后，再通过responseCapture的逻辑把最终响应发送给用户。

注意事项 #

• 根据你的业务场景调整Cookie的属性（比如Domain、MaxAge等），确保符合安全要求。
• 如果反向代理返回多个X-FOO头，使用capture.Header().Values("X-FOO")获取所有值。
• 如果你的代理需要支持流式响应（比如服务器推送），需要让responseCapture实现http.Flusher接口，把Flush()方法转发给原生ResponseWriter。

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