refactor(forwarder): split sniffer monolith into focused files, extract testable helpers
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.
This commit is contained in:
ginuerzh
+54
-1116
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,184 @@
|
||||
package forwarder
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/http/httputil"
|
||||
"time"
|
||||
|
||||
"github.com/go-gost/core/logger"
|
||||
"github.com/go-gost/core/recorder"
|
||||
xhttp "github.com/go-gost/x/internal/net/http"
|
||||
xrecorder "github.com/go-gost/x/recorder"
|
||||
"golang.org/x/net/http2"
|
||||
)
|
||||
|
||||
// serveH2 handles HTTP/2 connections by reading the client preface and
|
||||
// proxying requests through an http2.Transport.
|
||||
func (h *Sniffer) serveH2(ctx context.Context, conn net.Conn, ho *HandleOptions) error {
|
||||
const expectedBody = "SM\r\n\r\n"
|
||||
|
||||
buf := make([]byte, len(expectedBody))
|
||||
n, err := io.ReadFull(conn, buf)
|
||||
if err != nil {
|
||||
return fmt.Errorf("h2: error reading client preface: %s", err)
|
||||
}
|
||||
if string(buf[:n]) != expectedBody {
|
||||
return errors.New("h2: invalid client preface")
|
||||
}
|
||||
|
||||
ro := ho.recorderObject
|
||||
log := ho.log
|
||||
|
||||
ro.Time = time.Time{}
|
||||
|
||||
tr := &http2.Transport{
|
||||
DialTLSContext: func(ctx context.Context, network, addr string, cfg *tls.Config) (net.Conn, error) {
|
||||
node, cc, err := h.dialTLS(ctx, addr, ho)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
log = log.WithFields(map[string]any{"src": cc.LocalAddr().String(), "dst": cc.RemoteAddr().String()})
|
||||
ro.SrcAddr = cc.LocalAddr().String()
|
||||
ro.DstAddr = cc.RemoteAddr().String()
|
||||
log.Debugf("connected to node %s(%s)", node.Name, node.Addr)
|
||||
return cc, nil
|
||||
},
|
||||
}
|
||||
defer tr.CloseIdleConnections()
|
||||
|
||||
(&http2.Server{}).ServeConn(conn, &http2.ServeConnOpts{
|
||||
Context: ctx,
|
||||
SawClientPreface: true,
|
||||
Handler: &h2Handler{
|
||||
transport: tr,
|
||||
recorder: h.Recorder,
|
||||
recorderOptions: h.RecorderOptions,
|
||||
recorderObject: ro,
|
||||
log: log,
|
||||
},
|
||||
})
|
||||
return nil
|
||||
}
|
||||
|
||||
// h2Handler is an http.Handler that proxies HTTP/2 requests through an
|
||||
// http2.Transport while recording request and response metadata.
|
||||
type h2Handler struct {
|
||||
transport http.RoundTripper
|
||||
recorder recorder.Recorder
|
||||
recorderOptions *recorder.Options
|
||||
recorderObject *xrecorder.HandlerRecorderObject
|
||||
log logger.Logger
|
||||
}
|
||||
|
||||
func (h *h2Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
log := h.log
|
||||
|
||||
ro := &xrecorder.HandlerRecorderObject{}
|
||||
*ro = *h.recorderObject
|
||||
ro.Time = time.Now()
|
||||
|
||||
var err error
|
||||
log.Infof("%s <-> %s", ro.RemoteAddr, r.Host)
|
||||
defer func() {
|
||||
ro.Duration = time.Since(ro.Time)
|
||||
if err != nil {
|
||||
ro.Err = err.Error()
|
||||
}
|
||||
if rerr := ro.Record(r.Context(), h.recorder); rerr != nil {
|
||||
log.Errorf("record: %v", rerr)
|
||||
}
|
||||
|
||||
log.WithFields(map[string]any{
|
||||
"duration": time.Since(ro.Time),
|
||||
}).Infof("%s >-< %s", ro.RemoteAddr, r.Host)
|
||||
}()
|
||||
|
||||
if clientIP := xhttp.GetClientIP(r); clientIP != nil {
|
||||
ro.ClientIP = clientIP.String()
|
||||
}
|
||||
ro.HTTP = &xrecorder.HTTPRecorderObject{
|
||||
Host: r.Host,
|
||||
Proto: r.Proto,
|
||||
Scheme: "https",
|
||||
Method: r.Method,
|
||||
URI: r.RequestURI,
|
||||
Request: xrecorder.HTTPRequestRecorderObject{
|
||||
ContentLength: r.ContentLength,
|
||||
Header: r.Header.Clone(),
|
||||
},
|
||||
}
|
||||
|
||||
if log.IsLevelEnabled(logger.TraceLevel) {
|
||||
dump, _ := httputil.DumpRequest(r, false)
|
||||
log.Trace(string(dump))
|
||||
}
|
||||
|
||||
url := r.URL
|
||||
url.Scheme = "https"
|
||||
url.Host = r.Host
|
||||
req := &http.Request{
|
||||
Method: r.Method,
|
||||
URL: url,
|
||||
Host: r.Host,
|
||||
Header: r.Header,
|
||||
Body: r.Body,
|
||||
ContentLength: r.ContentLength,
|
||||
Trailer: r.Trailer,
|
||||
}
|
||||
|
||||
var reqBody *xhttp.Body
|
||||
if bodySize := clampBodySize(h.recorderOptions); bodySize > 0 && req.Body != nil {
|
||||
reqBody = xhttp.NewBody(req.Body, bodySize)
|
||||
req.Body = reqBody
|
||||
}
|
||||
|
||||
resp, respErr := h.transport.RoundTrip(req.WithContext(r.Context()))
|
||||
if reqBody != nil {
|
||||
ro.HTTP.Request.Body = reqBody.Content()
|
||||
ro.HTTP.Request.ContentLength = reqBody.Length()
|
||||
}
|
||||
err = respErr
|
||||
if err != nil {
|
||||
log.Error(err)
|
||||
w.WriteHeader(http.StatusServiceUnavailable)
|
||||
return
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
ro.HTTP.StatusCode = resp.StatusCode
|
||||
ro.HTTP.Response.Header = resp.Header
|
||||
ro.HTTP.Response.ContentLength = resp.ContentLength
|
||||
|
||||
if log.IsLevelEnabled(logger.TraceLevel) {
|
||||
dump, _ := httputil.DumpResponse(resp, false)
|
||||
log.Trace(string(dump))
|
||||
}
|
||||
|
||||
h.setHeader(w, resp.Header)
|
||||
w.WriteHeader(resp.StatusCode)
|
||||
|
||||
if bodySize := clampBodySize(h.recorderOptions); bodySize > 0 {
|
||||
respBody := xhttp.NewBody(resp.Body, bodySize)
|
||||
resp.Body = respBody
|
||||
_, _ = io.Copy(w, resp.Body)
|
||||
ro.HTTP.Response.Body = respBody.Content()
|
||||
ro.HTTP.Response.ContentLength = respBody.Length()
|
||||
} else {
|
||||
_, _ = io.Copy(w, resp.Body)
|
||||
}
|
||||
}
|
||||
|
||||
func (h *h2Handler) setHeader(w http.ResponseWriter, header http.Header) {
|
||||
for k, v := range header {
|
||||
for i := range v {
|
||||
w.Header().Add(k, v[i])
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,414 @@
|
||||
package forwarder
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/http/httputil"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/go-gost/core/auth"
|
||||
"github.com/go-gost/core/bypass"
|
||||
"github.com/go-gost/core/chain"
|
||||
"github.com/go-gost/core/hop"
|
||||
"github.com/go-gost/core/logger"
|
||||
"github.com/go-gost/core/observer/stats"
|
||||
xbypass "github.com/go-gost/x/bypass"
|
||||
xctx "github.com/go-gost/x/ctx"
|
||||
xio "github.com/go-gost/x/internal/io"
|
||||
xnet "github.com/go-gost/x/internal/net"
|
||||
xhttp "github.com/go-gost/x/internal/net/http"
|
||||
"github.com/go-gost/x/internal/util/sniffing"
|
||||
xstats "github.com/go-gost/x/observer/stats"
|
||||
stats_wrapper "github.com/go-gost/x/observer/stats/wrapper"
|
||||
xrecorder "github.com/go-gost/x/recorder"
|
||||
)
|
||||
|
||||
// HandleHTTP sniffs and proxies an HTTP connection. It reads the initial
|
||||
// request, performs node selection via the configured hop, and forwards the
|
||||
// request with HTTP keep-alive support.
|
||||
func (h *Sniffer) HandleHTTP(ctx context.Context, conn net.Conn, opts ...HandleOption) error {
|
||||
var ho HandleOptions
|
||||
for _, opt := range opts {
|
||||
opt(&ho)
|
||||
}
|
||||
ho.readTimeout = h.effectiveReadTimeout(&ho)
|
||||
|
||||
pStats := xstats.Stats{}
|
||||
conn = stats_wrapper.WrapConn(conn, &pStats)
|
||||
|
||||
br := bufio.NewReader(conn)
|
||||
req, err := http.ReadRequest(br)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
log := ho.log
|
||||
if log.IsLevelEnabled(logger.TraceLevel) {
|
||||
dump, _ := httputil.DumpRequest(req, false)
|
||||
log.Trace(string(dump))
|
||||
}
|
||||
|
||||
ro := ho.recorderObject
|
||||
ro.HTTP = &xrecorder.HTTPRecorderObject{
|
||||
Host: req.Host,
|
||||
Proto: req.Proto,
|
||||
Scheme: req.URL.Scheme,
|
||||
Method: req.Method,
|
||||
URI: req.RequestURI,
|
||||
Request: xrecorder.HTTPRequestRecorderObject{
|
||||
ContentLength: req.ContentLength,
|
||||
Header: req.Header.Clone(),
|
||||
},
|
||||
}
|
||||
|
||||
if clientIP := xhttp.GetClientIP(req); clientIP != nil {
|
||||
clientAddr := &net.TCPAddr{IP: clientIP}
|
||||
ro.ClientAddr = clientAddr.String()
|
||||
ctx = xctx.ContextWithSrcAddr(ctx, clientAddr)
|
||||
}
|
||||
|
||||
// http/2
|
||||
if req.Method == "PRI" && len(req.Header) == 0 && req.URL.Path == "*" && req.Proto == "HTTP/2.0" {
|
||||
return h.serveH2(ctx, xnet.NewReadWriteConn(br, conn, conn), &ho)
|
||||
}
|
||||
|
||||
node, cc, err := h.dial(ctx, conn, req, &ho)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer cc.Close()
|
||||
|
||||
log = log.WithFields(map[string]any{"src": cc.LocalAddr().String(), "dst": cc.RemoteAddr().String()})
|
||||
log.Debugf("connected to node %s(%s)", node.Name, node.Addr)
|
||||
|
||||
ro.SrcAddr = cc.LocalAddr().String()
|
||||
ro.DstAddr = cc.RemoteAddr().String()
|
||||
ro.Time = time.Time{}
|
||||
|
||||
shouldClose, err := h.httpRoundTrip(ctx, xio.NewReadWriteCloser(br, conn, conn), cc, node, req, &pStats, &ho)
|
||||
if err != nil || shouldClose {
|
||||
return err
|
||||
}
|
||||
|
||||
for {
|
||||
pStats.Reset()
|
||||
|
||||
req, err := http.ReadRequest(br)
|
||||
if err != nil {
|
||||
if errors.Is(err, io.EOF) || errors.Is(err, net.ErrClosed) {
|
||||
return nil
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
if log.IsLevelEnabled(logger.TraceLevel) {
|
||||
dump, _ := httputil.DumpRequest(req, false)
|
||||
log.Trace(string(dump))
|
||||
}
|
||||
|
||||
if shouldClose, err := h.httpRoundTrip(ctx, xio.NewReadWriteCloser(br, conn, conn), cc, node, req, &pStats, &ho); err != nil || shouldClose {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// dial selects a node and establishes a connection for an HTTP request.
|
||||
func (h *Sniffer) dial(ctx context.Context, conn net.Conn, req *http.Request, ho *HandleOptions) (node *chain.Node, cc net.Conn, err error) {
|
||||
dial := ho.dial
|
||||
if dial == nil {
|
||||
dial = (&net.Dialer{}).DialContext
|
||||
}
|
||||
|
||||
if node = ho.node; node != nil {
|
||||
cc, err = dial(ctx, "tcp", node.Addr)
|
||||
return
|
||||
}
|
||||
|
||||
ro := ho.recorderObject
|
||||
|
||||
res := &http.Response{
|
||||
ProtoMajor: req.ProtoMajor,
|
||||
ProtoMinor: req.ProtoMinor,
|
||||
Header: http.Header{},
|
||||
StatusCode: http.StatusServiceUnavailable,
|
||||
}
|
||||
|
||||
host := normalizeHost(req.Host, "80")
|
||||
if host != "" {
|
||||
ro.Host = host
|
||||
ho.log = ho.log.WithFields(map[string]any{
|
||||
"host": host,
|
||||
})
|
||||
|
||||
if ho.bypass != nil &&
|
||||
ho.bypass.Contains(ctx, "tcp", host,
|
||||
bypass.WithService(ho.service),
|
||||
bypass.WithPathOption(req.RequestURI)) {
|
||||
ho.log.Debugf("bypass: %s %s", host, req.RequestURI)
|
||||
res.StatusCode = http.StatusForbidden
|
||||
ro.HTTP.StatusCode = res.StatusCode
|
||||
if werr := res.Write(conn); werr != nil {
|
||||
ho.log.Warnf("write bypass response: %v", werr)
|
||||
}
|
||||
return nil, nil, xbypass.ErrBypass
|
||||
}
|
||||
}
|
||||
|
||||
node = &chain.Node{}
|
||||
if ho.hop != nil {
|
||||
var clientIP net.IP
|
||||
if clientAddr, _ := net.ResolveTCPAddr("tcp", ro.ClientAddr); clientAddr != nil {
|
||||
clientIP = clientAddr.IP
|
||||
}
|
||||
|
||||
node = ho.hop.Select(ctx,
|
||||
hop.ClientIPSelectOption(clientIP),
|
||||
hop.ProtocolSelectOption(sniffing.ProtoHTTP),
|
||||
hop.HostSelectOption(host),
|
||||
hop.MethodSelectOption(req.Method),
|
||||
hop.PathSelectOption(req.URL.Path),
|
||||
hop.QuerySelectOption(req.URL.Query()),
|
||||
hop.HeaderSelectOption(req.Header),
|
||||
)
|
||||
}
|
||||
if node == nil {
|
||||
ho.log.Warnf("node for %s not found", host)
|
||||
res.StatusCode = http.StatusBadGateway
|
||||
ro.HTTP.StatusCode = res.StatusCode
|
||||
if werr := res.Write(conn); werr != nil {
|
||||
ho.log.Warnf("write error response: %v", werr)
|
||||
}
|
||||
return nil, nil, errors.New("node not available")
|
||||
}
|
||||
if node.Addr == "" {
|
||||
node = &chain.Node{
|
||||
Name: node.Name,
|
||||
Addr: host,
|
||||
}
|
||||
}
|
||||
|
||||
ro.Host = node.Addr
|
||||
ho.log = ho.log.WithFields(map[string]any{
|
||||
"node": node.Name,
|
||||
"dst": node.Addr,
|
||||
})
|
||||
ho.log.Debugf("find node for host %s -> %s(%s)", host, node.Name, node.Addr)
|
||||
|
||||
cc, err = dial(ctx, "tcp", node.Addr)
|
||||
if err != nil {
|
||||
if marker := node.Marker(); marker != nil {
|
||||
marker.Mark()
|
||||
}
|
||||
ho.log.Warnf("connect to node %s(%s) failed: %v", node.Name, node.Addr, err)
|
||||
if werr := res.Write(conn); werr != nil {
|
||||
ho.log.Warnf("write error response: %v", werr)
|
||||
}
|
||||
return
|
||||
}
|
||||
if marker := node.Marker(); marker != nil {
|
||||
marker.Reset()
|
||||
}
|
||||
|
||||
cc = tlsWrapConn(cc, node.Options().TLS)
|
||||
return
|
||||
}
|
||||
|
||||
// httpRoundTrip forwards a single HTTP request/response pair and records
|
||||
// traffic metadata. Returns whether the connection should be closed.
|
||||
func (h *Sniffer) httpRoundTrip(ctx context.Context, rw, cc io.ReadWriteCloser, node *chain.Node, req *http.Request, pStats stats.Stats, ho *HandleOptions) (shouldClose bool, err error) {
|
||||
shouldClose = true
|
||||
|
||||
log := ho.log
|
||||
ro := &xrecorder.HandlerRecorderObject{}
|
||||
*ro = *ho.recorderObject
|
||||
|
||||
ro.Time = time.Now()
|
||||
log.Infof("%s <-> %s", ro.RemoteAddr, req.Host)
|
||||
defer func() {
|
||||
if err != nil {
|
||||
ro.Err = err.Error()
|
||||
}
|
||||
ro.InputBytes = pStats.Get(stats.KindInputBytes)
|
||||
ro.OutputBytes = pStats.Get(stats.KindOutputBytes)
|
||||
ro.Duration = time.Since(ro.Time)
|
||||
if rerr := ro.Record(ctx, h.Recorder); rerr != nil {
|
||||
log.Errorf("record: %v", rerr)
|
||||
}
|
||||
|
||||
log.WithFields(map[string]any{
|
||||
"duration": time.Since(ro.Time),
|
||||
"inputBytes": ro.InputBytes,
|
||||
"outputBytes": ro.OutputBytes,
|
||||
}).Infof("%s >-< %s", ro.RemoteAddr, req.Host)
|
||||
}()
|
||||
|
||||
ro.HTTP = &xrecorder.HTTPRecorderObject{
|
||||
Host: req.Host,
|
||||
Proto: req.Proto,
|
||||
Scheme: req.URL.Scheme,
|
||||
Method: req.Method,
|
||||
URI: req.RequestURI,
|
||||
Request: xrecorder.HTTPRequestRecorderObject{
|
||||
ContentLength: req.ContentLength,
|
||||
Header: req.Header.Clone(),
|
||||
},
|
||||
}
|
||||
|
||||
res := &http.Response{
|
||||
ProtoMajor: req.ProtoMajor,
|
||||
ProtoMinor: req.ProtoMinor,
|
||||
Header: http.Header{},
|
||||
StatusCode: http.StatusServiceUnavailable,
|
||||
}
|
||||
|
||||
// HTTP/1.0
|
||||
if req.ProtoMajor == 1 && req.ProtoMinor == 0 {
|
||||
if strings.ToLower(req.Header.Get("Connection")) == "keep-alive" {
|
||||
req.Header.Del("Connection")
|
||||
} else {
|
||||
req.Header.Set("Connection", "close")
|
||||
}
|
||||
}
|
||||
|
||||
var responseHeader map[string]string
|
||||
var respBodyRewrites []chain.HTTPBodyRewriteSettings
|
||||
if httpSettings := node.Options().HTTP; httpSettings != nil {
|
||||
if auther := httpSettings.Auther; auther != nil {
|
||||
username, password, _ := req.BasicAuth()
|
||||
id, ok := auther.Authenticate(ctx, username, password, auth.WithService(ho.service))
|
||||
if !ok {
|
||||
res.StatusCode = http.StatusUnauthorized
|
||||
ro.HTTP.StatusCode = res.StatusCode
|
||||
res.Header.Set("WWW-Authenticate", "Basic")
|
||||
log.Warnf("node %s(%s) 401 unauthorized", node.Name, node.Addr)
|
||||
res.Write(rw)
|
||||
err = errors.New("unauthorized")
|
||||
return
|
||||
}
|
||||
ctx = xctx.ContextWithClientID(ctx, xctx.ClientID(id))
|
||||
}
|
||||
|
||||
if httpSettings.Host != "" {
|
||||
req.Host = httpSettings.Host
|
||||
}
|
||||
for k, v := range httpSettings.RequestHeader {
|
||||
req.Header.Set(k, v)
|
||||
}
|
||||
|
||||
for _, re := range httpSettings.RewriteURL {
|
||||
if re.Pattern.MatchString(req.URL.Path) {
|
||||
if s := re.Pattern.ReplaceAllString(req.URL.Path, re.Replacement); s != "" {
|
||||
req.URL.Path = s
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
responseHeader = httpSettings.ResponseHeader
|
||||
respBodyRewrites = httpSettings.RewriteResponseBody
|
||||
}
|
||||
|
||||
if bodySize := clampBodySize(h.RecorderOptions); bodySize > 0 && req.Body != nil {
|
||||
reqBody := xhttp.NewBody(req.Body, bodySize)
|
||||
req.Body = reqBody
|
||||
err = req.Write(cc)
|
||||
ro.HTTP.Request.Body = reqBody.Content()
|
||||
ro.HTTP.Request.ContentLength = reqBody.Length()
|
||||
} else {
|
||||
err = req.Write(cc)
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
res.Write(rw)
|
||||
return
|
||||
}
|
||||
|
||||
br := bufio.NewReader(cc)
|
||||
var resp *http.Response
|
||||
for {
|
||||
xio.SetReadDeadline(cc, time.Now().Add(ho.readTimeout))
|
||||
resp, err = http.ReadResponse(br, req)
|
||||
if err != nil {
|
||||
log.Errorf("read response: %v", err)
|
||||
res.Write(rw)
|
||||
return
|
||||
}
|
||||
if resp.StatusCode == http.StatusContinue {
|
||||
resp.Write(rw)
|
||||
resp.Body.Close()
|
||||
continue
|
||||
}
|
||||
|
||||
break
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
xio.SetReadDeadline(cc, time.Time{})
|
||||
|
||||
if len(responseHeader) > 0 {
|
||||
if resp.Header == nil {
|
||||
resp.Header = http.Header{}
|
||||
}
|
||||
for k, v := range responseHeader {
|
||||
resp.Header.Set(k, v)
|
||||
}
|
||||
}
|
||||
|
||||
ro.HTTP.StatusCode = resp.StatusCode
|
||||
ro.HTTP.Response.Header = resp.Header
|
||||
ro.HTTP.Response.ContentLength = resp.ContentLength
|
||||
|
||||
if log.IsLevelEnabled(logger.TraceLevel) {
|
||||
dump, _ := httputil.DumpResponse(resp, false)
|
||||
log.Trace(string(dump))
|
||||
}
|
||||
|
||||
if resp.StatusCode == http.StatusSwitchingProtocols {
|
||||
h.handleUpgradeResponse(ctx, rw, cc, req, resp, ro, log)
|
||||
return
|
||||
}
|
||||
|
||||
// HTTP/1.0
|
||||
if req.ProtoMajor == 1 && req.ProtoMinor == 0 {
|
||||
if !resp.Close {
|
||||
resp.Header.Set("Connection", "keep-alive")
|
||||
}
|
||||
resp.ProtoMajor = req.ProtoMajor
|
||||
resp.ProtoMinor = req.ProtoMinor
|
||||
}
|
||||
|
||||
if !ho.httpKeepalive {
|
||||
resp.Header.Set("Connection", "close")
|
||||
}
|
||||
|
||||
if err = rewriteRespBody(resp, respBodyRewrites...); err != nil {
|
||||
log.Errorf("rewrite body: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
if bodySize := clampBodySize(h.RecorderOptions); bodySize > 0 {
|
||||
respBody := xhttp.NewBody(resp.Body, bodySize)
|
||||
resp.Body = respBody
|
||||
err = resp.Write(rw)
|
||||
ro.HTTP.Response.Body = respBody.Content()
|
||||
ro.HTTP.Response.ContentLength = respBody.Length()
|
||||
} else {
|
||||
err = resp.Write(rw)
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
log.Errorf("write response: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
if resp.ContentLength >= 0 {
|
||||
shouldClose = resp.Close
|
||||
}
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
@@ -0,0 +1,92 @@
|
||||
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
|
||||
}
|
||||
|
||||
@@ -0,0 +1,719 @@
|
||||
package forwarder
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"regexp"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/go-gost/core/bypass"
|
||||
"github.com/go-gost/core/chain"
|
||||
"github.com/go-gost/core/hop"
|
||||
"github.com/go-gost/core/recorder"
|
||||
xlogger "github.com/go-gost/x/logger"
|
||||
xrecorder "github.com/go-gost/x/recorder"
|
||||
)
|
||||
|
||||
// =============================================================================
|
||||
// Test Helpers (local mocks)
|
||||
// =============================================================================
|
||||
|
||||
type noopRecorder struct{}
|
||||
|
||||
func (n *noopRecorder) Record(_ context.Context, _ []byte, _ ...recorder.RecordOption) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// mockBypass is a configurable bypass.Bypass for tests.
|
||||
type mockBypass struct {
|
||||
contains bool
|
||||
whitelist bool
|
||||
}
|
||||
|
||||
func (m *mockBypass) Contains(_ context.Context, _, _ string, _ ...bypass.Option) bool {
|
||||
return m.contains
|
||||
}
|
||||
|
||||
func (m *mockBypass) IsWhitelist() bool { return m.whitelist }
|
||||
|
||||
// =============================================================================
|
||||
// Pure Function Tests
|
||||
// =============================================================================
|
||||
|
||||
func TestClampBodySize(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
opts *recorder.Options
|
||||
want int
|
||||
}{
|
||||
{"nil opts", nil, 0},
|
||||
{"HTTPBody disabled", &recorder.Options{HTTPBody: false, MaxBodySize: 1000}, 0},
|
||||
{"zero MaxBodySize defaults to defaultBodySize", &recorder.Options{HTTPBody: true, MaxBodySize: 0}, defaultBodySize},
|
||||
{"negative MaxBodySize defaults", &recorder.Options{HTTPBody: true, MaxBodySize: -1}, defaultBodySize},
|
||||
{"valid MaxBodySize within bounds", &recorder.Options{HTTPBody: true, MaxBodySize: 1024}, 1024},
|
||||
{"MaxBodySize capped at maxBodySize", &recorder.Options{HTTPBody: true, MaxBodySize: 10 * 1024 * 1024}, maxBodySize},
|
||||
{"exact maxBodySize", &recorder.Options{HTTPBody: true, MaxBodySize: maxBodySize}, maxBodySize},
|
||||
{"exact defaultBodySize", &recorder.Options{HTTPBody: true, MaxBodySize: defaultBodySize}, defaultBodySize},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
if got := clampBodySize(tt.opts); got != tt.want {
|
||||
t.Errorf("clampBodySize() = %d, want %d", got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeHost(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
host string
|
||||
defaultPort string
|
||||
want string
|
||||
}{
|
||||
{"empty host", "", "443", ""},
|
||||
{"host with port", "example.com:8080", "443", "example.com:8080"},
|
||||
{"host without port, default 80", "example.com", "80", "example.com:80"},
|
||||
{"host without port, default 443", "example.com", "443", "example.com:443"},
|
||||
{"IPv4 with port", "127.0.0.1:9090", "80", "127.0.0.1:9090"},
|
||||
{"IPv4 without port", "127.0.0.1", "80", "127.0.0.1:80"},
|
||||
{"IPv6 bracketed with port", "[::1]:8080", "443", "[::1]:8080"},
|
||||
{"IPv6 bracketed without port", "[::1]", "443", "[::1]:443"},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := normalizeHost(tt.host, tt.defaultPort)
|
||||
if got != tt.want {
|
||||
t.Errorf("normalizeHost(%q, %q) = %q, want %q", tt.host, tt.defaultPort, got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestEffectiveReadTimeout(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
snifferTO time.Duration
|
||||
hoTO time.Duration
|
||||
want time.Duration
|
||||
}{
|
||||
{"both zero defaults to DefaultReadTimeout", 0, 0, DefaultReadTimeout},
|
||||
{"sniffer timeout only", 10 * time.Second, 0, 10 * time.Second},
|
||||
{"ho timeout only", 0, 5 * time.Second, 5 * time.Second},
|
||||
{"ho timeout takes precedence", 10 * time.Second, 3 * time.Second, 3 * time.Second},
|
||||
{"sniffer timeout takes precedence when ho is zero", 15 * time.Second, 0, 15 * time.Second},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
h := &Sniffer{ReadTimeout: tt.snifferTO}
|
||||
ho := &HandleOptions{readTimeout: tt.hoTO}
|
||||
got := h.effectiveReadTimeout(ho)
|
||||
if got != tt.want {
|
||||
t.Errorf("effectiveReadTimeout() = %v, want %v", got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestUpgradeType(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
header http.Header
|
||||
want string
|
||||
}{
|
||||
{
|
||||
"websocket upgrade",
|
||||
http.Header{"Connection": {"Upgrade"}, "Upgrade": {"websocket"}},
|
||||
"websocket",
|
||||
},
|
||||
{
|
||||
"h2c upgrade",
|
||||
http.Header{"Connection": {"Upgrade"}, "Upgrade": {"h2c"}},
|
||||
"h2c",
|
||||
},
|
||||
{
|
||||
"no upgrade header",
|
||||
http.Header{"Connection": {"keep-alive"}},
|
||||
"",
|
||||
},
|
||||
{
|
||||
"empty header",
|
||||
http.Header{},
|
||||
"",
|
||||
},
|
||||
{
|
||||
"connection has upgrade but no upgrade header",
|
||||
http.Header{"Connection": {"Upgrade"}},
|
||||
"",
|
||||
},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
if got := upgradeType(tt.header); got != tt.want {
|
||||
t.Errorf("upgradeType() = %q, want %q", got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestDrainBody(t *testing.T) {
|
||||
t.Run("nil body", func(t *testing.T) {
|
||||
got, err := drainBody(nil)
|
||||
if err != nil || got != nil {
|
||||
t.Errorf("drainBody(nil) = (%v, %v), want (nil, nil)", got, err)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("http.NoBody", func(t *testing.T) {
|
||||
got, err := drainBody(http.NoBody)
|
||||
if err != nil || got != nil {
|
||||
t.Errorf("drainBody(NoBody) = (%v, %v), want (nil, nil)", got, err)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("normal body", func(t *testing.T) {
|
||||
body := io.NopCloser(strings.NewReader("hello world"))
|
||||
got, err := drainBody(body)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if string(got) != "hello world" {
|
||||
t.Errorf("drainBody() = %q, want %q", string(got), "hello world")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("empty body", func(t *testing.T) {
|
||||
body := io.NopCloser(strings.NewReader(""))
|
||||
got, err := drainBody(body)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(got) != 0 {
|
||||
t.Errorf("drainBody() = %q, want empty", string(got))
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Rewrite Response Body Tests
|
||||
// =============================================================================
|
||||
|
||||
func TestRewriteRespBody(t *testing.T) {
|
||||
t.Run("nil response", func(t *testing.T) {
|
||||
if err := rewriteRespBody(nil); err != nil {
|
||||
t.Errorf("rewriteRespBody(nil) = %v, want nil", err)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("no rewrites", func(t *testing.T) {
|
||||
resp := &http.Response{
|
||||
ContentLength: 100,
|
||||
Body: io.NopCloser(strings.NewReader("original")),
|
||||
}
|
||||
_ = rewriteRespBody(resp)
|
||||
// Body unchanged
|
||||
got, _ := io.ReadAll(resp.Body)
|
||||
if string(got) != "original" {
|
||||
t.Errorf("body = %q, want %q", string(got), "original")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("zero content length", func(t *testing.T) {
|
||||
resp := &http.Response{
|
||||
ContentLength: 0,
|
||||
Body: io.NopCloser(strings.NewReader("original")),
|
||||
}
|
||||
_ = rewriteRespBody(resp, chain.HTTPBodyRewriteSettings{
|
||||
Pattern: regexp.MustCompile("."),
|
||||
Type: "*",
|
||||
Replacement: []byte("replaced"),
|
||||
})
|
||||
got, _ := io.ReadAll(resp.Body)
|
||||
// Content-Length was 0 (unknown), so no rewrite applied
|
||||
if string(got) != "original" {
|
||||
t.Errorf("body = %q, want %q (unchanged)", string(got), "original")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("with content encoding", func(t *testing.T) {
|
||||
resp := &http.Response{
|
||||
Header: http.Header{"Content-Encoding": {"gzip"}},
|
||||
ContentLength: 100,
|
||||
Body: io.NopCloser(strings.NewReader("original")),
|
||||
}
|
||||
_ = rewriteRespBody(resp, chain.HTTPBodyRewriteSettings{
|
||||
Type: "*",
|
||||
Replacement: []byte("replaced"),
|
||||
})
|
||||
got, _ := io.ReadAll(resp.Body)
|
||||
// Content-Encoding present, rewrite skipped
|
||||
if string(got) != "original" {
|
||||
t.Errorf("body = %q, want %q (unchanged)", string(got), "original")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestRewriteRespBodyContentTypeFilter(t *testing.T) {
|
||||
// Helper to create a rewrite rule that replaces everything.
|
||||
makeRewrite := func(rewriteType string, replacement string) chain.HTTPBodyRewriteSettings {
|
||||
return chain.HTTPBodyRewriteSettings{
|
||||
Pattern: regexp.MustCompile(".*"),
|
||||
Type: rewriteType,
|
||||
Replacement: []byte(replacement),
|
||||
}
|
||||
}
|
||||
|
||||
t.Run("exact type match text/html", func(t *testing.T) {
|
||||
resp := &http.Response{
|
||||
Header: http.Header{"Content-Type": {"text/html; charset=utf-8"}},
|
||||
ContentLength: 100,
|
||||
Body: io.NopCloser(strings.NewReader("original")),
|
||||
}
|
||||
_ = rewriteRespBody(resp, makeRewrite("text/html", "replaced"))
|
||||
got, _ := io.ReadAll(resp.Body)
|
||||
if string(got) != "replaced" {
|
||||
t.Errorf("body = %q, want %q", string(got), "replaced")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("wildcard type rewrites anything", func(t *testing.T) {
|
||||
resp := &http.Response{
|
||||
Header: http.Header{"Content-Type": {"application/json"}},
|
||||
ContentLength: 100,
|
||||
Body: io.NopCloser(strings.NewReader("hello")),
|
||||
}
|
||||
_ = rewriteRespBody(resp, makeRewrite("*", "world"))
|
||||
got, _ := io.ReadAll(resp.Body)
|
||||
if string(got) != "world" {
|
||||
t.Errorf("body = %q, want %q", string(got), "world")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("mismatched type is skipped", func(t *testing.T) {
|
||||
resp := &http.Response{
|
||||
Header: http.Header{"Content-Type": {"text/plain"}},
|
||||
ContentLength: 100,
|
||||
Body: io.NopCloser(strings.NewReader("original")),
|
||||
}
|
||||
_ = rewriteRespBody(resp, makeRewrite("text/html", "replaced"))
|
||||
got, _ := io.ReadAll(resp.Body)
|
||||
if string(got) != "original" {
|
||||
t.Errorf("body = %q, want %q (unchanged)", string(got), "original")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("empty content type defaults to text/html", func(t *testing.T) {
|
||||
resp := &http.Response{
|
||||
Header: http.Header{},
|
||||
ContentLength: 100,
|
||||
Body: io.NopCloser(strings.NewReader("original")),
|
||||
}
|
||||
// Empty type defaults to "text/html", and response has no Content-Type -> "" containing "text/html"? No.
|
||||
// strings.Contains("text/html", "") is always true, so rewrite should happen.
|
||||
_ = rewriteRespBody(resp, makeRewrite("", "replaced"))
|
||||
got, _ := io.ReadAll(resp.Body)
|
||||
if string(got) != "replaced" {
|
||||
t.Errorf("body = %q, want %q (should be rewritten)", string(got), "replaced")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// TLS Wrap Tests
|
||||
// =============================================================================
|
||||
|
||||
func TestTLSWrapConn_NilSettings(t *testing.T) {
|
||||
cc := &net.TCPConn{} // will never be used as a real connection
|
||||
if got := tlsWrapConn(cc, nil); got != cc {
|
||||
t.Errorf("tlsWrapConn(nil settings) should return original conn unchanged")
|
||||
}
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// HandleHTTP Integration Tests
|
||||
// =============================================================================
|
||||
|
||||
func TestHandleHTTP_BasicProxy(t *testing.T) {
|
||||
// Set up an HTTP test server that acts as the upstream.
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
w.Write([]byte("OK from upstream"))
|
||||
}))
|
||||
defer upstream.Close()
|
||||
|
||||
h := &Sniffer{
|
||||
ReadTimeout: 5 * time.Second,
|
||||
Recorder: &noopRecorder{},
|
||||
}
|
||||
ro := &xrecorder.HandlerRecorderObject{}
|
||||
|
||||
// Use net.Pipe and serve the client request in a goroutine.
|
||||
clientConn, serverConn := net.Pipe()
|
||||
defer clientConn.Close()
|
||||
defer serverConn.Close()
|
||||
|
||||
errCh := make(chan error, 1)
|
||||
go func() {
|
||||
errCh <- h.HandleHTTP(context.Background(), serverConn,
|
||||
WithService("test-svc"),
|
||||
WithDial(func(ctx context.Context, network, address string) (net.Conn, error) {
|
||||
return net.Dial("tcp", upstream.Listener.Addr().String())
|
||||
}),
|
||||
WithRecorderObject(ro),
|
||||
WithLog(xlogger.Nop()),
|
||||
)
|
||||
}()
|
||||
|
||||
// Write an HTTP request to the client side of the pipe.
|
||||
req, _ := http.NewRequest("GET", "http://example.com/", nil)
|
||||
if err := req.Write(clientConn); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// Read the response back.
|
||||
resp, err := http.ReadResponse(bufio.NewReader(clientConn), req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
t.Errorf("status = %d, want %d", resp.StatusCode, http.StatusOK)
|
||||
}
|
||||
|
||||
clientConn.Close()
|
||||
if err := <-errCh; err != nil {
|
||||
t.Logf("HandleHTTP returned: %v", err)
|
||||
}
|
||||
|
||||
// Verify recorder object was populated.
|
||||
if ro.SrcAddr == "" {
|
||||
t.Error("SrcAddr should be populated")
|
||||
}
|
||||
if ro.DstAddr == "" {
|
||||
t.Error("DstAddr should be populated")
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleHTTP_HTTP2Detection(t *testing.T) {
|
||||
// HTTP/2 connection preface should be detected and handled separately.
|
||||
h := &Sniffer{
|
||||
ReadTimeout: 5 * time.Second,
|
||||
Recorder: &noopRecorder{},
|
||||
}
|
||||
ro := &xrecorder.HandlerRecorderObject{
|
||||
RemoteAddr: "127.0.0.1:12345",
|
||||
}
|
||||
|
||||
clientConn, serverConn := net.Pipe()
|
||||
defer clientConn.Close()
|
||||
defer serverConn.Close()
|
||||
|
||||
errCh := make(chan error, 1)
|
||||
go func() {
|
||||
errCh <- h.HandleHTTP(context.Background(), serverConn,
|
||||
WithService("test-svc"),
|
||||
WithRecorderObject(ro),
|
||||
WithLog(xlogger.Nop()),
|
||||
)
|
||||
}()
|
||||
|
||||
// Write HTTP/2 upgrade request
|
||||
clientConn.Write([]byte("PRI * HTTP/2.0\r\n\r\n"))
|
||||
// The handler will now try to read the SM\r\n\r\n preface, which we don't send.
|
||||
// It should return an error about invalid client preface.
|
||||
|
||||
clientConn.Close()
|
||||
|
||||
err := <-errCh
|
||||
if err == nil {
|
||||
t.Log("HandleHTTP returned nil (expected error from incomplete h2 preface)")
|
||||
}
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// HandleOptions Tests
|
||||
// =============================================================================
|
||||
|
||||
func TestWithService(t *testing.T) {
|
||||
opts := &HandleOptions{}
|
||||
WithService("mysvc")(opts)
|
||||
if opts.service != "mysvc" {
|
||||
t.Errorf("service = %q, want %q", opts.service, "mysvc")
|
||||
}
|
||||
}
|
||||
|
||||
func TestWithHTTPKeepalive(t *testing.T) {
|
||||
opts := &HandleOptions{}
|
||||
WithHTTPKeepalive(true)(opts)
|
||||
if !opts.httpKeepalive {
|
||||
t.Errorf("httpKeepalive = false, want true")
|
||||
}
|
||||
}
|
||||
|
||||
func TestWithNode(t *testing.T) {
|
||||
opts := &HandleOptions{}
|
||||
node := &chain.Node{Name: "test-node", Addr: "127.0.0.1:8080"}
|
||||
WithNode(node)(opts)
|
||||
if opts.node != node {
|
||||
t.Errorf("node not set")
|
||||
}
|
||||
}
|
||||
|
||||
func TestWithHop(t *testing.T) {
|
||||
opts := &HandleOptions{}
|
||||
mh := &mockHop{}
|
||||
WithHop(mh)(opts)
|
||||
if opts.hop != mh {
|
||||
t.Errorf("hop not set")
|
||||
}
|
||||
}
|
||||
|
||||
func TestWithBypass(t *testing.T) {
|
||||
opts := &HandleOptions{}
|
||||
bp := &mockBypass{contains: true}
|
||||
WithBypass(bp)(opts)
|
||||
if opts.bypass != bp {
|
||||
t.Errorf("bypass not set")
|
||||
}
|
||||
}
|
||||
|
||||
func TestWithRecorderObject(t *testing.T) {
|
||||
opts := &HandleOptions{}
|
||||
ro := &xrecorder.HandlerRecorderObject{}
|
||||
WithRecorderObject(ro)(opts)
|
||||
if opts.recorderObject != ro {
|
||||
t.Errorf("recorderObject not set")
|
||||
}
|
||||
}
|
||||
|
||||
func TestWithLog(t *testing.T) {
|
||||
opts := &HandleOptions{}
|
||||
log := xlogger.Nop()
|
||||
WithLog(log)(opts)
|
||||
if opts.log != log {
|
||||
t.Errorf("log not set")
|
||||
}
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// mockHop for testing
|
||||
// =============================================================================
|
||||
|
||||
type mockHop struct {
|
||||
nodes []*chain.Node
|
||||
}
|
||||
|
||||
func (m *mockHop) Select(ctx context.Context, opts ...hop.SelectOption) *chain.Node {
|
||||
if len(m.nodes) > 0 {
|
||||
return m.nodes[0]
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Ensure mockHop implements hop.Hop.
|
||||
var _ hop.Hop = (*mockHop)(nil)
|
||||
|
||||
// =============================================================================
|
||||
// HandleHTTP Edge Case Tests
|
||||
// =============================================================================
|
||||
|
||||
func TestHandleHTTP_DialError(t *testing.T) {
|
||||
h := &Sniffer{
|
||||
ReadTimeout: 5 * time.Second,
|
||||
Recorder: &noopRecorder{},
|
||||
}
|
||||
ro := &xrecorder.HandlerRecorderObject{}
|
||||
|
||||
clientConn, serverConn := net.Pipe()
|
||||
defer clientConn.Close()
|
||||
defer serverConn.Close()
|
||||
|
||||
errCh := make(chan error, 1)
|
||||
go func() {
|
||||
errCh <- h.HandleHTTP(context.Background(), serverConn,
|
||||
WithDial(func(ctx context.Context, network, address string) (net.Conn, error) {
|
||||
return nil, io.ErrUnexpectedEOF
|
||||
}),
|
||||
WithRecorderObject(ro),
|
||||
WithLog(xlogger.Nop()),
|
||||
)
|
||||
}()
|
||||
|
||||
// Write a valid HTTP request, then read the error response to avoid pipe deadlock.
|
||||
req, _ := http.NewRequest("GET", "http://example.com/", nil)
|
||||
req.Write(clientConn)
|
||||
|
||||
// Read the error response to unblock the server-side write.
|
||||
br := bufio.NewReader(clientConn)
|
||||
resp, respErr := http.ReadResponse(br, req)
|
||||
if respErr != nil {
|
||||
t.Errorf("reading error response: %v", respErr)
|
||||
} else {
|
||||
resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusServiceUnavailable {
|
||||
t.Errorf("status = %d, want %d", resp.StatusCode, http.StatusServiceUnavailable)
|
||||
}
|
||||
}
|
||||
|
||||
err := <-errCh
|
||||
if err == nil {
|
||||
t.Error("expected error from failed dial, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleHTTP_HTTP10Request(t *testing.T) {
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
w.Write([]byte("HTTP/1.0 OK"))
|
||||
}))
|
||||
defer upstream.Close()
|
||||
|
||||
h := &Sniffer{
|
||||
ReadTimeout: 5 * time.Second,
|
||||
Recorder: &noopRecorder{},
|
||||
}
|
||||
ro := &xrecorder.HandlerRecorderObject{}
|
||||
|
||||
clientConn, serverConn := net.Pipe()
|
||||
defer clientConn.Close()
|
||||
defer serverConn.Close()
|
||||
|
||||
errCh := make(chan error, 1)
|
||||
go func() {
|
||||
errCh <- h.HandleHTTP(context.Background(), serverConn,
|
||||
WithDial(func(ctx context.Context, network, address string) (net.Conn, error) {
|
||||
return net.Dial("tcp", upstream.Listener.Addr().String())
|
||||
}),
|
||||
WithRecorderObject(ro),
|
||||
WithLog(xlogger.Nop()),
|
||||
)
|
||||
}()
|
||||
|
||||
// Write an HTTP/1.0 request
|
||||
clientConn.Write([]byte("GET / HTTP/1.0\r\nHost: example.com\r\nConnection: keep-alive\r\n\r\n"))
|
||||
|
||||
resp, err := http.ReadResponse(bufio.NewReader(clientConn), nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
t.Errorf("status = %d, want %d", resp.StatusCode, http.StatusOK)
|
||||
}
|
||||
|
||||
clientConn.Close()
|
||||
<-errCh
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// WebSocket Frame Test
|
||||
// =============================================================================
|
||||
|
||||
func TestCopyWebsocketFrame_Basic(t *testing.T) {
|
||||
h := &Sniffer{
|
||||
Recorder: &noopRecorder{},
|
||||
RecorderOptions: &recorder.Options{HTTPBody: false},
|
||||
}
|
||||
|
||||
// Build a minimal WebSocket text frame: fin=true, opcode=text(1), masked, payload="hi"
|
||||
payload := []byte("hi")
|
||||
mask := []byte{0xAA, 0xBB, 0xCC, 0xDD}
|
||||
maskedPayload := make([]byte, len(payload))
|
||||
for i := range payload {
|
||||
maskedPayload[i] = payload[i] ^ mask[i%4]
|
||||
}
|
||||
|
||||
var frame bytes.Buffer
|
||||
frame.WriteByte(0x81) // FIN + text opcode
|
||||
frame.WriteByte(0x82) // MASK + len=2
|
||||
frame.Write(mask) // mask key
|
||||
frame.Write(maskedPayload) // masked payload
|
||||
|
||||
w := &bytes.Buffer{}
|
||||
ro := &xrecorder.HandlerRecorderObject{}
|
||||
|
||||
err := h.copyWebsocketFrame(w, &frame, &bytes.Buffer{}, "client", ro)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
if ro.Websocket == nil {
|
||||
t.Fatal("Websocket recorder object should be populated")
|
||||
}
|
||||
if ro.Websocket.OpCode != 1 {
|
||||
t.Errorf("opcode = %d, want 1 (text)", ro.Websocket.OpCode)
|
||||
}
|
||||
if ro.Websocket.Masked != true {
|
||||
t.Error("client frame should be marked as masked")
|
||||
}
|
||||
|
||||
// Client direction: InputBytes should be non-zero.
|
||||
if ro.InputBytes == 0 {
|
||||
t.Error("InputBytes should be non-zero for client direction")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCopyWebsocketFrame_WithBodyRecording(t *testing.T) {
|
||||
h := &Sniffer{
|
||||
Recorder: &noopRecorder{},
|
||||
RecorderOptions: &recorder.Options{HTTPBody: true, MaxBodySize: 1024},
|
||||
}
|
||||
|
||||
payload := []byte("hello ws")
|
||||
// Build a server frame (unmasked) with FIN + text opcode.
|
||||
var frame bytes.Buffer
|
||||
frame.WriteByte(0x81) // FIN + text
|
||||
frame.WriteByte(byte(len(payload))) // MASK=0, len
|
||||
// No mask for server frames
|
||||
frame.Write(payload)
|
||||
|
||||
w := &bytes.Buffer{}
|
||||
ro := &xrecorder.HandlerRecorderObject{}
|
||||
|
||||
err := h.copyWebsocketFrame(w, &frame, &bytes.Buffer{}, "server", ro)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
if ro.Websocket == nil {
|
||||
t.Fatal("Websocket recorder object should be populated")
|
||||
}
|
||||
if string(ro.Websocket.Payload) != "hello ws" {
|
||||
t.Errorf("payload = %q, want %q", ro.Websocket.Payload, "hello ws")
|
||||
}
|
||||
if ro.Websocket.Masked != false {
|
||||
t.Error("server frame should be marked as unmasked")
|
||||
}
|
||||
// Server frames: OutputBytes should be non-zero.
|
||||
if ro.OutputBytes == 0 {
|
||||
t.Error("OutputBytes should be non-zero for server direction")
|
||||
}
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// normalizeHost Edge Cases
|
||||
// =============================================================================
|
||||
|
||||
func TestNormalizeHost_EdgeCases(t *testing.T) {
|
||||
// IPv6 with brackets already has port
|
||||
got := normalizeHost("[::1]:8080", "443")
|
||||
if got != "[::1]:8080" {
|
||||
t.Errorf("got %q, want %q", got, "[::1]:8080")
|
||||
}
|
||||
|
||||
// Hostname with port
|
||||
got = normalizeHost("localhost:3000", "80")
|
||||
if got != "localhost:3000" {
|
||||
t.Errorf("got %q, want %q", got, "localhost:3000")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,320 @@
|
||||
package forwarder
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"time"
|
||||
|
||||
"github.com/go-gost/core/bypass"
|
||||
"github.com/go-gost/core/chain"
|
||||
"github.com/go-gost/core/hop"
|
||||
dissector "github.com/go-gost/tls-dissector"
|
||||
xbypass "github.com/go-gost/x/bypass"
|
||||
xio "github.com/go-gost/x/internal/io"
|
||||
xnet "github.com/go-gost/x/internal/net"
|
||||
"github.com/go-gost/x/internal/util/sniffing"
|
||||
tls_util "github.com/go-gost/x/internal/util/tls"
|
||||
xrecorder "github.com/go-gost/x/recorder"
|
||||
"crypto/x509"
|
||||
)
|
||||
|
||||
// HandleTLS sniffs and proxies a TLS connection. It parses the ClientHello
|
||||
// for SNI-based routing, optionally performs MITM TLS termination for HTTP
|
||||
// content inspection, and records TLS handshake metadata.
|
||||
func (h *Sniffer) HandleTLS(ctx context.Context, conn net.Conn, opts ...HandleOption) error {
|
||||
var ho HandleOptions
|
||||
for _, opt := range opts {
|
||||
opt(&ho)
|
||||
}
|
||||
ho.readTimeout = h.effectiveReadTimeout(&ho)
|
||||
|
||||
buf := new(bytes.Buffer)
|
||||
clientHello, err := dissector.ParseClientHello(io.TeeReader(conn, buf))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
ro := ho.recorderObject
|
||||
ro.TLS = &xrecorder.TLSRecorderObject{
|
||||
ServerName: clientHello.ServerName,
|
||||
ClientHello: hex.EncodeToString(buf.Bytes()),
|
||||
}
|
||||
if len(clientHello.SupportedProtos) > 0 {
|
||||
ro.TLS.Proto = clientHello.SupportedProtos[0]
|
||||
}
|
||||
|
||||
host := normalizeHost(clientHello.ServerName, "443")
|
||||
if host != "" {
|
||||
ro.Host = host
|
||||
|
||||
if ho.bypass != nil && ho.bypass.Contains(ctx, "tcp", host, bypass.WithService(ho.service)) {
|
||||
return xbypass.ErrBypass
|
||||
}
|
||||
}
|
||||
|
||||
node, cc, err := h.dialTLS(ctx, host, &ho)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer cc.Close()
|
||||
ho.node = node
|
||||
|
||||
log := ho.log.WithFields(map[string]any{"src": cc.LocalAddr().String(), "dst": cc.RemoteAddr().String()})
|
||||
log.Debugf("connected to node %s(%s)", node.Name, node.Addr)
|
||||
|
||||
ro.SrcAddr = cc.LocalAddr().String()
|
||||
ro.DstAddr = cc.RemoteAddr().String()
|
||||
|
||||
if h.Certificate != nil && h.PrivateKey != nil &&
|
||||
len(clientHello.SupportedProtos) > 0 && (clientHello.SupportedProtos[0] == "h2" || clientHello.SupportedProtos[0] == "http/1.1") {
|
||||
if host == "" {
|
||||
host = ro.Host
|
||||
}
|
||||
if h.MitmBypass == nil || !h.MitmBypass.Contains(ctx, "tcp", host) {
|
||||
return h.terminateTLS(ctx, xnet.NewReadWriteConn(io.MultiReader(buf, conn), conn, conn), cc, clientHello, &ho)
|
||||
}
|
||||
}
|
||||
|
||||
if _, err := buf.WriteTo(cc); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
xio.SetReadDeadline(cc, time.Now().Add(ho.readTimeout))
|
||||
serverHello, serverHelloErr := dissector.ParseServerHello(io.TeeReader(cc, buf))
|
||||
xio.SetReadDeadline(cc, time.Time{})
|
||||
|
||||
if serverHello != nil {
|
||||
ro.TLS.CipherSuite = tls_util.CipherSuite(serverHello.CipherSuite).String()
|
||||
ro.TLS.CompressionMethod = serverHello.CompressionMethod
|
||||
if serverHello.Proto != "" {
|
||||
ro.TLS.Proto = serverHello.Proto
|
||||
}
|
||||
if serverHello.Version > 0 {
|
||||
ro.TLS.Version = tls_util.Version(serverHello.Version).String()
|
||||
}
|
||||
}
|
||||
|
||||
if buf.Len() > 0 {
|
||||
ro.TLS.ServerHello = hex.EncodeToString(buf.Bytes())
|
||||
}
|
||||
|
||||
if _, err := buf.WriteTo(conn); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
log.Infof("%s <-> %s", ro.RemoteAddr, ro.Host)
|
||||
xnet.Pipe(ctx, conn, cc)
|
||||
log.WithFields(map[string]any{
|
||||
"duration": time.Since(ro.Time),
|
||||
}).Infof("%s >-< %s", ro.RemoteAddr, ro.Host)
|
||||
|
||||
// If server hello parsing failed, surface the error after proxy completes.
|
||||
if serverHelloErr != nil {
|
||||
return serverHelloErr
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// dialTLS selects a node and establishes a TLS connection.
|
||||
func (h *Sniffer) dialTLS(ctx context.Context, host string, ho *HandleOptions) (node *chain.Node, cc net.Conn, err error) {
|
||||
dial := ho.dial
|
||||
if dial == nil {
|
||||
dial = (&net.Dialer{}).DialContext
|
||||
}
|
||||
|
||||
if node = ho.node; node != nil {
|
||||
cc, err = dial(ctx, "tcp", node.Addr)
|
||||
return
|
||||
}
|
||||
|
||||
node = &chain.Node{}
|
||||
|
||||
ro := ho.recorderObject
|
||||
if ho.hop != nil {
|
||||
var clientIP net.IP
|
||||
if clientAddr, _ := net.ResolveTCPAddr("tcp", ro.ClientAddr); clientAddr != nil {
|
||||
clientIP = clientAddr.IP
|
||||
}
|
||||
|
||||
node = ho.hop.Select(ctx,
|
||||
hop.ClientIPSelectOption(clientIP),
|
||||
hop.HostSelectOption(host),
|
||||
hop.ProtocolSelectOption(sniffing.ProtoTLS),
|
||||
)
|
||||
}
|
||||
if node == nil {
|
||||
err = errors.New("node not available")
|
||||
return
|
||||
}
|
||||
if node.Addr == "" {
|
||||
node = &chain.Node{
|
||||
Name: node.Name,
|
||||
Addr: host,
|
||||
}
|
||||
}
|
||||
|
||||
addr := node.Addr
|
||||
network := "tcp"
|
||||
if opts := node.Options(); opts != nil {
|
||||
switch opts.Network {
|
||||
case "unix":
|
||||
network = opts.Network
|
||||
default:
|
||||
if _, _, splitErr := net.SplitHostPort(addr); splitErr != nil {
|
||||
addr += ":443"
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// No options — ensure port is present.
|
||||
if _, _, splitErr := net.SplitHostPort(addr); splitErr != nil {
|
||||
addr += ":443"
|
||||
}
|
||||
}
|
||||
ro.Host = addr
|
||||
|
||||
ho.log = ho.log.WithFields(map[string]any{
|
||||
"host": host,
|
||||
"node": node.Name,
|
||||
"dst": fmt.Sprintf("%s/%s", addr, network),
|
||||
})
|
||||
ho.log.Debugf("find node for host %s -> %s(%s)", host, node.Name, addr)
|
||||
|
||||
cc, err = dial(ctx, network, addr)
|
||||
if err != nil {
|
||||
if marker := node.Marker(); marker != nil {
|
||||
marker.Mark()
|
||||
}
|
||||
ho.log.Warnf("connect to node %s(%s) failed: %v", node.Name, node.Addr, err)
|
||||
return
|
||||
}
|
||||
|
||||
if marker := node.Marker(); marker != nil {
|
||||
marker.Reset()
|
||||
}
|
||||
|
||||
cc = tlsWrapConn(cc, node.Options().TLS)
|
||||
return
|
||||
}
|
||||
|
||||
// terminateTLS performs MITM TLS termination: handshakes with the upstream
|
||||
// server as a client, then with the downstream client as a server using a
|
||||
// dynamically generated certificate. The decrypted traffic is then handled
|
||||
// as HTTP.
|
||||
func (h *Sniffer) terminateTLS(ctx context.Context, conn, cc net.Conn, clientHello *dissector.ClientHelloInfo, ho *HandleOptions) error {
|
||||
ro := ho.recorderObject
|
||||
log := ho.log
|
||||
|
||||
nextProtos := clientHello.SupportedProtos
|
||||
if h.NegotiatedProtocol != "" {
|
||||
nextProtos = []string{h.NegotiatedProtocol}
|
||||
}
|
||||
|
||||
cfg := &tls.Config{
|
||||
ServerName: clientHello.ServerName,
|
||||
NextProtos: nextProtos,
|
||||
CipherSuites: clientHello.CipherSuites,
|
||||
}
|
||||
if cfg.ServerName == "" {
|
||||
cfg.InsecureSkipVerify = true
|
||||
}
|
||||
clientConn := tls.Client(cc, cfg)
|
||||
if err := clientConn.HandshakeContext(ctx); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
cs := clientConn.ConnectionState()
|
||||
ro.TLS.CipherSuite = tls_util.CipherSuite(cs.CipherSuite).String()
|
||||
ro.TLS.Proto = cs.NegotiatedProtocol
|
||||
ro.TLS.Version = tls_util.Version(cs.Version).String()
|
||||
|
||||
host := cfg.ServerName
|
||||
if host == "" {
|
||||
if len(cs.PeerCertificates) > 0 {
|
||||
host = cs.PeerCertificates[0].Subject.CommonName
|
||||
}
|
||||
if host == "" {
|
||||
host = ro.Host
|
||||
}
|
||||
}
|
||||
if splitHost, _, _ := net.SplitHostPort(host); splitHost != "" {
|
||||
host = splitHost
|
||||
}
|
||||
|
||||
negotiatedProtocol := cs.NegotiatedProtocol
|
||||
if h.NegotiatedProtocol != "" {
|
||||
negotiatedProtocol = h.NegotiatedProtocol
|
||||
}
|
||||
nextProtos = nil
|
||||
if negotiatedProtocol != "" {
|
||||
nextProtos = []string{negotiatedProtocol}
|
||||
}
|
||||
|
||||
// cache the tls server handshake record.
|
||||
wb := &bytes.Buffer{}
|
||||
conn = xnet.NewReadWriteConn(conn, io.MultiWriter(wb, conn), conn)
|
||||
|
||||
serverConn := tls.Server(conn, &tls.Config{
|
||||
NextProtos: nextProtos,
|
||||
GetCertificate: func(chi *tls.ClientHelloInfo) (*tls.Certificate, error) {
|
||||
certPool := h.CertPool
|
||||
if certPool == nil {
|
||||
certPool = DefaultCertPool
|
||||
}
|
||||
serverName := chi.ServerName
|
||||
if serverName == "" {
|
||||
serverName = host
|
||||
}
|
||||
cert, cerr := certPool.Get(serverName)
|
||||
if cert != nil {
|
||||
pool := x509.NewCertPool()
|
||||
pool.AddCert(h.Certificate)
|
||||
if _, cerr = cert.Verify(x509.VerifyOptions{
|
||||
DNSName: serverName,
|
||||
Roots: pool,
|
||||
}); cerr != nil {
|
||||
log.Warnf("verify cached certificate for %s: %v", serverName, cerr)
|
||||
cert = nil
|
||||
}
|
||||
}
|
||||
if cert == nil {
|
||||
cert, cerr = tls_util.GenerateCertificate(serverName, 7*24*time.Hour, h.Certificate, h.PrivateKey)
|
||||
certPool.Put(serverName, cert)
|
||||
}
|
||||
if cerr != nil {
|
||||
return nil, cerr
|
||||
}
|
||||
|
||||
return &tls.Certificate{
|
||||
Certificate: [][]byte{cert.Raw},
|
||||
PrivateKey: h.PrivateKey,
|
||||
}, nil
|
||||
},
|
||||
})
|
||||
handshakeErr := serverConn.HandshakeContext(ctx)
|
||||
if record, _ := dissector.ReadRecord(wb); record != nil {
|
||||
wb.Reset()
|
||||
record.WriteTo(wb)
|
||||
ro.TLS.ServerHello = hex.EncodeToString(wb.Bytes())
|
||||
}
|
||||
if handshakeErr != nil {
|
||||
return handshakeErr
|
||||
}
|
||||
|
||||
opts := []HandleOption{
|
||||
WithDial(func(ctx context.Context, network, address string) (net.Conn, error) {
|
||||
return clientConn, nil
|
||||
}),
|
||||
WithHTTPKeepalive(true),
|
||||
WithNode(ho.node),
|
||||
WithRecorderObject(ro),
|
||||
WithLog(log),
|
||||
}
|
||||
return h.HandleHTTP(ctx, serverConn, opts...)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,147 @@
|
||||
package forwarder
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"io"
|
||||
"math"
|
||||
"time"
|
||||
|
||||
"github.com/go-gost/core/logger"
|
||||
"github.com/go-gost/x/internal/util/sniffing"
|
||||
ws_util "github.com/go-gost/x/internal/util/ws"
|
||||
xrecorder "github.com/go-gost/x/recorder"
|
||||
"golang.org/x/time/rate"
|
||||
)
|
||||
|
||||
// sniffingWebsocketFrame copies WebSocket frames between rw and cc while
|
||||
// recording frame metadata. It runs two goroutines for bidirectional copy
|
||||
// and waits for the first error.
|
||||
func (h *Sniffer) sniffingWebsocketFrame(ctx context.Context, rw, cc io.ReadWriteCloser, ro *xrecorder.HandlerRecorderObject, log logger.Logger) error {
|
||||
errc := make(chan error, 2)
|
||||
|
||||
sampleRate := h.WebsocketSampleRate
|
||||
if sampleRate == 0 {
|
||||
sampleRate = sniffing.DefaultSampleRate
|
||||
}
|
||||
if sampleRate < 0 {
|
||||
sampleRate = math.MaxFloat64
|
||||
}
|
||||
|
||||
go func() {
|
||||
ro2 := &xrecorder.HandlerRecorderObject{}
|
||||
*ro2 = *ro
|
||||
ro := ro2
|
||||
|
||||
limiter := rate.NewLimiter(rate.Limit(sampleRate), int(sampleRate))
|
||||
|
||||
buf := &bytes.Buffer{}
|
||||
for {
|
||||
start := time.Now()
|
||||
|
||||
if err := h.copyWebsocketFrame(cc, rw, buf, "client", ro); err != nil {
|
||||
errc <- err
|
||||
return
|
||||
}
|
||||
|
||||
if limiter.Allow() {
|
||||
ro.Duration = time.Since(start)
|
||||
ro.Time = time.Now()
|
||||
if err := ro.Record(ctx, h.Recorder); err != nil {
|
||||
log.Errorf("record: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
go func() {
|
||||
ro2 := &xrecorder.HandlerRecorderObject{}
|
||||
*ro2 = *ro
|
||||
ro := ro2
|
||||
|
||||
limiter := rate.NewLimiter(rate.Limit(sampleRate), int(sampleRate))
|
||||
|
||||
buf := &bytes.Buffer{}
|
||||
for {
|
||||
start := time.Now()
|
||||
|
||||
if err := h.copyWebsocketFrame(rw, cc, buf, "server", ro); err != nil {
|
||||
errc <- err
|
||||
return
|
||||
}
|
||||
|
||||
if limiter.Allow() {
|
||||
ro.Duration = time.Since(start)
|
||||
ro.Time = time.Now()
|
||||
if err := ro.Record(ctx, h.Recorder); err != nil {
|
||||
log.Errorf("record: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
err := <-errc
|
||||
// Close both connections to unblock the other goroutine.
|
||||
rw.Close()
|
||||
cc.Close()
|
||||
<-errc // wait for the other goroutine to finish
|
||||
return err
|
||||
}
|
||||
|
||||
// copyWebsocketFrame reads one WebSocket frame from r, records its metadata,
|
||||
// and writes it to w.
|
||||
func (h *Sniffer) copyWebsocketFrame(w io.Writer, r io.Reader, buf *bytes.Buffer, from string, ro *xrecorder.HandlerRecorderObject) (err error) {
|
||||
fr := ws_util.Frame{}
|
||||
if _, err = fr.ReadFrom(r); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
ws := &xrecorder.WebsocketRecorderObject{
|
||||
From: from,
|
||||
Fin: fr.Header.Fin,
|
||||
Rsv1: fr.Header.Rsv1,
|
||||
Rsv2: fr.Header.Rsv2,
|
||||
Rsv3: fr.Header.Rsv3,
|
||||
OpCode: int(fr.Header.OpCode),
|
||||
Masked: fr.Header.Masked,
|
||||
MaskKey: fr.Header.MaskKey,
|
||||
Length: fr.Header.PayloadLength,
|
||||
}
|
||||
|
||||
if bodySize := clampBodySize(h.RecorderOptions); bodySize > 0 {
|
||||
buf.Reset()
|
||||
if _, err := io.Copy(buf, io.LimitReader(fr.Data, int64(bodySize))); err != nil {
|
||||
return err
|
||||
}
|
||||
ws.Payload = buf.Bytes()
|
||||
|
||||
ro.Websocket = ws
|
||||
length := uint64(fr.Header.Length()) + uint64(fr.Header.PayloadLength)
|
||||
if from == "client" {
|
||||
ro.InputBytes = length
|
||||
ro.OutputBytes = 0
|
||||
} else {
|
||||
ro.InputBytes = 0
|
||||
ro.OutputBytes = length
|
||||
}
|
||||
|
||||
fr.Data = io.MultiReader(bytes.NewReader(buf.Bytes()), fr.Data)
|
||||
} else {
|
||||
ro.Websocket = ws
|
||||
length := uint64(fr.Header.Length()) + uint64(fr.Header.PayloadLength)
|
||||
if from == "client" {
|
||||
ro.InputBytes = length
|
||||
ro.OutputBytes = 0
|
||||
} else {
|
||||
ro.InputBytes = 0
|
||||
ro.OutputBytes = length
|
||||
}
|
||||
}
|
||||
|
||||
if _, err := fr.WriteTo(w); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user