x/handler/router/handler_test.go

package router

import (
	"bytes"
	"context"
	"testing"

	"github.com/go-gost/core/auth"
	"github.com/go-gost/core/handler"
	"github.com/go-gost/relay"
	"github.com/google/uuid"
)

// ---------------------------------------------------------------------------
// newInitdHandler — convenience constructor with minimal Init
// ---------------------------------------------------------------------------

func newInitdHandler(t *testing.T, opts ...handler.Option) *routerHandler {
	t.Helper()
	h := NewHandler(append([]handler.Option{
		handler.LoggerOption(&testLogger{}),
	}, opts...)...).(*routerHandler)
	if err := h.parseMetadata(testMD(nil)); err != nil {
		t.Fatalf("parseMetadata: %v", err)
	}
	h.id = uuid.New().String()
	h.log = &testLogger{}
	h.pool = NewConnectorPool(h.id)
	_, cancel := context.WithCancel(context.Background())
	h.cancel = cancel
	return h
}

// ---------------------------------------------------------------------------
// NewHandler tests
// ---------------------------------------------------------------------------

func TestNewHandler_Minimal(t *testing.T) {
	h := NewHandler()
	if h == nil {
		t.Fatal("handler is nil")
	}
	rh := h.(*routerHandler)
	if rh.options.Service != "" {
		t.Errorf("Service = %q, want empty", rh.options.Service)
	}
	if rh.sdCache == nil {
		t.Error("sdCache is nil")
	}
	if rh.routeCache == nil {
		t.Error("routeCache is nil")
	}
}

func TestNewHandler_WithOptions(t *testing.T) {
	h := NewHandler(
		handler.LoggerOption(&testLogger{}),
		handler.ServiceOption("test-svc"),
	)
	rh := h.(*routerHandler)
	if rh.options.Logger == nil {
		t.Error("Logger not set")
	}
	if rh.options.Service != "test-svc" {
		t.Errorf("Service = %q, want test-svc", rh.options.Service)
	}
}

// ---------------------------------------------------------------------------
// Init tests
// ---------------------------------------------------------------------------

func TestInit_Minimal(t *testing.T) {
	rh := NewHandler(handler.LoggerOption(&testLogger{})).(*routerHandler)
	if err := rh.Init(testMD(nil)); err != nil {
		t.Fatalf("Init: %v", err)
	}
	if rh.id == "" {
		t.Error("id is empty")
	}
	if rh.cancel == nil {
		t.Error("cancel is nil")
	}
	if rh.stats != nil {
		t.Error("stats should be nil (no observer)")
	}
	if rh.limiter != nil {
		t.Error("limiter should be nil (no limiter)")
	}
}

func TestInit_WithObserver(t *testing.T) {
	obs := newFakeObserver(10)
	rh := NewHandler(
		handler.LoggerOption(&testLogger{}),
		handler.ObserverOption(obs),
		handler.ServiceOption("test-svc"),
	).(*routerHandler)

	if err := rh.Init(testMD(nil)); err != nil {
		t.Fatalf("Init: %v", err)
	}
	if rh.stats == nil {
		t.Error("stats should be non-nil when observer is set")
	}
}

func TestInit_WithLimiter(t *testing.T) {
	rh := NewHandler(
		handler.LoggerOption(&testLogger{}),
		handler.TrafficLimiterOption(&mockTrafficLimiter{}),
	).(*routerHandler)

	if err := rh.Init(testMD(nil)); err != nil {
		t.Fatalf("Init: %v", err)
	}
	if rh.limiter == nil {
		t.Error("limiter should be non-nil when limiter is set")
	}
}

// ---------------------------------------------------------------------------
// Handle tests
// ---------------------------------------------------------------------------

func TestHandle_BadVersion(t *testing.T) {
	h := newInitdHandler(t)

	// Version2 causes relay.Request.ReadFrom to return an error because the
	// relay protocol only supports Version1. The Handle method returns the
	// read error directly — it never reaches the version check.
	req := relay.Request{
		Version: 0x02,
		Cmd:     relay.CmdAssociate,
	}
	var buf bytes.Buffer
	req.WriteTo(&buf)

	conn := &fakeConn{buf: buf.Bytes()}
	err := h.Handle(context.Background(), conn)
	if err == nil {
		t.Fatal("expected error for bad version")
	}

	// Connection must be closed via defer.
	if !conn.closed {
		t.Error("connection was not closed")
	}
}

// TestHandle_BadRequestVersion verifies that Handle rejects requests where
// the version byte has been tampered with. relay.Request.ReadFrom checks the
// version byte before the handler does — the handler's own version check
// (req.Version != relay.Version1) is a defensive secondary guard. Since
// ReadFrom catches bad versions first, Handle returns relay.ErrBadVersion
// from ReadFrom without writing a response.
func TestHandle_BadRequestVersion(t *testing.T) {
	h := newInitdHandler(t)

	// Build a valid relay request with Version1, then override the version byte.
	rid := relay.NewTunnelID([]byte("0123456789abcdef"))
	req := relay.Request{
		Version: relay.Version1,
		Cmd:     relay.CmdAssociate,
		Features: []relay.Feature{
			&relay.TunnelFeature{ID: rid},
			&relay.NetworkFeature{Network: relay.NetworkIP},
		},
	}
	var buf bytes.Buffer
	req.WriteTo(&buf)

	// Override the version byte to something that is not Version1.
	b := buf.Bytes()
	if len(b) > 0 {
		b[0] = 0x02
	}

	conn := &fakeConn{buf: b}
	err := h.Handle(context.Background(), conn)
	if err == nil {
		t.Fatal("expected error for bad request version")
	}

	// Connection must be closed via defer.
	if !conn.closed {
		t.Error("connection was not closed")
	}
}

func TestHandle_UnknownCmd(t *testing.T) {
	h := newInitdHandler(t)

	req := relay.Request{
		Version: relay.Version1,
		Cmd:     0xFF,
	}
	var buf bytes.Buffer
	req.WriteTo(&buf)

	conn := &fakeConn{buf: buf.Bytes()}
	err := h.Handle(context.Background(), conn)
	if err != ErrUnknownCmd {
		t.Errorf("err = %v, want ErrUnknownCmd", err)
	}

	resp := readRelayResponse(t, conn.writeBuf.Bytes())
	if resp.Status != relay.StatusBadRequest {
		t.Errorf("status = %d, want BadRequest", resp.Status)
	}
}

func TestHandle_AuthSuccess(t *testing.T) {
	h := newInitdHandler(t,
		handler.AutherOption(&mockAuther{
			authenticateFn: func(ctx context.Context, user, pass string, opts ...auth.Option) (string, bool) {
				return "client-1", true
			},
		}),
	)

	rid := relay.NewTunnelID([]byte("0123456789abcdef"))
	req := relay.Request{
		Version: relay.Version1,
		Cmd:     relay.CmdAssociate,
		Features: []relay.Feature{
			&relay.UserAuthFeature{Username: "user", Password: "pass"},
			&relay.TunnelFeature{ID: rid},
			&relay.NetworkFeature{Network: relay.NetworkIP},
		},
	}
	var buf bytes.Buffer
	req.WriteTo(&buf)

	conn := &fakeConn{buf: buf.Bytes()}
	err := h.Handle(context.Background(), conn)
	if err != nil {
		t.Fatalf("Handle: %v", err)
	}
}

func TestHandle_AuthFailure(t *testing.T) {
	h := newInitdHandler(t,
		handler.AutherOption(&mockAuther{
			authenticateFn: func(ctx context.Context, user, pass string, opts ...auth.Option) (string, bool) {
				return "", false
			},
		}),
	)

	rid := relay.NewTunnelID([]byte("0123456789abcdef"))
	req := relay.Request{
		Version: relay.Version1,
		Cmd:     relay.CmdAssociate,
		Features: []relay.Feature{
			&relay.UserAuthFeature{Username: "user", Password: "wrong"},
			&relay.TunnelFeature{ID: rid},
			&relay.NetworkFeature{Network: relay.NetworkIP},
		},
	}
	var buf bytes.Buffer
	req.WriteTo(&buf)

	conn := &fakeConn{buf: buf.Bytes()}
	err := h.Handle(context.Background(), conn)
	if err != ErrUnauthorized {
		t.Errorf("err = %v, want ErrUnauthorized", err)
	}

	resp := readRelayResponse(t, conn.writeBuf.Bytes())
	if resp.Status != relay.StatusUnauthorized {
		t.Errorf("status = %d, want Unauthorized", resp.Status)
	}
}

func TestHandle_AssociateValid(t *testing.T) {
	h := newInitdHandler(t)
	h.md.ingress = nil

	rid := relay.NewTunnelID([]byte("0123456789abcdef"))
	reqData := buildRelayAssociateRequest(t, "10.0.0.1:0", rid, "ip")

	conn := &fakeConn{buf: reqData}
	err := h.Handle(context.Background(), conn)
	if err != nil {
		t.Fatalf("Handle: %v", err)
	}
}

// ---------------------------------------------------------------------------
// Close tests
// ---------------------------------------------------------------------------

func TestClose(t *testing.T) {
	h := newInitdHandler(t)
	h.Close()
	// should not panic
}

func TestClose_Double(t *testing.T) {
	h := newInitdHandler(t)
	h.Close()
	h.Close() // should not panic
}