fix(middleware): correct and simplify HTML modification / buffer management logic, correct Accept-Encoding header

This commit is contained in:
yusing
2025-10-27 15:08:29 +08:00
parent 39c8cc2820
commit 5e00e1c437
2 changed files with 10 additions and 19 deletions

View File

@@ -15,20 +15,15 @@ import (
)
type modifyHTML struct {
Target string // css selector
HTML string // html to inject
Replace bool // replace the target element with the new html instead of appending it
bytesPool synk.UnsizedBytesPool
Target string // css selector
HTML string // html to inject
Replace bool // replace the target element with the new html instead of appending it
}
var ModifyHTML = NewMiddleware[modifyHTML]()
func (m *modifyHTML) setup() {
m.bytesPool = synk.GetUnsizedBytesPool()
}
func (m *modifyHTML) before(_ http.ResponseWriter, req *http.Request) bool {
req.Header.Set("Accept-Encoding", "")
req.Header.Set("Accept-Encoding", "identity")
return true
}
@@ -55,7 +50,6 @@ func (m *modifyHTML) modifyResponse(resp *http.Response) error {
resp.Body.Close()
if err != nil {
log.Err(err).Str("url", fullURL(resp.Request)).Msg("failed to read response body")
release(content)
resp.Body = eofReader{}
return err
}
@@ -83,23 +77,24 @@ func (m *modifyHTML) modifyResponse(resp *http.Response) error {
ele.First().AppendHtml(m.HTML)
}
// should not use content (from sized pool) directly for bytes.Buffer
buf := m.bytesPool.GetBuffer()
buf.Write(content)
release(content)
pool := synk.GetUnsizedBytesPool()
buf := pool.GetBuffer()
err = buildHTML(doc, buf)
if err != nil {
pool.PutBuffer(buf)
log.Err(err).Str("url", fullURL(resp.Request)).Msg("failed to build html")
// invalid html, restore the original body
resp.Body = readerWithRelease(content, release)
return err
}
release(content)
resp.ContentLength = int64(buf.Len())
resp.Header.Set("Content-Length", strconv.Itoa(buf.Len()))
resp.Header.Set("Content-Type", "text/html; charset=utf-8")
resp.Body = readerWithRelease(buf.Bytes(), func(_ []byte) {
m.bytesPool.PutBuffer(buf)
pool.PutBuffer(buf)
})
return nil
}

View File

@@ -45,10 +45,6 @@ var (
var fontCSSTemplate = template.Must(template.New("fontCSS").Parse(fontCSS))
func (m *themed) setup() {
m.m.setup()
}
func (m *themed) before(w http.ResponseWriter, req *http.Request) bool {
return m.m.before(w, req)
}