gost-x/chain/node_test.go

package chain

import (
	"context"
	"errors"
	"net"
	"sync/atomic"
	"testing"
	"time"

	"github.com/go-gost/core/chain"
	"github.com/go-gost/core/connector"

	xlogger "github.com/go-gost/x/logger"
)

// testTransport implements chain.Transporter with real TCP dials.
type testTransport struct {
	dialer  *net.Dialer
	options chain.TransportOptions
}

func (t *testTransport) Dial(ctx context.Context, addr string) (net.Conn, error) {
	return t.dialer.DialContext(ctx, "tcp", addr)
}

func (t *testTransport) Handshake(ctx context.Context, conn net.Conn) (net.Conn, error) {
	return conn, nil
}

func (t *testTransport) Connect(ctx context.Context, conn net.Conn, network, address string) (net.Conn, error) {
	return nil, errors.New("not implemented")
}

func (t *testTransport) Bind(ctx context.Context, conn net.Conn, network, address string, opts ...connector.BindOption) (net.Listener, error) {
	return nil, errors.New("not implemented")
}

func (t *testTransport) Multiplex() bool { return false }

func (t *testTransport) Options() *chain.TransportOptions { return &t.options }

func (t *testTransport) Copy() chain.Transporter { return t }

func TestNodeProbeTCP(t *testing.T) {
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	defer l.Close()

	// accept one connection, then close immediately (TCP probe only needs dial+handshake)
	go func() {
		for {
			conn, _ := l.Accept()
			if conn != nil {
				conn.Close()
			}
		}
	}()

	addr := l.Addr().String()
	tr := &testTransport{dialer: &net.Dialer{Timeout: 2 * time.Second}}
	node := chain.NewNode("test", addr, chain.TransportNodeOption(tr))

	cfg := &chain.ProbeConfig{
		Type:     chain.ProbeTypeTCP,
		Addr:     addr,
		Interval: 100 * time.Millisecond,
		Timeout:  2 * time.Second,
	}

	StartNodeProbe(node, cfg, xlogger.Nop())

	// Wait for first probe
	time.Sleep(300 * time.Millisecond)

	pr := node.ProbeResult()
	if pr == nil || !pr.Success {
		t.Fatalf("expected successful probe, got %+v", pr)
	}

	// Close listener → next probe fails
	l.Close()
	time.Sleep(300 * time.Millisecond)

	pr = node.ProbeResult()
	if pr == nil || pr.Success {
		t.Fatalf("expected failed probe after listener close, got %+v", pr)
	}
	if node.Marker().Count() == 0 {
		t.Error("expected marker count > 0 after failure")
	}

	node.Close()
}

func TestNodeProbeClose(t *testing.T) {
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	defer l.Close()

	go func() {
		for {
			conn, _ := l.Accept()
			if conn != nil {
				conn.Close()
			}
		}
	}()

	addr := l.Addr().String()
	tr := &testTransport{dialer: &net.Dialer{Timeout: 2 * time.Second}}
	node := chain.NewNode("test", addr, chain.TransportNodeOption(tr))

	cfg := &chain.ProbeConfig{
		Type:     chain.ProbeTypeTCP,
		Addr:     addr,
		Interval: 100 * time.Millisecond,
		Timeout:  2 * time.Second,
	}

	StartNodeProbe(node, cfg, xlogger.Nop())
	time.Sleep(150 * time.Millisecond)

	// Snapshot result before close
	lastSuccess := node.ProbeResult().Success
	if !lastSuccess {
		t.Fatal("expected initial success before close")
	}

	node.Close()

	// Capture marker count right after close
	countAfterClose := node.Marker().Count()

	// Wait longer than probe interval — no more updates after close
	time.Sleep(300 * time.Millisecond)

	// Marker count should not increase after close (no more probes running)
	if node.Marker().Count() != countAfterClose {
		t.Errorf("marker changed after close: was %d, now %d", countAfterClose, node.Marker().Count())
	}
}

func TestNodeProbeDisabled(t *testing.T) {
	node := chain.NewNode("no-probe", "addr")
	// nil config → no-op
	StartNodeProbe(node, nil, xlogger.Nop())
	if pr := node.ProbeResult(); pr != nil {
		t.Error("expected nil probe result for disabled probe")
	}
	// empty addr → no-op
	StartNodeProbe(node, &chain.ProbeConfig{}, xlogger.Nop())
	if pr := node.ProbeResult(); pr != nil {
		t.Error("expected nil probe result for empty addr")
	}
}

func TestNodeProbeCmd(t *testing.T) {
	// Cmd probe doesn't need a transport — StartNodeProbe skips the addr check.
	node := chain.NewNode("cmd-test", "")
	cfg := &chain.ProbeConfig{
		Type:     chain.ProbeTypeCmd,
		Command:  "true",
		Interval: 100 * time.Millisecond,
		Timeout:  2 * time.Second,
	}

	StartNodeProbe(node, cfg, xlogger.Nop())
	time.Sleep(300 * time.Millisecond)

	pr := node.ProbeResult()
	if pr == nil || !pr.Success {
		t.Fatalf("expected cmd probe success, got %+v", pr)
	}
	if node.Marker().Count() != 0 {
		t.Errorf("expected marker count 0 after success, got %d", node.Marker().Count())
	}

	// Switch to failing command
	cfg2 := &chain.ProbeConfig{
		Type:     chain.ProbeTypeCmd,
		Command:  "false",
		Interval: 100 * time.Millisecond,
		Timeout:  2 * time.Second,
	}
	node.Close() // stop old probe
	node2 := chain.NewNode("cmd-test-2", "")
	StartNodeProbe(node2, cfg2, xlogger.Nop())
	time.Sleep(300 * time.Millisecond)

	pr2 := node2.ProbeResult()
	if pr2 == nil || pr2.Success {
		t.Fatalf("expected cmd probe failure for exit 1, got %+v", pr2)
	}
	if node2.Marker().Count() == 0 {
		t.Error("expected marker count > 0 after cmd failure")
	}

	node2.Close()
}

func TestNodeProbeCmdDisabled(t *testing.T) {
	// cmd probe without command → no-op
	node := chain.NewNode("cmd-disabled", "")
	cfg := &chain.ProbeConfig{
		Type:    chain.ProbeTypeCmd,
		Command: "",
	}
	StartNodeProbe(node, cfg, xlogger.Nop())
	if pr := node.ProbeResult(); pr != nil {
		t.Error("expected nil probe result for empty command")
	}
}

func TestNodeProbeHTTP(t *testing.T) {
	// Start a small HTTP server
	mux := &httpTestMux{}
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	defer l.Close()

	go func() {
		for {
			conn, err := l.Accept()
			if err != nil {
				return
			}
			go mux.serve(conn)
		}
	}()

	addr := l.Addr().String()
	tr := &testTransport{dialer: &net.Dialer{Timeout: 2 * time.Second}}
	node := chain.NewNode("test-http", addr, chain.TransportNodeOption(tr))

	cfg := &chain.ProbeConfig{
		Type:           chain.ProbeTypeHTTP,
		Addr:           addr,
		Interval:       100 * time.Millisecond,
		Timeout:        2 * time.Second,
		HTTPPath:       "/health",
		ExpectedStatus: 200,
	}

	// Case 1: server returns 200
	mux.status.Store(200)
	StartNodeProbe(node, cfg, xlogger.Nop())
	time.Sleep(300 * time.Millisecond)

	pr := node.ProbeResult()
	if pr == nil || !pr.Success {
		t.Fatalf("expected HTTP 200 probe success, got %+v", pr)
	}

	// Case 2: server returns 500
	mux.status.Store(500)
	time.Sleep(300 * time.Millisecond)

	pr = node.ProbeResult()
	if pr == nil || pr.Success {
		t.Fatalf("expected HTTP 500 probe failure, got %+v", pr)
	}

	node.Close()
}

// httpTestMux is a minimal HTTP server that returns a configurable status.
type httpTestMux struct {
	status atomic.Int32
}

func (m *httpTestMux) serve(conn net.Conn) {
	defer conn.Close()
	buf := make([]byte, 1024)
	n, _ := conn.Read(buf)
	// Parse just enough to construct a response
	body := "HTTP/1.1 000 OK\r\nContent-Length: 0\r\nConnection: close\r\n\r\n"
	switch m.status.Load() {
	case 200:
		body = "HTTP/1.1 200 OK\r\nContent-Length: 0\r\nConnection: close\r\n\r\n"
	case 500:
		body = "HTTP/1.1 500 Internal Server Error\r\nContent-Length: 0\r\nConnection: close\r\n\r\n"
	}
	_ = n
	conn.Write([]byte(body))
}