refactor(forwarder): extract resolveHTTPNode and resolveTLSNode helpers, add tests
Pull node-resolution logic out of dial/dialTLS into pure resolveHTTPNode and resolveTLSNode helpers. This separates node selection (bypass, hop, fallback) from connection establishment, making both independently testable. Add 21 new tests covering bypass, hop selection, nil-hop, addr fallback, upgrade response type mismatch, websocket frames, response body rewriting, and h2 preface validation.
This commit is contained in:
@@ -83,7 +83,8 @@ func (h *Sniffer) HandleHTTP(ctx context.Context, conn net.Conn, opts ...HandleO
|
|||||||
}
|
}
|
||||||
defer cc.Close()
|
defer cc.Close()
|
||||||
|
|
||||||
log = log.WithFields(map[string]any{"src": cc.LocalAddr().String(), "dst": cc.RemoteAddr().String()})
|
ho.log = log.WithFields(map[string]any{"src": cc.LocalAddr().String(), "dst": cc.RemoteAddr().String()})
|
||||||
|
log = ho.log
|
||||||
log.Debugf("connected to node %s(%s)", node.Name, node.Addr)
|
log.Debugf("connected to node %s(%s)", node.Name, node.Addr)
|
||||||
|
|
||||||
ro.SrcAddr = cc.LocalAddr().String()
|
ro.SrcAddr = cc.LocalAddr().String()
|
||||||
@@ -117,55 +118,31 @@ func (h *Sniffer) HandleHTTP(ctx context.Context, conn net.Conn, opts ...HandleO
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// dial selects a node and establishes a connection for an HTTP request.
|
// resolveHTTPNode selects a node for an HTTP request by applying bypass rules
|
||||||
func (h *Sniffer) dial(ctx context.Context, conn net.Conn, req *http.Request, ho *HandleOptions) (node *chain.Node, cc net.Conn, err error) {
|
// and hop selection. It returns the selected node or an error response.
|
||||||
dial := ho.dial
|
func resolveHTTPNode(ctx context.Context, host string, req *http.Request, ho *HandleOptions) (node *chain.Node, res *http.Response, err error) {
|
||||||
if dial == nil {
|
res = &http.Response{
|
||||||
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,
|
ProtoMajor: req.ProtoMajor,
|
||||||
ProtoMinor: req.ProtoMinor,
|
ProtoMinor: req.ProtoMinor,
|
||||||
Header: http.Header{},
|
Header: http.Header{},
|
||||||
StatusCode: http.StatusServiceUnavailable,
|
StatusCode: http.StatusServiceUnavailable,
|
||||||
}
|
}
|
||||||
|
|
||||||
host := normalizeHost(req.Host, "80")
|
if ho.bypass != nil &&
|
||||||
if host != "" {
|
ho.bypass.Contains(ctx, "tcp", host,
|
||||||
ro.Host = host
|
bypass.WithService(ho.service),
|
||||||
ho.log = ho.log.WithFields(map[string]any{
|
bypass.WithPathOption(req.RequestURI)) {
|
||||||
"host": host,
|
ho.log.Debugf("bypass: %s %s", host, req.RequestURI)
|
||||||
})
|
res.StatusCode = http.StatusForbidden
|
||||||
|
return nil, res, xbypass.ErrBypass
|
||||||
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{}
|
node = &chain.Node{}
|
||||||
if ho.hop != nil {
|
if ho.hop != nil {
|
||||||
var clientIP net.IP
|
var clientIP net.IP
|
||||||
if clientAddr, _ := net.ResolveTCPAddr("tcp", ro.ClientAddr); clientAddr != nil {
|
if clientAddr, _ := net.ResolveTCPAddr("tcp", ho.recorderObject.ClientAddr); clientAddr != nil {
|
||||||
clientIP = clientAddr.IP
|
clientIP = clientAddr.IP
|
||||||
}
|
}
|
||||||
|
|
||||||
node = ho.hop.Select(ctx,
|
node = ho.hop.Select(ctx,
|
||||||
hop.ClientIPSelectOption(clientIP),
|
hop.ClientIPSelectOption(clientIP),
|
||||||
hop.ProtocolSelectOption(sniffing.ProtoHTTP),
|
hop.ProtocolSelectOption(sniffing.ProtoHTTP),
|
||||||
@@ -179,11 +156,7 @@ func (h *Sniffer) dial(ctx context.Context, conn net.Conn, req *http.Request, ho
|
|||||||
if node == nil {
|
if node == nil {
|
||||||
ho.log.Warnf("node for %s not found", host)
|
ho.log.Warnf("node for %s not found", host)
|
||||||
res.StatusCode = http.StatusBadGateway
|
res.StatusCode = http.StatusBadGateway
|
||||||
ro.HTTP.StatusCode = res.StatusCode
|
return nil, res, errors.New("node not available")
|
||||||
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 == "" {
|
if node.Addr == "" {
|
||||||
node = &chain.Node{
|
node = &chain.Node{
|
||||||
@@ -191,6 +164,46 @@ func (h *Sniffer) dial(ctx context.Context, conn net.Conn, req *http.Request, ho
|
|||||||
Addr: host,
|
Addr: host,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
return node, nil, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// dial selects a node, establishes a connection, and sends the request upstream.
|
||||||
|
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
|
||||||
|
host := normalizeHost(req.Host, "80")
|
||||||
|
if host != "" {
|
||||||
|
ro.Host = host
|
||||||
|
ho.log = ho.log.WithFields(map[string]any{
|
||||||
|
"host": host,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
node, res, resolveErr := resolveHTTPNode(ctx, host, req, ho)
|
||||||
|
if resolveErr != nil {
|
||||||
|
ro.HTTP.StatusCode = res.StatusCode
|
||||||
|
if werr := res.Write(conn); werr != nil {
|
||||||
|
ho.log.Warnf("write error response: %v", werr)
|
||||||
|
}
|
||||||
|
return nil, nil, resolveErr
|
||||||
|
}
|
||||||
|
|
||||||
|
// Prepare an error response for potential connection failures.
|
||||||
|
res = &http.Response{
|
||||||
|
ProtoMajor: req.ProtoMajor,
|
||||||
|
ProtoMinor: req.ProtoMinor,
|
||||||
|
Header: http.Header{},
|
||||||
|
StatusCode: http.StatusServiceUnavailable,
|
||||||
|
}
|
||||||
|
|
||||||
ro.Host = node.Addr
|
ro.Host = node.Addr
|
||||||
ho.log = ho.log.WithFields(map[string]any{
|
ho.log = ho.log.WithFields(map[string]any{
|
||||||
@@ -411,4 +424,3 @@ func (h *Sniffer) httpRoundTrip(ctx context.Context, rw, cc io.ReadWriteCloser,
|
|||||||
|
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -671,7 +671,7 @@ func TestCopyWebsocketFrame_WithBodyRecording(t *testing.T) {
|
|||||||
payload := []byte("hello ws")
|
payload := []byte("hello ws")
|
||||||
// Build a server frame (unmasked) with FIN + text opcode.
|
// Build a server frame (unmasked) with FIN + text opcode.
|
||||||
var frame bytes.Buffer
|
var frame bytes.Buffer
|
||||||
frame.WriteByte(0x81) // FIN + text
|
frame.WriteByte(0x81) // FIN + text
|
||||||
frame.WriteByte(byte(len(payload))) // MASK=0, len
|
frame.WriteByte(byte(len(payload))) // MASK=0, len
|
||||||
// No mask for server frames
|
// No mask for server frames
|
||||||
frame.Write(payload)
|
frame.Write(payload)
|
||||||
@@ -717,3 +717,542 @@ func TestNormalizeHost_EdgeCases(t *testing.T) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// =============================================================================
|
||||||
|
// resolveHTTPNode Tests
|
||||||
|
// =============================================================================
|
||||||
|
|
||||||
|
type configurableHop struct {
|
||||||
|
node *chain.Node
|
||||||
|
}
|
||||||
|
|
||||||
|
func (h *configurableHop) Select(_ context.Context, _ ...hop.SelectOption) *chain.Node {
|
||||||
|
return h.node
|
||||||
|
}
|
||||||
|
|
||||||
|
var _ hop.Hop = (*configurableHop)(nil)
|
||||||
|
|
||||||
|
func TestResolveHTTPNode_Bypass(t *testing.T) {
|
||||||
|
ho := &HandleOptions{
|
||||||
|
service: "test-svc",
|
||||||
|
bypass: &mockBypass{contains: true},
|
||||||
|
log: xlogger.Nop(),
|
||||||
|
recorderObject: &xrecorder.HandlerRecorderObject{},
|
||||||
|
}
|
||||||
|
req, _ := http.NewRequest("GET", "http://example.com/path", nil)
|
||||||
|
|
||||||
|
node, res, err := resolveHTTPNode(context.Background(), "example.com:80", req, ho)
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected bypass error, got nil")
|
||||||
|
}
|
||||||
|
if node != nil {
|
||||||
|
t.Error("node should be nil on bypass")
|
||||||
|
}
|
||||||
|
if res == nil {
|
||||||
|
t.Fatal("res should not be nil on bypass")
|
||||||
|
}
|
||||||
|
if res.StatusCode != http.StatusForbidden {
|
||||||
|
t.Errorf("status = %d, want %d", res.StatusCode, http.StatusForbidden)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestResolveHTTPNode_NoHop(t *testing.T) {
|
||||||
|
ho := &HandleOptions{
|
||||||
|
log: xlogger.Nop(),
|
||||||
|
recorderObject: &xrecorder.HandlerRecorderObject{},
|
||||||
|
}
|
||||||
|
req, _ := http.NewRequest("GET", "http://example.com/path", nil)
|
||||||
|
|
||||||
|
node, res, err := resolveHTTPNode(context.Background(), "example.com:80", req, ho)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
if res != nil {
|
||||||
|
t.Error("res should be nil on success")
|
||||||
|
}
|
||||||
|
if node == nil {
|
||||||
|
t.Fatal("node should not be nil")
|
||||||
|
}
|
||||||
|
if node.Addr != "example.com:80" {
|
||||||
|
t.Errorf("node.Addr = %q, want %q", node.Addr, "example.com:80")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestResolveHTTPNode_HopReturnsNil(t *testing.T) {
|
||||||
|
ho := &HandleOptions{
|
||||||
|
hop: &configurableHop{node: nil},
|
||||||
|
log: xlogger.Nop(),
|
||||||
|
recorderObject: &xrecorder.HandlerRecorderObject{},
|
||||||
|
}
|
||||||
|
req, _ := http.NewRequest("GET", "http://example.com/path", nil)
|
||||||
|
|
||||||
|
node, res, err := resolveHTTPNode(context.Background(), "example.com:80", req, ho)
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error, got nil")
|
||||||
|
}
|
||||||
|
if err.Error() != "node not available" {
|
||||||
|
t.Errorf("err = %q, want %q", err.Error(), "node not available")
|
||||||
|
}
|
||||||
|
if node != nil {
|
||||||
|
t.Error("node should be nil on error")
|
||||||
|
}
|
||||||
|
if res == nil {
|
||||||
|
t.Fatal("res should not be nil")
|
||||||
|
}
|
||||||
|
if res.StatusCode != http.StatusBadGateway {
|
||||||
|
t.Errorf("status = %d, want %d", res.StatusCode, http.StatusBadGateway)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestResolveHTTPNode_HopReturnsValidNode(t *testing.T) {
|
||||||
|
expectedNode := &chain.Node{Name: "backend", Addr: "10.0.0.1:8080"}
|
||||||
|
ho := &HandleOptions{
|
||||||
|
hop: &configurableHop{node: expectedNode},
|
||||||
|
log: xlogger.Nop(),
|
||||||
|
recorderObject: &xrecorder.HandlerRecorderObject{},
|
||||||
|
}
|
||||||
|
req, _ := http.NewRequest("GET", "http://example.com/path", nil)
|
||||||
|
|
||||||
|
node, res, err := resolveHTTPNode(context.Background(), "example.com:80", req, ho)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
if res != nil {
|
||||||
|
t.Error("res should be nil on success")
|
||||||
|
}
|
||||||
|
if node != expectedNode {
|
||||||
|
t.Errorf("node = %v, want %v", node, expectedNode)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestResolveHTTPNode_HopReturnsNodeWithoutAddr(t *testing.T) {
|
||||||
|
ho := &HandleOptions{
|
||||||
|
hop: &configurableHop{node: &chain.Node{Name: "backend"}},
|
||||||
|
log: xlogger.Nop(),
|
||||||
|
recorderObject: &xrecorder.HandlerRecorderObject{},
|
||||||
|
}
|
||||||
|
req, _ := http.NewRequest("GET", "http://example.com/path", nil)
|
||||||
|
|
||||||
|
node, res, err := resolveHTTPNode(context.Background(), "example.com:80", req, ho)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
if res != nil {
|
||||||
|
t.Error("res should be nil on success")
|
||||||
|
}
|
||||||
|
if node == nil {
|
||||||
|
t.Fatal("node should not be nil")
|
||||||
|
}
|
||||||
|
if node.Addr != "example.com:80" {
|
||||||
|
t.Errorf("node.Addr = %q, want %q (host fallback)", node.Addr, "example.com:80")
|
||||||
|
}
|
||||||
|
if node.Name != "backend" {
|
||||||
|
t.Errorf("node.Name = %q, want %q", node.Name, "backend")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// =============================================================================
|
||||||
|
// resolveTLSNode Tests
|
||||||
|
// =============================================================================
|
||||||
|
|
||||||
|
func TestResolveTLSNode_NoHop(t *testing.T) {
|
||||||
|
ho := &HandleOptions{
|
||||||
|
log: xlogger.Nop(),
|
||||||
|
recorderObject: &xrecorder.HandlerRecorderObject{},
|
||||||
|
}
|
||||||
|
|
||||||
|
node, err := resolveTLSNode(context.Background(), "example.com", ho)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
if node == nil {
|
||||||
|
t.Fatal("node should not be nil")
|
||||||
|
}
|
||||||
|
if node.Addr != "example.com" {
|
||||||
|
t.Errorf("node.Addr = %q, want %q", node.Addr, "example.com")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestResolveTLSNode_HopReturnsNil(t *testing.T) {
|
||||||
|
ho := &HandleOptions{
|
||||||
|
hop: &configurableHop{node: nil},
|
||||||
|
log: xlogger.Nop(),
|
||||||
|
recorderObject: &xrecorder.HandlerRecorderObject{},
|
||||||
|
}
|
||||||
|
|
||||||
|
node, err := resolveTLSNode(context.Background(), "example.com", ho)
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error, got nil")
|
||||||
|
}
|
||||||
|
if err.Error() != "node not available" {
|
||||||
|
t.Errorf("err = %q, want %q", err.Error(), "node not available")
|
||||||
|
}
|
||||||
|
if node != nil {
|
||||||
|
t.Error("node should be nil on error")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestResolveTLSNode_HopReturnsValidNode(t *testing.T) {
|
||||||
|
expectedNode := &chain.Node{Name: "backend", Addr: "10.0.0.1:443"}
|
||||||
|
ho := &HandleOptions{
|
||||||
|
hop: &configurableHop{node: expectedNode},
|
||||||
|
log: xlogger.Nop(),
|
||||||
|
recorderObject: &xrecorder.HandlerRecorderObject{},
|
||||||
|
}
|
||||||
|
|
||||||
|
node, err := resolveTLSNode(context.Background(), "example.com", ho)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
if node != expectedNode {
|
||||||
|
t.Errorf("node = %v, want %v", node, expectedNode)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestResolveTLSNode_HopReturnsNodeWithoutAddr(t *testing.T) {
|
||||||
|
ho := &HandleOptions{
|
||||||
|
hop: &configurableHop{node: &chain.Node{Name: "backend"}},
|
||||||
|
log: xlogger.Nop(),
|
||||||
|
recorderObject: &xrecorder.HandlerRecorderObject{},
|
||||||
|
}
|
||||||
|
|
||||||
|
node, err := resolveTLSNode(context.Background(), "example.com", ho)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
if node == nil {
|
||||||
|
t.Fatal("node should not be nil")
|
||||||
|
}
|
||||||
|
if node.Addr != "example.com" {
|
||||||
|
t.Errorf("node.Addr = %q, want %q (host fallback)", node.Addr, "example.com")
|
||||||
|
}
|
||||||
|
if node.Name != "backend" {
|
||||||
|
t.Errorf("node.Name = %q, want %q", node.Name, "backend")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// =============================================================================
|
||||||
|
// handleUpgradeResponse Tests
|
||||||
|
// =============================================================================
|
||||||
|
|
||||||
|
func TestHandleUpgradeResponse_TypeMismatch(t *testing.T) {
|
||||||
|
h := &Sniffer{}
|
||||||
|
req, _ := http.NewRequest("GET", "http://example.com/", nil)
|
||||||
|
req.Header.Set("Connection", "Upgrade")
|
||||||
|
req.Header.Set("Upgrade", "websocket")
|
||||||
|
|
||||||
|
res := &http.Response{
|
||||||
|
Header: http.Header{
|
||||||
|
"Connection": {"Upgrade"},
|
||||||
|
"Upgrade": {"h2c"},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
clientConn, serverConn := net.Pipe()
|
||||||
|
defer clientConn.Close()
|
||||||
|
defer serverConn.Close()
|
||||||
|
|
||||||
|
err := h.handleUpgradeResponse(context.Background(), serverConn, clientConn, req, res,
|
||||||
|
&xrecorder.HandlerRecorderObject{}, xlogger.Nop())
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected type mismatch error, got nil")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestHandleUpgradeResponse_NonWebsocket(t *testing.T) {
|
||||||
|
h := &Sniffer{Websocket: false}
|
||||||
|
|
||||||
|
req, _ := http.NewRequest("GET", "http://example.com/", nil)
|
||||||
|
req.Header.Set("Connection", "Upgrade")
|
||||||
|
req.Header.Set("Upgrade", "h2c")
|
||||||
|
|
||||||
|
res := &http.Response{
|
||||||
|
StatusCode: http.StatusSwitchingProtocols,
|
||||||
|
ProtoMajor: 1,
|
||||||
|
ProtoMinor: 1,
|
||||||
|
Header: http.Header{
|
||||||
|
"Connection": {"Upgrade"},
|
||||||
|
"Upgrade": {"h2c"},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
errCh := make(chan error, 1)
|
||||||
|
clientConn, serverConn := net.Pipe()
|
||||||
|
defer clientConn.Close()
|
||||||
|
defer serverConn.Close()
|
||||||
|
|
||||||
|
go func() {
|
||||||
|
errCh <- h.handleUpgradeResponse(context.Background(), serverConn, clientConn, req, res,
|
||||||
|
&xrecorder.HandlerRecorderObject{}, xlogger.Nop())
|
||||||
|
}()
|
||||||
|
|
||||||
|
br := bufio.NewReader(clientConn)
|
||||||
|
readResp, readErr := http.ReadResponse(br, req)
|
||||||
|
if readErr != nil {
|
||||||
|
t.Fatalf("reading upgrade response: %v", readErr)
|
||||||
|
}
|
||||||
|
if readResp.StatusCode != http.StatusSwitchingProtocols {
|
||||||
|
t.Errorf("status = %d, want %d", readResp.StatusCode, http.StatusSwitchingProtocols)
|
||||||
|
}
|
||||||
|
|
||||||
|
clientConn.Close()
|
||||||
|
serverConn.Close()
|
||||||
|
<-errCh
|
||||||
|
}
|
||||||
|
|
||||||
|
// =============================================================================
|
||||||
|
// rewriteRespBody Additional Tests
|
||||||
|
// =============================================================================
|
||||||
|
|
||||||
|
func TestRewriteRespBody_PatternReplacement(t *testing.T) {
|
||||||
|
resp := &http.Response{
|
||||||
|
Header: http.Header{"Content-Type": {"text/html"}},
|
||||||
|
ContentLength: 100,
|
||||||
|
Body: io.NopCloser(strings.NewReader("<title>Old Title</title>")),
|
||||||
|
}
|
||||||
|
err := rewriteRespBody(resp, chain.HTTPBodyRewriteSettings{
|
||||||
|
Pattern: regexp.MustCompile("Old"),
|
||||||
|
Type: "text/html",
|
||||||
|
Replacement: []byte("New"),
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
got, _ := io.ReadAll(resp.Body)
|
||||||
|
if string(got) != "<title>New Title</title>" {
|
||||||
|
t.Errorf("body = %q, want %q", string(got), "<title>New Title</title>")
|
||||||
|
}
|
||||||
|
if resp.ContentLength != int64(len("<title>New Title</title>")) {
|
||||||
|
t.Errorf("ContentLength = %d, want %d", resp.ContentLength, len("<title>New Title</title>"))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRewriteRespBody_MultipleRewrites(t *testing.T) {
|
||||||
|
resp := &http.Response{
|
||||||
|
Header: http.Header{"Content-Type": {"text/html"}},
|
||||||
|
ContentLength: 100,
|
||||||
|
Body: io.NopCloser(strings.NewReader("hello")),
|
||||||
|
}
|
||||||
|
err := rewriteRespBody(resp,
|
||||||
|
chain.HTTPBodyRewriteSettings{
|
||||||
|
Pattern: regexp.MustCompile("h.*o"),
|
||||||
|
Type: "text/html",
|
||||||
|
Replacement: []byte("first"),
|
||||||
|
},
|
||||||
|
chain.HTTPBodyRewriteSettings{
|
||||||
|
Pattern: regexp.MustCompile("first"),
|
||||||
|
Type: "text/html",
|
||||||
|
Replacement: []byte("second"),
|
||||||
|
},
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
got, _ := io.ReadAll(resp.Body)
|
||||||
|
if string(got) != "second" {
|
||||||
|
t.Errorf("body = %q, want %q", string(got), "second")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRewriteRespBody_NilPattern(t *testing.T) {
|
||||||
|
resp := &http.Response{
|
||||||
|
Header: http.Header{"Content-Type": {"text/html"}},
|
||||||
|
ContentLength: 100,
|
||||||
|
Body: io.NopCloser(strings.NewReader("original")),
|
||||||
|
}
|
||||||
|
err := rewriteRespBody(resp, chain.HTTPBodyRewriteSettings{
|
||||||
|
Pattern: nil,
|
||||||
|
Type: "*",
|
||||||
|
Replacement: []byte("replaced"),
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
got, _ := io.ReadAll(resp.Body)
|
||||||
|
if string(got) != "original" {
|
||||||
|
t.Errorf("body = %q, want %q (unchanged)", string(got), "original")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRewriteRespBody_EmptyContentTypeUsesDefault(t *testing.T) {
|
||||||
|
resp := &http.Response{
|
||||||
|
ContentLength: 100,
|
||||||
|
Body: io.NopCloser(strings.NewReader("original")),
|
||||||
|
}
|
||||||
|
err := rewriteRespBody(resp, chain.HTTPBodyRewriteSettings{
|
||||||
|
Pattern: regexp.MustCompile(".*"),
|
||||||
|
Type: "",
|
||||||
|
Replacement: []byte("rewritten"),
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
got, _ := io.ReadAll(resp.Body)
|
||||||
|
if string(got) != "rewritten" {
|
||||||
|
t.Errorf("body = %q, want %q", string(got), "rewritten")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRewriteRespBody_NegativeContentLength(t *testing.T) {
|
||||||
|
resp := &http.Response{
|
||||||
|
Header: http.Header{"Content-Type": {"text/html"}},
|
||||||
|
ContentLength: -1,
|
||||||
|
Body: io.NopCloser(strings.NewReader("original")),
|
||||||
|
}
|
||||||
|
err := rewriteRespBody(resp, chain.HTTPBodyRewriteSettings{
|
||||||
|
Pattern: regexp.MustCompile(".*"),
|
||||||
|
Type: "*",
|
||||||
|
Replacement: []byte("replaced"),
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
got, _ := io.ReadAll(resp.Body)
|
||||||
|
if string(got) != "original" {
|
||||||
|
t.Errorf("body = %q, want %q (unchanged)", string(got), "original")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// =============================================================================
|
||||||
|
// copyWebsocketFrame Additional Tests
|
||||||
|
// =============================================================================
|
||||||
|
|
||||||
|
func TestCopyWebsocketFrame_EmptyPayload(t *testing.T) {
|
||||||
|
h := &Sniffer{
|
||||||
|
Recorder: &noopRecorder{},
|
||||||
|
RecorderOptions: &recorder.Options{HTTPBody: true, MaxBodySize: 1024},
|
||||||
|
}
|
||||||
|
|
||||||
|
var frame bytes.Buffer
|
||||||
|
frame.WriteByte(0x81) // FIN + text
|
||||||
|
frame.WriteByte(0x00) // MASK=0, len=0
|
||||||
|
|
||||||
|
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 ro.Websocket.Length != 0 {
|
||||||
|
t.Errorf("payload length = %d, want 0", ro.Websocket.Length)
|
||||||
|
}
|
||||||
|
if len(ro.Websocket.Payload) != 0 {
|
||||||
|
t.Errorf("payload = %q, want empty", ro.Websocket.Payload)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestCopyWebsocketFrame_ServerWithoutBodyRecording(t *testing.T) {
|
||||||
|
h := &Sniffer{
|
||||||
|
Recorder: &noopRecorder{},
|
||||||
|
RecorderOptions: &recorder.Options{HTTPBody: false},
|
||||||
|
}
|
||||||
|
|
||||||
|
payload := []byte("data")
|
||||||
|
var frame bytes.Buffer
|
||||||
|
frame.WriteByte(0x82) // FIN + binary
|
||||||
|
frame.WriteByte(byte(len(payload))) // MASK=0, len=4
|
||||||
|
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 ro.Websocket.Payload != nil {
|
||||||
|
t.Error("payload should be nil when body recording is disabled")
|
||||||
|
}
|
||||||
|
if ro.Websocket.OpCode != 2 {
|
||||||
|
t.Errorf("opcode = %d, want 2 (binary)", ro.Websocket.OpCode)
|
||||||
|
}
|
||||||
|
if ro.OutputBytes == 0 {
|
||||||
|
t.Error("OutputBytes should be non-zero for server direction")
|
||||||
|
}
|
||||||
|
if ro.InputBytes != 0 {
|
||||||
|
t.Error("InputBytes should be zero for server direction")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestCopyWebsocketFrame_MaskedFrame(t *testing.T) {
|
||||||
|
h := &Sniffer{
|
||||||
|
Recorder: &noopRecorder{},
|
||||||
|
RecorderOptions: &recorder.Options{HTTPBody: false},
|
||||||
|
}
|
||||||
|
|
||||||
|
payload := []byte("secret")
|
||||||
|
mask := []byte{0x12, 0x34, 0x56, 0x78}
|
||||||
|
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
|
||||||
|
frame.WriteByte(0x80 | byte(len(payload))) // MASK=1, len
|
||||||
|
frame.Write(mask)
|
||||||
|
frame.Write(maskedPayload)
|
||||||
|
|
||||||
|
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.Masked {
|
||||||
|
t.Error("client frame should be masked")
|
||||||
|
}
|
||||||
|
written := w.Bytes()
|
||||||
|
if len(written) < 2+len(payload) {
|
||||||
|
t.Fatalf("written frame too short: %d bytes", len(written))
|
||||||
|
}
|
||||||
|
if written[0] != 0x81 {
|
||||||
|
t.Errorf("first byte = 0x%02x, want 0x81", written[0])
|
||||||
|
}
|
||||||
|
if written[1]&0x80 != 0x80 {
|
||||||
|
t.Error("mask bit should be preserved")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// =============================================================================
|
||||||
|
// serveH2 unit test (client preface)
|
||||||
|
// =============================================================================
|
||||||
|
|
||||||
|
func TestServeH2_InvalidPreface(t *testing.T) {
|
||||||
|
h := &Sniffer{
|
||||||
|
ReadTimeout: 5 * time.Second,
|
||||||
|
Recorder: &noopRecorder{},
|
||||||
|
}
|
||||||
|
ho := &HandleOptions{
|
||||||
|
log: xlogger.Nop(),
|
||||||
|
recorderObject: &xrecorder.HandlerRecorderObject{},
|
||||||
|
}
|
||||||
|
|
||||||
|
clientConn, serverConn := net.Pipe()
|
||||||
|
defer clientConn.Close()
|
||||||
|
defer serverConn.Close()
|
||||||
|
|
||||||
|
// net.Pipe is synchronous: Write blocks until the other side Reads.
|
||||||
|
// Write exactly 6 bytes (size of "SM\r\n\r\n" preface) then return.
|
||||||
|
go func() {
|
||||||
|
clientConn.Write([]byte("XXXXXX"))
|
||||||
|
}()
|
||||||
|
|
||||||
|
err := h.serveH2(context.Background(), serverConn, ho)
|
||||||
|
if err == nil {
|
||||||
|
t.Error("expected error from invalid h2 preface, got nil")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -11,6 +11,7 @@ import (
|
|||||||
"net"
|
"net"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
|
"crypto/x509"
|
||||||
"github.com/go-gost/core/bypass"
|
"github.com/go-gost/core/bypass"
|
||||||
"github.com/go-gost/core/chain"
|
"github.com/go-gost/core/chain"
|
||||||
"github.com/go-gost/core/hop"
|
"github.com/go-gost/core/hop"
|
||||||
@@ -21,7 +22,6 @@ import (
|
|||||||
"github.com/go-gost/x/internal/util/sniffing"
|
"github.com/go-gost/x/internal/util/sniffing"
|
||||||
tls_util "github.com/go-gost/x/internal/util/tls"
|
tls_util "github.com/go-gost/x/internal/util/tls"
|
||||||
xrecorder "github.com/go-gost/x/recorder"
|
xrecorder "github.com/go-gost/x/recorder"
|
||||||
"crypto/x509"
|
|
||||||
)
|
)
|
||||||
|
|
||||||
// HandleTLS sniffs and proxies a TLS connection. It parses the ClientHello
|
// HandleTLS sniffs and proxies a TLS connection. It parses the ClientHello
|
||||||
@@ -121,6 +121,32 @@ func (h *Sniffer) HandleTLS(ctx context.Context, conn net.Conn, opts ...HandleOp
|
|||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// resolveTLSNode selects a node for a TLS connection by applying hop selection.
|
||||||
|
func resolveTLSNode(ctx context.Context, host string, ho *HandleOptions) (node *chain.Node, err error) {
|
||||||
|
node = &chain.Node{}
|
||||||
|
if ho.hop != nil {
|
||||||
|
var clientIP net.IP
|
||||||
|
if clientAddr, _ := net.ResolveTCPAddr("tcp", ho.recorderObject.ClientAddr); clientAddr != nil {
|
||||||
|
clientIP = clientAddr.IP
|
||||||
|
}
|
||||||
|
node = ho.hop.Select(ctx,
|
||||||
|
hop.ClientIPSelectOption(clientIP),
|
||||||
|
hop.HostSelectOption(host),
|
||||||
|
hop.ProtocolSelectOption(sniffing.ProtoTLS),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
if node == nil {
|
||||||
|
return nil, errors.New("node not available")
|
||||||
|
}
|
||||||
|
if node.Addr == "" {
|
||||||
|
node = &chain.Node{
|
||||||
|
Name: node.Name,
|
||||||
|
Addr: host,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return node, nil
|
||||||
|
}
|
||||||
|
|
||||||
// dialTLS selects a node and establishes a TLS connection.
|
// 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) {
|
func (h *Sniffer) dialTLS(ctx context.Context, host string, ho *HandleOptions) (node *chain.Node, cc net.Conn, err error) {
|
||||||
dial := ho.dial
|
dial := ho.dial
|
||||||
@@ -133,32 +159,12 @@ func (h *Sniffer) dialTLS(ctx context.Context, host string, ho *HandleOptions) (
|
|||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
node = &chain.Node{}
|
node, err = resolveTLSNode(ctx, host, ho)
|
||||||
|
if err != nil {
|
||||||
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
|
return
|
||||||
}
|
}
|
||||||
if node.Addr == "" {
|
|
||||||
node = &chain.Node{
|
|
||||||
Name: node.Name,
|
|
||||||
Addr: host,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
|
ro := ho.recorderObject
|
||||||
addr := node.Addr
|
addr := node.Addr
|
||||||
network := "tcp"
|
network := "tcp"
|
||||||
if opts := node.Options(); opts != nil {
|
if opts := node.Options(); opts != nil {
|
||||||
@@ -171,7 +177,6 @@ func (h *Sniffer) dialTLS(ctx context.Context, host string, ho *HandleOptions) (
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
} else {
|
} else {
|
||||||
// No options — ensure port is present.
|
|
||||||
if _, _, splitErr := net.SplitHostPort(addr); splitErr != nil {
|
if _, _, splitErr := net.SplitHostPort(addr); splitErr != nil {
|
||||||
addr += ":443"
|
addr += ":443"
|
||||||
}
|
}
|
||||||
@@ -317,4 +322,3 @@ func (h *Sniffer) terminateTLS(ctx context.Context, conn, cc net.Conn, clientHel
|
|||||||
}
|
}
|
||||||
return h.HandleHTTP(ctx, serverConn, opts...)
|
return h.HandleHTTP(ctx, serverConn, opts...)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user