Traefik中间件开发问题：POST请求异常及拦截需求

问题分析 #

你的代码存在几个核心问题，导致POST请求响应异常、返回500，以及拦截逻辑不符合预期：

• 重复发送响应头：captureResponseWriter.Write方法中主动调用w.rw.WriteHeader(w.StatusCode)，如果源站已调用过WriteHeader，重复调用会触发HTTP协议错误，直接导致连接关闭。
• 请求重复处理：ServeHTTP末尾调用h.next.ServeHTTP(rw, req)，但h.source.ServeHTTP已经完成了请求处理，重复调用会让客户端收到重复响应，引发连接异常。
• 拦截头信息混乱：拦截时直接修改原响应头，但未重置之前复制的源站头，导致403响应混有源站的头信息。
• 状态码未正确转发：WriteHeader方法注释掉了w.rw.WriteHeader(statusCode)，源站的状态码无法正常传递到客户端。

修复后的代码 #

1. 修正captureResponseWriter实现 #

type captureResponseWriter struct {
    Headers    http.Header
    StatusCode int
    Buffer     bytes.Buffer
    rw         http.ResponseWriter
    headerSent bool // 标记响应头是否已发送，避免重复发送
}

func NewCapture(rw http.ResponseWriter) *captureResponseWriter {
    return &captureResponseWriter{
        rw:         rw,
        Headers:    http.Header{},
        StatusCode: http.StatusOK,
        Buffer:     bytes.Buffer{},
        headerSent: false,
    }
}

func (w *captureResponseWriter) Header() http.Header {
    // 返回自身维护的Headers，避免直接修改原响应头
    return w.Headers
}

func (w *captureResponseWriter) Write(b []byte) (int, error) {
    if !w.headerSent {
        // 第一次写入内容时自动发送状态码
        w.WriteHeader(w.StatusCode)
    }
    return w.Buffer.Write(b)
}

func (w *captureResponseWriter) WriteHeader(statusCode int) {
    if w.headerSent {
        return // 已发送过头，直接返回
    }
    w.StatusCode = statusCode
    w.headerSent = true
    // 克隆当前的头信息（此时为源站设置的头）
    w.Headers = w.Headers.Clone()
}

func (w *captureResponseWriter) Flush() {
    if f, ok := w.rw.(http.Flusher); ok {
        f.Flush()
    }
}

func (w *captureResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
    if h, ok := w.rw.(http.Hijacker); ok {
        return h.Hijack()
    }
    return nil, nil, fmt.Errorf("not a hijacker: %T", w.rw)
}

2. 修正ServeHTTP逻辑 #

func (h *accessControlHandler) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
    capture := NewCapture(rw)
    // 调用源站处理请求，捕获响应内容、头和状态码
    h.source.ServeHTTP(capture, req)

    // 遍历检查拦截规则
    for _, rule := range h.config.Rules {
        r := NewRule(rule.Expr)
        matchRule := r.Run()
        if !matchRule {
            continue
        }

        h.logger.Infof("match rule: %s, with action %s", r.RawExpr, r.Action)
        if r.Action == "reject" {
            // 拦截时重置响应头，确保403响应的头信息干净
            rw.Header().Reset()
            content := []byte("Block by AccessControl!")
            rw.Header().Set("Content-Length", fmt.Sprintf("%d", len(content)))
            rw.WriteHeader(http.StatusForbidden)
            _, _ = rw.Write(content)
            
            // 异步发送攻击日志
            go requestlog.SendAttackLog(req, "精准访问控制", makeAttackLogInfo(host, h.domainHost, ForwardedURL, "拦截", r.RawExpr, ""), h.domainID, h.localaddr, h.domainHost)
            return
        }
    }

    // 未触发拦截，将捕获的源站响应转发给客户端
    // 复制源站响应头
    for k, v := range capture.Headers {
        rw.Header()[k] = v
    }
    // 发送状态码
    rw.WriteHeader(capture.StatusCode)
    // 发送响应内容
    _, _ = rw.Write(capture.Buffer.Bytes())
}

关键修复点说明 #

• 独立维护响应头：captureResponseWriter现在维护自身的Headers，避免直接修改原响应头，确保拦截时可以完全重置头信息。
• 避免重复发送头：新增headerSent标记，防止WriteHeader和Write方法重复触发响应头发送，规避HTTP协议错误。
• 取消重复请求处理：移除了末尾的h.next.ServeHTTP(rw, req)，改为手动转发捕获的源站响应，解决重复处理请求导致的连接异常。
• 拦截时的头重置：触发拦截时先调用rw.Header().Reset()，确保403响应的头信息不混有源站内容。

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