重写中间件中的内容长度

Question

below code rewrites the http body response of some queries.

However, it fails to update the "content length" header field, it always remains the same original value.

How can i update the content length header field of the http response ?

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)
    })
}

I receive two errors message, the first one from nginx

2019/03/08 05:58:37 [error] 31194#0: *19 
upstream prematurely closed connection while reading upstream, 
client: 82.21.18.16, server: buycoffee.online, request: 
"GET / HTTP/1.1", upstream: "http://127.0.0.1:8081/", host: "buycoffee.online"

the second one from curl

curl: (18) transfer closed with 6237 bytes remaining to read

Answer 1

as mkopriva suggested the trick was to rewrite ResponseWriter.WriteHeader.

I believe it has to do with the fact that write might call writeheader at first call, and from there i was not putting the instructions at the right place.

in order to prevent further difficulties i prefer to use chunked transfer.

the code change is:

func (w *writeReplacer) WriteHeader(statusCode int) {
    w.Header().Del("Content-length")
    w.Header().Set("Transfer-Encoding", "chunked")
    w.ResponseWriter.WriteHeader(statusCode)
}