x/handler/relay/forward_test.go

package relay

import (
	"context"
	"net"
	"testing"

	"github.com/go-gost/core/chain"
	"github.com/go-gost/core/handler"
	"github.com/go-gost/core/observer"
	"github.com/go-gost/relay"
)

func newForwardHandler(t *testing.T, opts ...handler.Option) *relayHandler {
	t.Helper()
	// Default: provide a working router
	defaultOpts := []handler.Option{
		handler.LoggerOption(&testLogger{}),
		handler.RouterOption(&mockRouter{
			dialFn: func(ctx context.Context, network, address string, opts ...chain.DialOption) (net.Conn, error) {
				pr, pw := net.Pipe()
				pw.Close()
				return pr, nil
			},
		}),
	}
	rh := newInitdHandler(t, append(defaultOpts, opts...)...)
	return rh
}

func TestHandleForward_NoTarget(t *testing.T) {
	mh := &mockHop{
		selectFn: func(ctx context.Context) *chain.Node {
			return nil
		},
	}
	rh := newForwardHandler(t)
	rh.Forward(mh)

	fc := &fakeConn{buf: buildRelayConnectRequest(t, "example.com:80", "")}
	err := rh.Handle(context.Background(), fc)
	if err == nil || err.Error() != "target not available" {
		t.Errorf("err = %v, want target not available", err)
	}

	resp := readRelayResponse(t, fc.writeBuf.Bytes())
	if resp.Status != relay.StatusServiceUnavailable {
		t.Errorf("Status = %d, want %d", resp.Status, relay.StatusServiceUnavailable)
	}
}

func TestHandleForward_DialSucceeds(t *testing.T) {
	mh := &mockHop{
		selectFn: func(ctx context.Context) *chain.Node {
			return makeTestNode("example.com:80")
		},
	}
	rh := newForwardHandler(t)
	rh.md.noDelay = true
	rh.Forward(mh)

	fc := &fakeConn{buf: buildRelayConnectRequest(t, "example.com:80", "")}
	err := rh.Handle(context.Background(), fc)
	if err != nil {
		t.Fatalf("Handle: %v", err)
	}

	resp := readRelayResponse(t, fc.writeBuf.Bytes())
	if resp.Status != relay.StatusOK {
		t.Errorf("Status = %d, want %d", resp.Status, relay.StatusOK)
	}
}

func TestHandleForward_DialFails(t *testing.T) {
	mh := &mockHop{
		selectFn: func(ctx context.Context) *chain.Node {
			return makeTestNode("example.com:80")
		},
	}
	rh := newInitdHandler(t,
		handler.LoggerOption(&testLogger{}),
		handler.RouterOption(&mockRouter{
			dialFn: func(ctx context.Context, network, address string, opts ...chain.DialOption) (net.Conn, error) {
				return nil, net.UnknownNetworkError("mock failure")
			},
		}),
	)
	rh.Forward(mh)

	fc := &fakeConn{buf: buildRelayConnectRequest(t, "example.com:80", "")}
	err := rh.Handle(context.Background(), fc)
	if err == nil {
		t.Fatal("expected error")
	}

	resp := readRelayResponse(t, fc.writeBuf.Bytes())
	if resp.Status != relay.StatusHostUnreachable {
		t.Errorf("Status = %d, want %d", resp.Status, relay.StatusHostUnreachable)
	}
}

func TestHandleForward_NoDelay(t *testing.T) {
	mh := &mockHop{
		selectFn: func(ctx context.Context) *chain.Node {
			return makeTestNode("example.com:80")
		},
	}
	rh := newForwardHandler(t)
	rh.md.noDelay = true
	rh.Forward(mh)

	fc := &fakeConn{buf: buildRelayConnectRequest(t, "example.com:80", "")}
	err := rh.Handle(context.Background(), fc)
	if err != nil {
		t.Fatalf("Handle: %v", err)
	}

	resp := readRelayResponse(t, fc.writeBuf.Bytes())
	if resp.Status != relay.StatusOK {
		t.Errorf("Status = %d, want %d", resp.Status, relay.StatusOK)
	}
}

func TestHandleForward_UDP(t *testing.T) {
	mh := &mockHop{
		selectFn: func(ctx context.Context) *chain.Node {
			return makeTestNode("example.com:53")
		},
	}
	rh := newForwardHandler(t)
	rh.md.noDelay = true
	rh.Forward(mh)

	fc := &fakeConn{buf: buildRelayConnectRequest(t, "example.com:53", "udp")}
	err := rh.Handle(context.Background(), fc)
	if err != nil {
		t.Fatalf("Handle: %v", err)
	}

	resp := readRelayResponse(t, fc.writeBuf.Bytes())
	if resp.Status != relay.StatusOK {
		t.Errorf("Status = %d, want %d", resp.Status, relay.StatusOK)
	}
}

func TestHandleForward_NilMarker(t *testing.T) {
	// Create a node with no marker (nil Marker)
	mh := &mockHop{
		selectFn: func(ctx context.Context) *chain.Node {
			return chain.NewNode("test", "example.com:80")
		},
	}
	rh := newForwardHandler(t)
	rh.md.noDelay = true
	rh.Forward(mh)

	fc := &fakeConn{buf: buildRelayConnectRequest(t, "example.com:80", "")}
	err := rh.Handle(context.Background(), fc)
	if err != nil {
		t.Fatalf("Handle: %v", err)
	}
}

func TestHandleForward_WithObserver(t *testing.T) {
	obs := &fakeObserver{eventsCh: make(chan []observer.Event, 10)}
	mh := &mockHop{
		selectFn: func(ctx context.Context) *chain.Node {
			return makeTestNode("example.com:80")
		},
	}
	rh := newForwardHandler(t,
		handler.ObserverOption(obs),
		handler.ServiceOption("test-svc"),
	)
	rh.md.noDelay = true
	rh.Forward(mh)

	fc := &fakeConn{buf: buildRelayConnectRequest(t, "example.com:80", "")}
	err := rh.Handle(context.Background(), fc)
	if err != nil {
		t.Fatalf("Handle: %v", err)
	}
}