在中间件中重写内容长度

Question

下面的代码重写了一些查询的 http body 响应。

但是，它无法更新“内容长度”标头字段，它始终保持相同的原始值。

如何更新 http 响应的内容长度标头字段？

type writeReplacer struct {

    http.ResponseWriter

    search  []byte

    replace func(*http.Request) string

    buf     []byte

    r       *http.Request

    dir     string

}

func (w *writeReplacer) Write(in []byte) (int, error) {

    if w.buf == nil {

        w.buf = []byte{}

    }

    w.buf = append(w.buf, in...)

    n := len(in)

    if index := bytes.LastIndex(w.buf, w.search); index > -1 {

        var r []byte

        if w.dir == "before" {

            g := []byte(w.replace(w.r))

            n += len(g)

            r = append(g, w.buf[index:]...)

            w.buf = append(w.buf[:index], r...)

        } else {

            g := []byte(w.replace(w.r))

            n += len(g)

            r = append(r, w.buf[:index+len(w.search)]...)

            r = append(r, g...)

            r = append(r, w.buf[index:]...)

            w.buf = r

        }

    }

    return n, nil

}

func (w *writeReplacer) Flush() {

    w.ResponseWriter.Header().Set("Content-Length", fmt.Sprint(len(w.buf)))

    w.ResponseWriter.Write(w.buf[:])

    w.buf = w.buf[:0]

}

func InsertAfter(h http.Handler, path string, search []byte, replace func(*http.Request) string) http.Handler {

    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {

        if r.URL.Path == path {

            w.Header().Del("Content-length")

            w = &writeReplacer{ResponseWriter: w, search: search, replace: replace, r: r, dir: "after"}

            defer w.(http.Flusher).Flush()

        }

        h.ServeHTTP(w, r)

    })

}

func InsertBefore(h http.Handler, path string, search []byte, replace func(*http.Request) string) http.Handler {

    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {

        if r.URL.Path == path {

            w.Header().Del("Content-length")

            w = &writeReplacer{ResponseWriter: w, search: search, replace: replace, r: r, dir: "before"}

            defer w.(http.Flusher).Flush()

        }

        h.ServeHTTP(w, r)

    })

}

Answer 1

诀窍是重写ResponseWriter.WriteHeader.我相信这与 write 可能会在第一次调用时调用 writeheader 这一事实有关，从那里我没有将指令放在正确的位置。为了防止进一步的困难，我更喜欢使用分块传输。代码更改是：func (w *writeReplacer) WriteHeader(statusCode int) {    w.Header().Del("Content-length")    w.Header().Set("Transfer-Encoding", "chunked")    w.ResponseWriter.WriteHeader(statusCode)}