使用go-chi中间件修改HTTP响应码时始终返回404求助

问题背景 #

使用go-chi开发API网关，需求是当请求的查询参数包含_always200=true时，强制将响应状态码设置为200，但当前实现后请求始终返回404。

现有代码问题分析 #

1. 自定义ResponseWriter逻辑缺陷：
   • Write方法仅将响应内容写入内部缓存Buf，未输出到原始的http.ResponseWriter，导致客户端无法收到响应内容。
   • 当后续处理器（如404处理器）调用WriteHeader(404)后，再调用WriteHeader(200)会被忽略——HTTP响应头只能发送一次，一旦已发送就无法修改。
2. 参数判断逻辑未生效：
   • 中间件中从r.Context().Value(always200QueryKey)获取always200值，但代码中未体现将查询参数解析后存入Context的逻辑，导致always200始终为false，强制200的逻辑从未触发。
3. 初始化默认值缺失：
   • NewCustomResponseWriter未初始化StatusCode和WroteHeader字段，后续状态码判断可能出现异常。

修复后的代码实现 #

1. 修正自定义ResponseWriter (custom_response_writer.go) #

package main

import (
	"bytes"
	"net/http"
)

type CustomResponseWriter struct {
	http.ResponseWriter
	Buf         *bytes.Buffer
	StatusCode  int
	wroteHeader bool
}

func NewCustomResponseWriter(w http.ResponseWriter) *CustomResponseWriter {
	return &CustomResponseWriter{
		ResponseWriter: w,
		Buf:            new(bytes.Buffer),
		StatusCode:     http.StatusOK, // 默认状态码设为200
		wroteHeader:    false,
	}
}

func (c *CustomResponseWriter) WriteHeader(code int) {
	if !c.wroteHeader {
		c.StatusCode = code
		c.wroteHeader = true
		// 暂不发送响应头，留到最后统一处理
	}
}

func (c *CustomResponseWriter) Write(b []byte) (int, error) {
	return c.Buf.Write(b)
}

// Flush 将缓存的响应内容和最终状态码发送给客户端
func (c *CustomResponseWriter) Flush() {
	if !c.wroteHeader {
		c.ResponseWriter.WriteHeader(c.StatusCode)
		c.wroteHeader = true
	}
	c.ResponseWriter.Write(c.Buf.Bytes())
}

2. 修正中间件逻辑 #

package main

import "net/http"

func HeaderFilterMiddleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// 直接解析查询参数，无需依赖Context
		always200 := r.URL.Query().Get("_always200") == "true"
		
		crw := NewCustomResponseWriter(w)
		next.ServeHTTP(crw, r)

		// 如果需要强制200，覆盖状态码
		if always200 {
			crw.StatusCode = http.StatusOK
		}

		// 统一发送响应头和内容
		crw.Flush()
	})
}

关键修复点说明 #

• 延迟发送响应头：自定义ResponseWriter的WriteHeader仅记录状态码，不立即发送，直到最后Flush时才统一发送，确保状态码可以被修改。
• 直接解析查询参数：中间件内直接读取查询参数判断_always200=true，避免依赖未实现的Context存储逻辑。
• 统一输出响应：通过Flush方法将缓存的响应内容和最终状态码发送给客户端，确保客户端能收到正确的响应。

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