gost-x/service/service_test.go

package service

import (
	"bytes"
	"context"
	"errors"
	"io"
	"net"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/go-gost/core/admission"
	"github.com/go-gost/core/handler"
	"github.com/go-gost/core/listener"
	"github.com/go-gost/core/logger"
	"github.com/go-gost/core/metadata"
	"github.com/go-gost/core/observer"
	"github.com/go-gost/core/observer/stats"
	"github.com/go-gost/core/recorder"
	xctx "github.com/go-gost/x/ctx"
	connwrapper "github.com/go-gost/x/limiter/conn/wrapper"
	xlogger "github.com/go-gost/x/logger"
)

// --- Mocks ---

type mockListener struct {
	addr    net.Addr
	connCh  chan net.Conn // connections to deliver via Accept
	errCh   chan error    // errors to deliver via Accept
	closed  atomic.Bool
	closeMu sync.Mutex
	closeCh chan struct{}
}

func newMockListener() *mockListener {
	return &mockListener{
		addr:    &net.TCPAddr{IP: net.ParseIP("127.0.0.1"), Port: 0},
		connCh:  make(chan net.Conn, 10),
		errCh:   make(chan error, 10),
		closeCh: make(chan struct{}),
	}
}

func (l *mockListener) Init(metadata.Metadata) error { return nil }

func (l *mockListener) Accept() (net.Conn, error) {
	select {
	case c := <-l.connCh:
		return c, nil
	case e := <-l.errCh:
		return nil, e
	case <-l.closeCh:
		return nil, net.ErrClosed
	}
}

func (l *mockListener) Addr() net.Addr { return l.addr }

func (l *mockListener) Close() error {
	l.closeMu.Lock()
	defer l.closeMu.Unlock()
	if l.closed.CompareAndSwap(false, true) {
		close(l.closeCh)
	}
	return nil
}

type mockHandler struct {
	handleFn func(ctx context.Context, conn net.Conn) error
	closed   atomic.Bool
}

func newMockHandler(fn func(ctx context.Context, conn net.Conn) error) *mockHandler {
	return &mockHandler{handleFn: fn}
}

func (h *mockHandler) Init(metadata.Metadata) error { return nil }

func (h *mockHandler) Handle(ctx context.Context, conn net.Conn, opts ...handler.HandleOption) error {
	if h.handleFn != nil {
		return h.handleFn(ctx, conn)
	}
	return nil
}

func (h *mockHandler) Close() error {
	h.closed.Store(true)
	return nil
}

type mockConn struct {
	net.Conn
	remoteAddr net.Addr
	localAddr  net.Addr
	ctx        context.Context
	closed     atomic.Bool
}

func (c *mockConn) RemoteAddr() net.Addr { return c.remoteAddr }
func (c *mockConn) LocalAddr() net.Addr  { return c.localAddr }
func (c *mockConn) Close() error         { c.closed.Store(true); return nil }

type mockAdmission struct {
	allow bool
	calls atomic.Int32
}

func (a *mockAdmission) Admit(ctx context.Context, network, addr string, opts ...admission.Option) bool {
	a.calls.Add(1)
	return a.allow
}

type mockObserver struct {
	events []observer.Event
	mu     sync.Mutex
	err    error
	closed atomic.Bool
}

func (o *mockObserver) Observe(ctx context.Context, events []observer.Event, opts ...observer.Option) error {
	o.mu.Lock()
	defer o.mu.Unlock()
	o.events = append(o.events, events...)
	return o.err
}

func (o *mockObserver) Close() error {
	o.closed.Store(true)
	return nil
}

func (o *mockObserver) getEvents() []observer.Event {
	o.mu.Lock()
	defer o.mu.Unlock()
	out := make([]observer.Event, len(o.events))
	copy(out, o.events)
	return out
}

type mockRecorder struct {
	recorded [][]byte
	mu       sync.Mutex
	err      error
}

func (r *mockRecorder) Record(ctx context.Context, b []byte, opts ...recorder.RecordOption) error {
	r.mu.Lock()
	defer r.mu.Unlock()
	cp := make([]byte, len(b))
	copy(cp, b)
	r.recorded = append(r.recorded, cp)
	return r.err
}

func (r *mockRecorder) getRecorded() [][]byte {
	r.mu.Lock()
	defer r.mu.Unlock()
	return r.recorded
}

// --- NewService tests ---

func TestNewServiceDefaultsNopLogger(t *testing.T) {
	ln := newMockListener()
	h := newMockHandler(nil)
	svc := NewService("test", ln, h).(*defaultService)
	// Should not panic when logging
	svc.options.logger.Info("test")
}

func TestNewServiceSetsRunningState(t *testing.T) {
	ln := newMockListener()
	h := newMockHandler(nil)
	svc := NewService("test", ln, h).(*defaultService)
	if svc.Status().State() != StateRunning {
		t.Errorf("State() = %q, want %q", svc.Status().State(), StateRunning)
	}
}

func TestNewServiceRecordsCreateTime(t *testing.T) {
	before := time.Now()
	ln := newMockListener()
	h := newMockHandler(nil)
	svc := NewService("test", ln, h).(*defaultService)
	after := time.Now()

	ct := svc.Status().CreateTime()
	if ct.Before(before) || ct.After(after) {
		t.Errorf("CreateTime = %v, want between %v and %v", ct, before, after)
	}
}

// --- Addr tests ---

func TestAddr(t *testing.T) {
	addr := &net.TCPAddr{IP: net.ParseIP("10.0.0.1"), Port: 8080}
	ln := newMockListener()
	ln.addr = addr
	svc := NewService("test", ln, newMockHandler(nil)).(*defaultService)
	if got := svc.Addr(); got.String() != addr.String() {
		t.Errorf("Addr() = %v, want %v", got, addr)
	}
}

// --- Serve tests ---

func TestServeAcceptsConnection(t *testing.T) {
	ln := newMockListener()
	var handled atomic.Int32
	h := newMockHandler(func(ctx context.Context, conn net.Conn) error {
		handled.Add(1)
		return nil
	})
	svc := NewService("test", ln, h).(*defaultService)

	done := make(chan error, 1)
	go func() { done <- svc.Serve() }()

	ln.connCh <- &mockConn{
		remoteAddr: &net.TCPAddr{IP: net.ParseIP("192.168.1.1"), Port: 1234},
		localAddr:  ln.addr,
	}

	// Give handler goroutine time to run
	time.Sleep(100 * time.Millisecond)

	ln.Close()
	<-done

	if handled.Load() != 1 {
		t.Errorf("handled = %d, want 1", handled.Load())
	}
}

// gost#914: a single rejected connection must not end the accept loop. The conn
// limiter's wrapper signals an over-limit connection by returning a plain error
// from Accept, which the loop used to treat as fatal — so one rejection took the
// whole service down and every later connection was refused until a restart.
// The connection that was over the limit is refused; the service is not.
func TestServeSurvivesRejectedConnection(t *testing.T) {
	ln := newMockListener()
	var handled atomic.Int32
	h := newMockHandler(func(ctx context.Context, conn net.Conn) error {
		handled.Add(1)
		return nil
	})
	svc := NewService("test", ln, h).(*defaultService)

	done := make(chan error, 1)
	go func() { done <- svc.Serve() }()

	newAddr := func(port int) *net.TCPAddr {
		return &net.TCPAddr{IP: net.ParseIP("192.168.1.1"), Port: port}
	}

	// One connection is accepted, then the limiter rejects the next.
	ln.connCh <- &mockConn{remoteAddr: newAddr(1234), localAddr: ln.addr}
	ln.errCh <- connwrapper.ErrConnLimit

	// The loop must still be running: a later connection has to be served.
	ln.connCh <- &mockConn{remoteAddr: newAddr(1235), localAddr: ln.addr}

	select {
	case err := <-done:
		t.Fatalf("accept loop exited after a rejected connection: %v", err)
	case <-time.After(200 * time.Millisecond):
	}

	ln.Close()
	select {
	case <-done:
	case <-time.After(time.Second):
		t.Fatal("accept loop did not exit after the listener closed")
	}

	if got := handled.Load(); got != 2 {
		t.Errorf("handled = %d, want 2 (the connection after the rejection must be served)", got)
	}
}

func TestServeSetsReadyState(t *testing.T) {
	ln := newMockListener()
	svc := NewService("test", ln, newMockHandler(nil)).(*defaultService)

	done := make(chan error, 1)
	go func() { done <- svc.Serve() }()

	time.Sleep(50 * time.Millisecond)
	if got := svc.Status().State(); got != StateReady {
		t.Errorf("State() = %q, want %q", got, StateReady)
	}

	ln.Close()
	<-done
}

func TestServeReturnsOnListenerClose(t *testing.T) {
	ln := newMockListener()
	svc := NewService("test", ln, newMockHandler(nil)).(*defaultService)

	done := make(chan error, 1)
	go func() { done <- svc.Serve() }()

	ln.Close()
	err := <-done
	if !errors.Is(err, net.ErrClosed) {
		t.Errorf("Serve() = %v, want net.ErrClosed", err)
	}

	if svc.Status().State() != StateClosed {
		t.Errorf("State() = %q, want %q", svc.Status().State(), StateClosed)
	}
}

func TestServeTemporaryErrorRetry(t *testing.T) {
	ln := newMockListener()
	svc := NewService("test", ln, newMockHandler(nil)).(*defaultService)

	done := make(chan error, 1)
	go func() { done <- svc.Serve() }()

	// Send a temporary error, then close
	ln.errCh <- &tempErr{}
	time.Sleep(50 * time.Millisecond)

	if svc.Status().State() != StateFailed {
		t.Errorf("State() = %q, want %q", svc.Status().State(), StateFailed)
	}

	ln.Close()
	<-done
}

type tempErr struct{}

func (tempErr) Error() string   { return "temporary error" }
func (tempErr) Temporary() bool { return true }
func (tempErr) Timeout() bool   { return false }

func TestServeAdmissionDeny(t *testing.T) {
	ln := newMockListener()
	var handled atomic.Int32
	h := newMockHandler(func(ctx context.Context, conn net.Conn) error {
		handled.Add(1)
		return nil
	})
	adm := &mockAdmission{allow: false}
	svc := NewService("test", ln, h, AdmissionOption(adm)).(*defaultService)

	done := make(chan error, 1)
	go func() { done <- svc.Serve() }()

	ln.connCh <- &mockConn{
		remoteAddr: &net.TCPAddr{IP: net.ParseIP("10.0.0.1"), Port: 9999},
		localAddr:  ln.addr,
	}

	time.Sleep(100 * time.Millisecond)
	ln.Close()
	<-done

	if handled.Load() != 0 {
		t.Error("handler should not be called for denied connection")
	}
	if adm.calls.Load() != 1 {
		t.Errorf("admission calls = %d, want 1", adm.calls.Load())
	}
}

func TestServeAdmissionAllow(t *testing.T) {
	ln := newMockListener()
	var handled atomic.Int32
	h := newMockHandler(func(ctx context.Context, conn net.Conn) error {
		handled.Add(1)
		return nil
	})
	adm := &mockAdmission{allow: true}
	svc := NewService("test", ln, h, AdmissionOption(adm)).(*defaultService)

	done := make(chan error, 1)
	go func() { done <- svc.Serve() }()

	ln.connCh <- &mockConn{
		remoteAddr: &net.TCPAddr{IP: net.ParseIP("10.0.0.1"), Port: 9999},
		localAddr:  ln.addr,
	}

	time.Sleep(100 * time.Millisecond)
	ln.Close()
	<-done

	if handled.Load() != 1 {
		t.Errorf("handled = %d, want 1", handled.Load())
	}
}

func TestServeHandlerErrorIncrementsStats(t *testing.T) {
	ln := newMockListener()
	h := newMockHandler(func(ctx context.Context, conn net.Conn) error {
		return errors.New("handler error")
	})
	ms := newMockStats()
	svc := NewService("test", ln, h, StatsOption(ms)).(*defaultService)

	done := make(chan error, 1)
	go func() { done <- svc.Serve() }()

	ln.connCh <- &mockConn{
		remoteAddr: &net.TCPAddr{IP: net.ParseIP("10.0.0.1"), Port: 9999},
		localAddr:  ln.addr,
	}

	time.Sleep(100 * time.Millisecond)
	ln.Close()
	<-done

	if got := ms.Get(stats.KindTotalErrs); got != 1 {
		t.Errorf("TotalErrs = %d, want 1", got)
	}
}

func TestServeContextFromConn(t *testing.T) {
	ln := newMockListener()
	var gotSid string
	h := newMockHandler(func(ctx context.Context, conn net.Conn) error {
		if s := xctx.SidFromContext(ctx); s != "" {
			gotSid = string(s)
		}
		return nil
	})
	svc := NewService("test", ln, h).(*defaultService)

	done := make(chan error, 1)
	go func() { done <- svc.Serve() }()

	ln.connCh <- &mockConn{
		remoteAddr: &net.TCPAddr{IP: net.ParseIP("10.0.0.1"), Port: 9999},
		localAddr:  ln.addr,
	}

	time.Sleep(100 * time.Millisecond)
	ln.Close()
	<-done

	if gotSid == "" {
		t.Error("expected session ID in context, got empty")
	}
}

func TestServeCtxConnContext(t *testing.T) {
	ln := newMockListener()
	wantCtx := context.Background()
	var gotCtx context.Context
	h := newMockHandler(func(ctx context.Context, conn net.Conn) error {
		gotCtx = ctx
		return nil
	})
	svc := NewService("test", ln, h).(*defaultService)

	done := make(chan error, 1)
	go func() { done <- svc.Serve() }()

	ln.connCh <- &ctxConn{
		Conn: &mockConn{
			remoteAddr: &net.TCPAddr{IP: net.ParseIP("10.0.0.1"), Port: 9999},
			localAddr:  ln.addr,
		},
		ctx: wantCtx,
	}

	time.Sleep(100 * time.Millisecond)
	ln.Close()
	<-done

	if gotCtx == nil {
		t.Fatal("handler context is nil")
	}
	// The ctx from the conn should be used as the base
	if s := xctx.SidFromContext(gotCtx); s == "" {
		t.Error("expected session ID to be added to conn context")
	}
}

type ctxConn struct {
	net.Conn
	ctx context.Context
}

func (c *ctxConn) Context() context.Context { return c.ctx }

func TestServeRecordsClientAddress(t *testing.T) {
	ln := newMockListener()
	rec := &mockRecorder{}
	h := newMockHandler(nil)

	svc := NewService("test", ln, h,
		RecordersOption(recorder.RecorderObject{
			Recorder: rec,
			Record:   recorder.RecorderServiceClientAddress,
		}),
	).(*defaultService)

	done := make(chan error, 1)
	go func() { done <- svc.Serve() }()

	ln.connCh <- &mockConn{
		remoteAddr: &net.TCPAddr{IP: net.ParseIP("192.168.1.100"), Port: 54321},
		localAddr:  ln.addr,
	}

	time.Sleep(100 * time.Millisecond)
	ln.Close()
	<-done

	recorded := rec.getRecorded()
	if len(recorded) != 1 {
		t.Fatalf("recorded = %d entries, want 1", len(recorded))
	}
	if string(recorded[0]) != "192.168.1.100" {
		t.Errorf("recorded = %q, want %q", string(recorded[0]), "192.168.1.100")
	}
}

func TestServeRecordsClientAddressError(t *testing.T) {
	ln := newMockListener()
	rec := &mockRecorder{err: errors.New("record fail")}
	h := newMockHandler(nil)

	svc := NewService("test", ln, h,
		RecordersOption(recorder.RecorderObject{
			Recorder: rec,
			Record:   recorder.RecorderServiceClientAddress,
		}),
	).(*defaultService)

	done := make(chan error, 1)
	go func() { done <- svc.Serve() }()

	ln.connCh <- &mockConn{
		remoteAddr: &net.TCPAddr{IP: net.ParseIP("192.168.1.100"), Port: 54321},
		localAddr:  ln.addr,
	}

	time.Sleep(100 * time.Millisecond)
	ln.Close()
	<-done
	// Should not panic; error is just logged
}

func TestServeObserverStats(t *testing.T) {
	ln := newMockListener()
	obs := &mockObserver{}
	ms := newMockStats()
	h := newMockHandler(nil)

	svc := NewService("test", ln, h,
		ObserverOption(obs),
		StatsOption(ms),
		ObserverPeriodOption(50*time.Millisecond),
	).(*defaultService)

	done := make(chan error, 1)
	go func() { done <- svc.Serve() }()

	// Add some stats to trigger observation
	ms.Add(stats.KindTotalConns, 10)

	time.Sleep(200 * time.Millisecond)
	ln.Close()
	<-done

	events := obs.getEvents()
	if len(events) == 0 {
		t.Error("expected observer to receive events")
	}
}

// --- Close tests ---

func TestCloseClosesListener(t *testing.T) {
	ln := newMockListener()
	svc := NewService("test", ln, newMockHandler(nil))

	if err := svc.Close(); err != nil {
		t.Errorf("Close() = %v", err)
	}
	if !ln.closed.Load() {
		t.Error("listener was not closed")
	}
}

func TestCloseClosesHandler(t *testing.T) {
	ln := newMockListener()
	h := newMockHandler(nil)
	svc := NewService("test", ln, h).(*defaultService)

	if err := svc.Close(); err != nil {
		t.Errorf("Close() = %v", err)
	}
	if !h.closed.Load() {
		t.Error("handler was not closed")
	}
}

func TestCloseClosesObserver(t *testing.T) {
	ln := newMockListener()
	obs := &mockObserver{}
	svc := NewService("test", ln, newMockHandler(nil), ObserverOption(obs)).(*defaultService)

	if err := svc.Close(); err != nil {
		t.Errorf("Close() = %v", err)
	}
	if !obs.closed.Load() {
		t.Error("observer was not closed")
	}
}

func TestCloseJoinsErrors(t *testing.T) {
	ln := newMockListener()
	h := &closeErrHandler{err: errors.New("handler close error")}
	svc := NewService("test", ln, h).(*defaultService)

	err := svc.Close()
	if err == nil {
		t.Fatal("Close() = nil, want error")
	}
	if err.Error() != "handler close error" {
		t.Errorf("Close() = %v, want 'handler close error'", err)
	}
}

type closeErrHandler struct {
	err error
}

func (h *closeErrHandler) Init(metadata.Metadata) error { return nil }
func (h *closeErrHandler) Handle(context.Context, net.Conn, ...handler.HandleOption) error {
	return nil
}
func (h *closeErrHandler) Close() error { return h.err }

type mockCloser struct {
	closed atomic.Bool
}

func (c *mockCloser) Close() error {
	c.closed.Store(true)
	return nil
}

func TestCloseClosesClosers(t *testing.T) {
	ln := newMockListener()
	c1, c2 := &mockCloser{}, &mockCloser{}
	svc := NewService("test", ln, newMockHandler(nil), ClosersOption(c1, c2)).(*defaultService)

	if err := svc.Close(); err != nil {
		t.Errorf("Close() = %v", err)
	}
	if !c1.closed.Load() || !c2.closed.Load() {
		t.Error("closers were not closed")
	}
}

// --- Status tests ---

func TestStatusReturnsStatus(t *testing.T) {
	ln := newMockListener()
	svc := NewService("test", ln, newMockHandler(nil)).(*defaultService)
	st := svc.Status()
	if st == nil {
		t.Fatal("Status() = nil")
	}
	if st.State() != StateRunning {
		t.Errorf("State() = %q, want %q", st.State(), StateRunning)
	}
}

// --- ServiceEvent tests ---

func TestServiceEventType(t *testing.T) {
	ev := ServiceEvent{Kind: "service", Service: "test", State: StateReady}
	if ev.Type() != observer.EventStatus {
		t.Errorf("Type() = %v, want %v", ev.Type(), observer.EventStatus)
	}
}

// --- Integration: multiple connections ---

func TestServeMultipleConnections(t *testing.T) {
	ln := newMockListener()
	var handled atomic.Int32
	h := newMockHandler(func(ctx context.Context, conn net.Conn) error {
		handled.Add(1)
		return nil
	})
	svc := NewService("test", ln, h).(*defaultService)

	done := make(chan error, 1)
	go func() { done <- svc.Serve() }()

	for i := 0; i < 5; i++ {
		ln.connCh <- &mockConn{
			remoteAddr: &net.TCPAddr{IP: net.ParseIP("10.0.0.1"), Port: 1000 + i},
			localAddr:  ln.addr,
		}
	}

	time.Sleep(200 * time.Millisecond)
	ln.Close()
	<-done

	if got := handled.Load(); got != 5 {
		t.Errorf("handled = %d, want 5", got)
	}
}

// --- execCmd tests ---

func TestExecCmdInvalidCommand(t *testing.T) {
	svc := NewService("test", newMockListener(), newMockHandler(nil)).(*defaultService)
	err := svc.execCmd("")
	if err == nil {
		t.Error("execCmd('') = nil, want error")
	}
}

// --- Interface assertions ---

var _ io.Closer = (*closeErrHandler)(nil)
var _ listener.Listener = (*mockListener)(nil)
var _ handler.Handler = (*mockHandler)(nil)
var _ handler.Handler = (*closeErrHandler)(nil)
var _ admission.Admission = (*mockAdmission)(nil)
var _ observer.Observer = (*mockObserver)(nil)
var _ stats.Stats = (*mockStats)(nil)
var _ recorder.Recorder = (*mockRecorder)(nil)
var _ xctx.Context = (*ctxConn)(nil)

// TestServeDoesNotLogAClosedListenerAsAnError pins the classification a
// deliberate stop depends on: a closed listener answers Accept with core's
// listener.ErrClosed, which is how stopping a service ends, not a failure.
// net.ErrClosed alone was not enough — that sentinel never matched it, so every
// stop reached the log at error level and read like a crash.
func TestServeDoesNotLogAClosedListenerAsAnError(t *testing.T) {
	var buf bytes.Buffer
	ln := newMockListener()
	ln.errCh <- listener.ErrClosed

	svc := NewService("test", ln, newMockHandler(nil),
		LoggerOption(xlogger.NewLogger(
			xlogger.OutputOption(&buf),
			xlogger.LevelOption(logger.ErrorLevel),
		)))

	// The error still comes back to the caller: classification changes how it
	// is reported, not what the service returns.
	if err := svc.Serve(); !errors.Is(err, listener.ErrClosed) {
		t.Fatalf("Serve() = %v, want the listener's error back", err)
	}
	if got := buf.String(); strings.Contains(got, "accept:") {
		t.Fatalf("a closed listener reached the log at error level:\n%s", got)
	}
}

// TestServeRetriesAClosedPipeQuietly pins both halves of the fix: the retry
// stays (a remote listener wraps an accept failure in core's AcceptError — which
// always claims Temporary() — to say "rebuild the transport", and a closed pipe
// is one of the failures that recovers that way), and the *log* stops reading
// like a hiccup. A stop arrives by the same path, so the closed error must not
// be logged at warning level.
func TestServeRetriesAClosedPipeQuietly(t *testing.T) {
	var buf bytes.Buffer
	ln := newMockListener()
	err := listener.NewAcceptError(io.ErrClosedPipe)
	ln.errCh <- err
	ln.errCh <- err
	ln.errCh <- errors.New("gave up after the retries")

	svc := NewService("test", ln, newMockHandler(nil),
		LoggerOption(xlogger.NewLogger(
			xlogger.OutputOption(&buf),
			xlogger.LevelOption(logger.WarnLevel),
		)))

	// Serve returns only after the queue is drained, which is what proves the
	// closed pipe was retried rather than ending the service.
	if got := svc.Serve(); got == nil || !strings.Contains(got.Error(), "gave up after") {
		t.Fatalf("Serve() = %v, want the error after the retries", got)
	}
	// The error that really ended it is logged, as it should be; the closed pipe
	// on the way must not be.
	if got := buf.String(); strings.Contains(got, "closed pipe") {
		t.Fatalf("a closed pipe was logged at warning level or above:\n%s", got)
	}
}