重写中间件中的内容长度

时间:2019-03-08 05:50:57

标签: go

下面的代码重写了某些查询的http正文响应。

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

如何更新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)
    })
}

我收到两个错误消息,第一个错误消息来自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"

卷曲的第二个

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

1 个答案:

答案 0 :(得分:0)

正如mkopriva所建议的那样,窍门是重写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)
}