add http body for handler recorder

This commit is contained in:
ginuerzh
2024-09-16 19:44:15 +08:00
parent a618c5556e
commit e37213ac01
13 changed files with 292 additions and 31 deletions
+38 -3
View File
@@ -35,6 +35,10 @@ import (
"github.com/go-gost/x/registry"
)
const (
defaultBodySize = 1024 * 1024 * 10 // 10MB
)
func init() {
registry.HandlerRegistry().Register("tcp", NewHandler)
registry.HandlerRegistry().Register("udp", NewHandler)
@@ -375,6 +379,18 @@ func (h *forwardHandler) handleHTTP(ctx context.Context, rw io.ReadWriteCloser,
cc = tls.Client(cc, cfg)
}
var reqBody *xhttp.Body
if opts := h.recorder.Options; opts != nil && opts.HTTPBody {
if req.Body != nil {
maxSize := opts.MaxBodySize
if maxSize <= 0 {
maxSize = defaultBodySize
}
reqBody = xhttp.NewBody(req.Body, maxSize)
req.Body = reqBody
}
}
if err = req.Write(cc); err != nil {
cc.Close()
log.Warnf("send request to node %s(%s): %v", target.Name, target.Addr, err)
@@ -382,6 +398,11 @@ func (h *forwardHandler) handleHTTP(ctx context.Context, rw io.ReadWriteCloser,
return err
}
if reqBody != nil {
ro.HTTP.Request.Body = reqBody.Content()
ro.HTTP.Request.ContentLength = reqBody.Length()
}
if req.Header.Get("Upgrade") == "websocket" {
err = xnet.Transport(cc, xio.NewReadWriter(br, rw))
if err == nil {
@@ -395,16 +416,21 @@ func (h *forwardHandler) handleHTTP(ctx context.Context, rw io.ReadWriteCloser,
var err error
var res *http.Response
var respBody *xhttp.Body
defer func() {
ro.Duration = time.Since(start)
if err != nil {
ro.Err = err.Error()
}
if res != nil && ro.HTTP != nil {
ro.HTTP.Response.ContentLength = res.ContentLength
ro.HTTP.Response.Header = res.Header
if res != nil {
ro.HTTP.StatusCode = res.StatusCode
ro.HTTP.Response.Header = res.Header
ro.HTTP.Response.ContentLength = res.ContentLength
if respBody != nil {
ro.HTTP.Response.Body = respBody.Content()
ro.HTTP.Response.ContentLength = respBody.Length()
}
}
ro.Record(ctx, h.recorder.Recorder)
}()
@@ -432,6 +458,15 @@ func (h *forwardHandler) handleHTTP(ctx context.Context, rw io.ReadWriteCloser,
return
}
if opts := h.recorder.Options; opts != nil && opts.HTTPBody {
maxSize := opts.MaxBodySize
if maxSize <= 0 {
maxSize = defaultBodySize
}
respBody = xhttp.NewBody(res.Body, maxSize)
res.Body = respBody
}
if err = res.Write(rw); err != nil {
rw.Close()
log.Errorf("write response from node %s(%s): %v", target.Name, target.Addr, err)