54 lines
1.0 KiB
Go
54 lines
1.0 KiB
Go
package middleware
|
|
|
|
import (
|
|
"io"
|
|
"net/http"
|
|
|
|
"yh_web/server/pkg/traffic"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
type countReadCloser struct {
|
|
io.ReadCloser
|
|
}
|
|
|
|
func (c *countReadCloser) Read(p []byte) (int, error) {
|
|
n, err := c.ReadCloser.Read(p)
|
|
if n > 0 {
|
|
traffic.AddIn(n)
|
|
}
|
|
return n, err
|
|
}
|
|
|
|
type meterResponseWriter struct {
|
|
gin.ResponseWriter
|
|
}
|
|
|
|
func (w *meterResponseWriter) Write(p []byte) (int, error) {
|
|
n, err := w.ResponseWriter.Write(p)
|
|
if n > 0 {
|
|
traffic.AddOut(n)
|
|
}
|
|
return n, err
|
|
}
|
|
|
|
func (w *meterResponseWriter) WriteString(s string) (int, error) {
|
|
n, err := w.ResponseWriter.WriteString(s)
|
|
if n > 0 {
|
|
traffic.AddOut(n)
|
|
}
|
|
return n, err
|
|
}
|
|
|
|
// TrafficMeter 统计 HTTP 请求体与响应体字节量（进程级，非网卡级）。
|
|
func TrafficMeter() gin.HandlerFunc {
|
|
return func(c *gin.Context) {
|
|
if c.Request.Body != nil && c.Request.Body != http.NoBody {
|
|
c.Request.Body = &countReadCloser{ReadCloser: c.Request.Body}
|
|
}
|
|
c.Writer = &meterResponseWriter{ResponseWriter: c.Writer}
|
|
c.Next()
|
|
}
|
|
}
|