Hope to output replace data, but output origin data replace data or no output after trying

I want to rewrite response body in middleware,Only want to output replace data,Refer to How to rewrite response body in middleware?, modify my code as Requests are routed to

    ToolsGroup := Router.Group("")
    ToolsGroup .Use(middleware.ToolsGroupPermission())
    {
        ToolsGroup .GET("/ptr", func(c *gin.Context) {
            c.Data(http.StatusOK, "text/plain", []byte("orign data"))
        })
    }

Middleware is

package middleware

import (
    "bytes"
    "github.com/gin-gonic/gin"
)

func ToolsGroupPermission() gin.HandlerFunc {
    return func(c *gin.Context) {
        wb := &toolBodyWriter{
            body:           &bytes.Buffer{},
            ResponseWriter: c.Writer,
        }
        c.Writer = wb

        c.Next()

        wb.body.Reset()
        wb.Write([]byte("replace data"))
        //c.Data(http.StatusOK, "text/plain", []byte("replace data"))
    }
}

type toolBodyWriter struct {
    gin.ResponseWriter
    body *bytes.Buffer
}

func (r toolBodyWriter) Write(b []byte) (int, error) {
    return r.body.Write(b)
}

When the above code is modified, all [] bytes cannot be output. But the reference is changed to how do i get response body in after router middleware? when commenting:

func (r toolBodyWriter) Write(b []byte) (int, error) {
    r.body.Write(b)
    return r.ResponseWriter.Write(b)
}

Will output origin data replace data at the same time. But the requirement is to output replace data

CodePudding user response：

You cannot take back what you have written to r.ResponseWriter. Don't call ResponseWriter.Write before you know what you want to send to the client.

ToolsGroupPermission should either copy from the buffer to c.Writer, or call c.Data. There is no reason to even have the toolBodyWriter type; other handlers shouldn't have access to the original ResponseWriter if you plan to replace the response:

func ToolsGroupPermission() gin.HandlerFunc {
    return func(c *gin.Context) {
        buf := &bytes.Buffer{}
        rw := c.Writer        
        c.Writer = buf

        c.Next()

        var replaceResponse bool

        if replaceResponse {
            c.Data(http.StatusOK, "text/plain", []byte("replace data"))
        } else {
            // TODO: Copy header and status code somehow. I'm not familiar enough with gin to know how to do that nicely.
            io.Copy(rw, buf)
        }
    }
}

CodePudding user response：

package middleware

import (
    "bytes"
    "github.com/gin-gonic/gin"
)

func ToolsGroupPermission() gin.HandlerFunc {
    return func(c *gin.Context) {
        wb := &toolBodyWriter{
            body:           &bytes.Buffer{},
            ResponseWriter: c.Writer,
        }
        c.Writer = wb

        c.Next()

        originBytes := wb.body
        fmt.Printf("%s", originBytes)
        // clear Origin Buffer
        wb.body = &bytes.Buffer{}
        wb.Write([]byte("replace data"))
        wb.ResponseWriter.Write(wb.body.Bytes())
        //c.Data(http.StatusOK, "text/plain", []byte("replace data"))
    }
}

type toolBodyWriter struct {
    gin.ResponseWriter
    body *bytes.Buffer
}

func (r toolBodyWriter) Write(b []byte) (int, error) {
    return r.body.Write(b)
}

Or

package middleware

import (
    "bytes"
    "github.com/gin-gonic/gin"
)

func ToolsGroupPermission() gin.HandlerFunc {
    return func(c *gin.Context) {
        wb := &toolBodyWriter{
            body:           &bytes.Buffer{},
            ResponseWriter: c.Writer,
            status:         Origin,
        }
        c.Writer = wb

        c.Next()
        wb.status = Replace
        originBytes := wb.body
        fmt.Printf("%s", originBytes)
        // clear Origin Buffer
        wb.body = &bytes.Buffer{}
        wb.Write([]byte("replace data"))
        //c.Data(http.StatusOK, "text/plain", []byte("replace data"))
    }
}

type toolBodyWriter struct {
    gin.ResponseWriter
    body *bytes.Buffer
    status byte
}
const (
    Origin  byte = 0x0
    Replace      = 0x1
)

func (r toolBodyWriter) Write(b []byte) (int, error) {
    if r.status == 0x1 {
        r.body.Write(b)
        return r.ResponseWriter.Write(b)
    } else {
        return r.body.Write(b) //r.ResponseWriter.Write(b)
    }
}