nezha_domains/service/rpc/receipt_gate_agentcompat_test.go

//go:build agentcompat

package rpc

import (
	"bufio"
	"fmt"
	"net"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
)

func installReceiptGateForTest(conn net.Conn) *receiptGate {
	activeReceiptGateMu.Lock()
	receiptGateGeneration++
	generation := receiptGateGeneration
	activeReceiptGateMu.Unlock()
	gate := newReceiptGate(conn, generation)
	activeReceiptGateMu.Lock()
	activeReceiptGate = gate
	activeReceiptGateMu.Unlock()
	return gate
}

func clearReceiptGateForTest() {
	activeReceiptGateMu.Lock()
	gate := activeReceiptGate
	activeReceiptGate = nil
	activeReceiptGateMu.Unlock()
	if gate != nil {
		gate.close()
	}
}

func TestReceiptGate_EOFResetsGate(t *testing.T) {
	// Given
	serverConn, clientConn := net.Pipe()
	defer clientConn.Close()
	installReceiptGateForTest(serverConn)
	defer clearReceiptGateForTest()
	gate := currentReceiptGate()
	require.NotNil(t, gate)
	go func() {
		reader := bufio.NewReader(clientConn)
		_, _ = reader.ReadString('\n')
		_ = clientConn.Close()
	}()

	// When
	err := notifyReceiptAccepted(7, "uuid", 1, 1)

	// Then
	require.Error(t, err)
	activeReceiptGateMu.RLock()
	active := activeReceiptGate
	activeReceiptGateMu.RUnlock()
	require.Nil(t, active)
}

func TestReceiptGate_ListenerAcceptsAndReplacesConnections(t *testing.T) {
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	require.NoError(t, err)
	defer CloseReceiptGate()
	SetReceiptGateListener(listener)

	oldClient, err := net.Dial("tcp", listener.Addr().String())
	require.NoError(t, err)
	oldReader := bufio.NewReader(oldClient)
	require.Equal(t, "ready\n", mustReadLine(t, oldReader))

	newClient, err := net.Dial("tcp", listener.Addr().String())
	require.NoError(t, err)
	defer newClient.Close()
	newReader := bufio.NewReader(newClient)
	require.Equal(t, "ready\n", mustReadLine(t, newReader))
	_ = oldClient.SetReadDeadline(time.Now().Add(time.Second))
	_, oldErr := oldReader.ReadString('\n')
	require.Error(t, oldErr)
}

func TestReceiptGate_CloseInterruptsHeldReadAndQueuedWrite(t *testing.T) {
	// Given
	serverConn, clientConn := net.Pipe()
	defer clientConn.Close()
	installReceiptGateForTest(serverConn)
	acceptedStarted := make(chan struct{})
	acceptedDone := make(chan error, 1)
	go func() {
		close(acceptedStarted)
		acceptedDone <- notifyReceiptAccepted(7, "uuid", 1, 1)
	}()
	<-acceptedStarted
	reader := bufio.NewReader(clientConn)
	require.Equal(t, "accepted 7 uuid "+fmt.Sprint(currentReceiptGate().generation)+" 1 1\n", mustReadLine(t, reader))

	infoStarted := make(chan struct{})
	infoDone := make(chan error, 1)
	go func() {
		close(infoStarted)
		infoDone <- notifyInfo2(9, "held")
	}()
	<-infoStarted

	// When
	CloseReceiptGate()

	// Then
	select {
	case <-acceptedDone:
	case <-time.After(time.Second):
		t.Fatal("held receipt read was not interrupted")
	}
	select {
	case <-infoDone:
	case <-time.After(time.Second):
		t.Fatal("queued notification write was not released")
	}
}

func mustReadLine(t *testing.T, reader *bufio.Reader) string {
	t.Helper()
	line, err := reader.ReadString('\n')
	require.NoError(t, err)
	return line
}

func TestReceiptGate_MalformedCommandResetsGate(t *testing.T) {
	// Given
	serverConn, clientConn := net.Pipe()
	defer clientConn.Close()
	installReceiptGateForTest(serverConn)
	defer clearReceiptGateForTest()
	go func() {
		reader := bufio.NewReader(clientConn)
		_, _ = reader.ReadString('\n')
		_, _ = clientConn.Write([]byte("hold\n"))
	}()

	// When
	err := notifyReceiptAccepted(7, "uuid", 1, 1)

	// Then
	require.EqualError(t, err, "receipt gate received unexpected command")
	activeReceiptGateMu.RLock()
	active := activeReceiptGate
	activeReceiptGateMu.RUnlock()
	require.Nil(t, active)
}

func TestReceiptGate_ReplacementClosesOldConnection(t *testing.T) {
	t.Run("replacement closes old connection", func(t *testing.T) {
		// Given
		oldServer, oldClient := net.Pipe()
		newServer, newClient := net.Pipe()
		t.Cleanup(func() { require.NoError(t, oldClient.Close()) })
		t.Cleanup(func() { require.NoError(t, newClient.Close()) })
		oldGate := installReceiptGateForTest(oldServer)
		t.Cleanup(oldGate.close)
		t.Cleanup(clearReceiptGateForTest)
		newGate := newReceiptGate(newServer, oldGate.generation+1)
		activeReceiptGateMu.Lock()
		activeReceiptGate = newGate
		activeReceiptGateMu.Unlock()
		oldDone := make(chan error, 1)
		go func() { oldDone <- oldGate.sendAccepted(7, "uuid", 1, 1) }()
		reader := bufio.NewReader(oldClient)
		_, _ = reader.ReadString('\n')

		// When
		oldGate.close()

		// Then
		select {
		case err := <-oldDone:
			require.Error(t, err)
		case <-time.After(time.Second):
			t.Fatal("old receipt gate remained blocked after replacement")
		}
	})

	require.Nil(t, currentReceiptGate())
}

func TestReceiptGate_Info2AndReceiptNotificationsSerialize(t *testing.T) {
	// Given
	serverConn, clientConn := net.Pipe()
	defer clientConn.Close()
	installReceiptGateForTest(serverConn)
	defer clearReceiptGateForTest()
	gate := currentReceiptGate()
	require.NotNil(t, gate)
	lines := make(chan string, 2)
	go func() {
		reader := bufio.NewReader(clientConn)
		line, err := reader.ReadString('\n')
		if err != nil {
			return
		}
		lines <- strings.TrimSpace(line)
		_, _ = clientConn.Write([]byte("release\n"))
		line, err = reader.ReadString('\n')
		if err == nil {
			lines <- strings.TrimSpace(line)
		}
	}()

	// When
	acceptedDone := make(chan error, 1)
	go func() { acceptedDone <- notifyReceiptAccepted(7, "uuid", 1, 1) }()
	select {
	case err := <-acceptedDone:
		require.NoError(t, err)
	case <-time.After(time.Second):
		require.NoError(t, <-acceptedDone)
	}
	require.NoError(t, notifyInfo2(7, "uuid"))

	// Then
	require.Equal(t, "accepted 7 uuid "+fmt.Sprint(gate.generation)+" 1 1", <-lines)
	require.Equal(t, "info2 "+fmt.Sprint(gate.generation)+" 7 uuid", <-lines)
}