notify/send_state_test.go

package notify

import (
	"context"
	"errors"
	"net"
	"os"
	"sync"
	"testing"
	"time"
)

func TestClientSendReturnsServiceShutdownWhenNotRunning(t *testing.T) {
	client := NewClient().(*ClientCommon)
	UseLegacySecurityClient(client)

	err := client.Send("notify", []byte("hello"))
	if !errors.Is(err, errServiceShutdown) {
		t.Fatalf("client Send error = %v, want %v", err, errServiceShutdown)
	}
}

func TestServerSendReturnsServiceShutdownWhenNotRunning(t *testing.T) {
	server := NewServer().(*ServerCommon)
	UseLegacySecurityServer(server)

	err := server.Send(nil, "notify", []byte("hello"))
	if !errors.Is(err, errServiceShutdown) {
		t.Fatalf("server Send error = %v, want %v", err, errServiceShutdown)
	}
}

func TestClientSendReturnsTransportDetachedWhenSessionAliveWithoutTransport(t *testing.T) {
	client := NewClient().(*ClientCommon)
	client.markSessionStarted()
	client.clearClientSessionRuntimeTransport()

	err := client.Send("notify", []byte("hello"))
	if !errors.Is(err, errTransportDetached) {
		t.Fatalf("client Send error = %v, want %v", err, errTransportDetached)
	}
}

func TestServerSendReturnsTransportDetachedWhenPeerTransportMissing(t *testing.T) {
	server := NewServer().(*ServerCommon)
	server.markSessionStarted()

	err := server.Send(&ClientConn{}, "notify", []byte("hello"))
	if !errors.Is(err, errTransportDetached) {
		t.Fatalf("server Send error = %v, want %v", err, errTransportDetached)
	}
}

func TestPendingWaitClosedErrorReturnsTransportDetachedWhenServiceStillRunning(t *testing.T) {
	if err := pendingWaitClosedError(nil); !errors.Is(err, errTransportDetached) {
		t.Fatalf("pendingWaitClosedError(nil) = %v, want %v", err, errTransportDetached)
	}
	openCh := make(chan struct{})
	if err := pendingWaitClosedError(openCh); !errors.Is(err, errTransportDetached) {
		t.Fatalf("pendingWaitClosedError(open) = %v, want %v", err, errTransportDetached)
	}
}

func TestPendingWaitClosedErrorReturnsServiceShutdownWhenStopped(t *testing.T) {
	stopCh := make(chan struct{})
	close(stopCh)
	if err := pendingWaitClosedError(stopCh); !errors.Is(err, errServiceShutdown) {
		t.Fatalf("pendingWaitClosedError(stopped) = %v, want %v", err, errServiceShutdown)
	}
}

func TestClientSendCtxReturnsContextCanceled(t *testing.T) {
	client := NewClient().(*ClientCommon)
	secret := []byte("0123456789abcdef0123456789abcdef")
	client.SetSecretKey(secret)
	server := newRunningPeerAttachServerForTest(t, func(server *ServerCommon) {
		server.SetSecretKey(secret)
	})
	left, right := net.Pipe()
	defer right.Close()
	bootstrapPeerAttachConnForTest(t, server, right)
	if err := client.ConnectByConn(left); err != nil {
		t.Fatalf("client ConnectByConn failed: %v", err)
	}
	defer func() {
		client.setByeFromServer(true)
		_ = client.Stop()
	}()

	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	_, err := client.SendCtx(ctx, "ctx-canceled", []byte("payload"))
	if !errors.Is(err, context.Canceled) {
		t.Fatalf("client SendCtx error = %v, want %v", err, context.Canceled)
	}
}

func TestServerSendCtxReturnsContextCanceled(t *testing.T) {
	client := NewClient().(*ClientCommon)
	secret := []byte("0123456789abcdef0123456789abcdef")
	client.SetSecretKey(secret)
	server := newRunningPeerAttachServerForTest(t, func(server *ServerCommon) {
		server.SetSecretKey(secret)
	})
	left, right := net.Pipe()
	defer right.Close()
	bootstrapPeerAttachConnForTest(t, server, right)
	if err := client.ConnectByConn(left); err != nil {
		t.Fatalf("client ConnectByConn failed: %v", err)
	}
	defer func() {
		client.setByeFromServer(true)
		_ = client.Stop()
	}()

	var logical *LogicalConn
	deadline := time.Now().Add(time.Second)
	for time.Now().Before(deadline) {
		logical = server.GetLogicalConn(client.peerIdentity)
		if logical != nil {
			break
		}
		time.Sleep(10 * time.Millisecond)
	}
	if logical == nil {
		t.Fatal("server logical conn not found")
	}

	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	_, err := server.SendCtxLogical(ctx, logical, "ctx-canceled", []byte("payload"))
	if !errors.Is(err, context.Canceled) {
		t.Fatalf("server SendCtxLogical error = %v, want %v", err, context.Canceled)
	}
}

func TestReplyWaitPreservesLegacyDeadlineAndCancelSentinels(t *testing.T) {
	client := NewClient().(*ClientCommon)
	secret := []byte("0123456789abcdef0123456789abcdef")
	client.SetSecretKey(secret)
	server := newRunningPeerAttachServerForTest(t, func(server *ServerCommon) {
		server.SetSecretKey(secret)
	})
	clientCancelCtx, clientCancel := context.WithCancel(context.Background())
	serverCancelCtx, serverCancel := context.WithCancel(context.Background())
	server.SetLink("client-wait-timeout", func(*Message) {})
	server.SetLink("client-ctx-cancel", func(*Message) { clientCancel() })
	client.SetLink("server-wait-timeout", func(*Message) {})
	client.SetLink("server-ctx-cancel", func(*Message) { serverCancel() })

	left, right := net.Pipe()
	t.Cleanup(func() { _ = right.Close() })
	bootstrapPeerAttachConnForTest(t, server, right)
	if err := client.ConnectByConn(left); err != nil {
		t.Fatalf("client ConnectByConn failed: %v", err)
	}
	t.Cleanup(func() {
		client.setByeFromServer(true)
		_ = client.Stop()
	})

	var logical *LogicalConn
	deadline := time.Now().Add(time.Second)
	for time.Now().Before(deadline) {
		logical = server.GetLogicalConn(client.peerIdentity)
		if logical != nil {
			break
		}
		time.Sleep(time.Millisecond)
	}
	if logical == nil {
		t.Fatal("server logical conn not found")
	}

	if _, err := client.SendWait("client-wait-timeout", []byte("payload"), 50*time.Millisecond); err != os.ErrDeadlineExceeded {
		t.Fatalf("client SendWait error=%#v, want exact os.ErrDeadlineExceeded", err)
	}

	if _, err := client.SendCtx(clientCancelCtx, "client-ctx-cancel", []byte("payload")); err != context.Canceled {
		t.Fatalf("client SendCtx error=%#v, want exact context.Canceled", err)
	}

	if _, err := server.SendWaitLogical(logical, "server-wait-timeout", []byte("payload"), 50*time.Millisecond); err != os.ErrDeadlineExceeded {
		t.Fatalf("server SendWaitLogical error=%#v, want exact os.ErrDeadlineExceeded", err)
	}

	if _, err := server.SendCtxLogical(serverCancelCtx, logical, "server-ctx-cancel", []byte("payload")); err != context.Canceled {
		t.Fatalf("server SendCtxLogical error=%#v, want exact context.Canceled", err)
	}
}

func TestSetLinkIsSafeDuringConcurrentDispatch(t *testing.T) {
	t.Run("client", func(t *testing.T) {
		client := NewClient().(*ClientCommon)
		testConcurrentHandlerUpdateAndDispatch(
			t,
			client.SetLink,
			client.SetDefaultLink,
			client.dispatchMsg,
		)
	})
	t.Run("server", func(t *testing.T) {
		server := NewServer().(*ServerCommon)
		testConcurrentHandlerUpdateAndDispatch(
			t,
			server.SetLink,
			server.SetDefaultLink,
			server.dispatchMsg,
		)
	})
}

func testConcurrentHandlerUpdateAndDispatch(
	t *testing.T,
	setLink func(string, func(*Message)),
	setDefault func(func(*Message)),
	dispatch func(Message),
) {
	t.Helper()
	const iterations = 1000
	start := make(chan struct{})
	var wg sync.WaitGroup
	wg.Add(2)
	go func() {
		defer wg.Done()
		<-start
		for i := 0; i < iterations; i++ {
			setLink("concurrent-handler", func(*Message) {})
			setDefault(func(*Message) {})
		}
	}()
	go func() {
		defer wg.Done()
		<-start
		for i := 0; i < iterations; i++ {
			dispatch(Message{TransferMsg: TransferMsg{
				Key:  "concurrent-handler",
				Type: MSG_ASYNC,
			}})
		}
	}()
	close(start)
	wg.Wait()
}