ginuerzh
aba7e2a256
Split ~1250-line sniffer.go into sniffer_http.go (HTTP handling, dial, round-trip), sniffer_tls.go (TLS MITM, termination), sniffer_h2.go (HTTP/2), sniffer_ws.go (WebSocket frame copy/record), and sniffer_rewrite.go (upgrade, body rewrite). Extract clampBodySize, normalizeHost, effectiveReadTimeout as package-level functions to simplify testing. Add 22 tests covering pure functions, option setters, HTTP proxy integration, WebSocket frame copy, and edge cases.
93 lines
2.2 KiB
Go
93 lines
2.2 KiB
Go
package forwarder
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"strings"
|
|
|
|
"github.com/go-gost/core/chain"
|
|
"github.com/go-gost/core/logger"
|
|
xnet "github.com/go-gost/x/internal/net"
|
|
xrecorder "github.com/go-gost/x/recorder"
|
|
"golang.org/x/net/http/httpguts"
|
|
)
|
|
|
|
func upgradeType(h http.Header) string {
|
|
if !httpguts.HeaderValuesContainsToken(h["Connection"], "Upgrade") {
|
|
return ""
|
|
}
|
|
return h.Get("Upgrade")
|
|
}
|
|
|
|
func (h *Sniffer) handleUpgradeResponse(ctx context.Context, rw, cc io.ReadWriteCloser, req *http.Request, res *http.Response, ro *xrecorder.HandlerRecorderObject, log logger.Logger) error {
|
|
reqUpType := upgradeType(req.Header)
|
|
resUpType := upgradeType(res.Header)
|
|
if !strings.EqualFold(reqUpType, resUpType) {
|
|
return fmt.Errorf("backend tried to switch protocol %q when %q was requested", resUpType, reqUpType)
|
|
}
|
|
|
|
res.Body = nil
|
|
if err := res.Write(rw); err != nil {
|
|
return fmt.Errorf("response write: %v", err)
|
|
}
|
|
|
|
if reqUpType == "websocket" && h.Websocket {
|
|
return h.sniffingWebsocketFrame(ctx, rw, cc, ro, log)
|
|
}
|
|
|
|
return xnet.Pipe(ctx, rw, cc)
|
|
}
|
|
|
|
func rewriteRespBody(resp *http.Response, rewrites ...chain.HTTPBodyRewriteSettings) error {
|
|
if resp == nil || len(rewrites) == 0 || resp.ContentLength <= 0 {
|
|
return nil
|
|
}
|
|
|
|
if encoding := resp.Header.Get("Content-Encoding"); encoding != "" {
|
|
return nil
|
|
}
|
|
|
|
body, err := drainBody(resp.Body)
|
|
if err != nil || body == nil {
|
|
return err
|
|
}
|
|
|
|
contentType, _, _ := strings.Cut(resp.Header.Get("Content-Type"), ";")
|
|
for _, rewrite := range rewrites {
|
|
rewriteType := rewrite.Type
|
|
if rewriteType == "" {
|
|
rewriteType = "text/html"
|
|
}
|
|
if rewriteType != "*" && !strings.Contains(rewriteType, contentType) {
|
|
continue
|
|
}
|
|
|
|
if rewrite.Pattern != nil {
|
|
body = rewrite.Pattern.ReplaceAll(body, rewrite.Replacement)
|
|
}
|
|
}
|
|
|
|
resp.Body = io.NopCloser(bytes.NewReader(body))
|
|
resp.ContentLength = int64(len(body))
|
|
|
|
return nil
|
|
}
|
|
|
|
func drainBody(b io.ReadCloser) (body []byte, err error) {
|
|
if b == nil || b == http.NoBody {
|
|
return nil, nil
|
|
}
|
|
var buf bytes.Buffer
|
|
if _, err = buf.ReadFrom(b); err != nil {
|
|
return nil, err
|
|
}
|
|
if err = b.Close(); err != nil {
|
|
return nil, err
|
|
}
|
|
return buf.Bytes(), nil
|
|
}
|
|
|