feat(middleware): optimize HTML modification with byte pool

This commit is contained in:
yusing 2025-07-20 12:47:32 +08:00
parent f39513483b
commit 46c7ee4d84

View file

@ -9,6 +9,7 @@ import (
"github.com/PuerkitoBio/goquery" "github.com/PuerkitoBio/goquery"
"github.com/rs/zerolog/log" "github.com/rs/zerolog/log"
gphttp "github.com/yusing/go-proxy/internal/net/gphttp" gphttp "github.com/yusing/go-proxy/internal/net/gphttp"
"github.com/yusing/go-proxy/internal/utils/synk"
"golang.org/x/net/html" "golang.org/x/net/html"
) )
@ -20,7 +21,9 @@ type modifyHTML struct {
var ModifyHTML = NewMiddleware[modifyHTML]() var ModifyHTML = NewMiddleware[modifyHTML]()
func (m *modifyHTML) before(w http.ResponseWriter, req *http.Request) bool { var bytePool = synk.NewBytesPool()
func (m *modifyHTML) before(_ http.ResponseWriter, req *http.Request) bool {
req.Header.Set("Accept-Encoding", "") req.Header.Set("Accept-Encoding", "")
return true return true
} }
@ -75,7 +78,7 @@ func (m *modifyHTML) modifyResponse(resp *http.Response) error {
// copied and modified from (*goquery.Selection).Html() // copied and modified from (*goquery.Selection).Html()
func buildHTML(s *goquery.Document) (ret []byte, err error) { func buildHTML(s *goquery.Document) (ret []byte, err error) {
var buf bytes.Buffer buf := bytes.NewBuffer(bytePool.Get())
// Merge all head nodes into one // Merge all head nodes into one
headNodes := s.Find("head") headNodes := s.Find("head")
@ -97,7 +100,7 @@ func buildHTML(s *goquery.Document) (ret []byte, err error) {
if len(s.Nodes) > 0 { if len(s.Nodes) > 0 {
for c := s.Nodes[0].FirstChild; c != nil; c = c.NextSibling { for c := s.Nodes[0].FirstChild; c != nil; c = c.NextSibling {
err = html.Render(&buf, c) err = html.Render(buf, c)
if err != nil { if err != nil {
return return
} }