// Copyright 2026 The frp Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package server

import (
	"context"
	"errors"
	"math"
	"net"
	"net/http"
	"runtime"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/fatedier/golib/net/mux"
	"github.com/stretchr/testify/require"

	"github.com/fatedier/frp/pkg/auth"
	v1 "github.com/fatedier/frp/pkg/config/v1"
	"github.com/fatedier/frp/pkg/msg"
	plugin "github.com/fatedier/frp/pkg/plugin/server"
	"github.com/fatedier/frp/pkg/proto/wire"
	"github.com/fatedier/frp/pkg/util/util"
	"github.com/fatedier/frp/server/controller"
	"github.com/fatedier/frp/server/proxy"
	"github.com/fatedier/frp/server/registry"
	"github.com/fatedier/frp/server/visitor"
)

func TestWriteWithDeadlineTimesOutAndClearsDeadline(t *testing.T) {
	serverConn, clientConn := net.Pipe()
	defer serverConn.Close()
	defer clientConn.Close()

	err := writeWithDeadline(serverConn, 50*time.Millisecond, func() error {
		_, writeErr := serverConn.Write([]byte("x"))
		return writeErr
	})
	require.Error(t, err)

	var netErr net.Error
	require.True(t, errors.As(err, &netErr))
	require.True(t, netErr.Timeout())

	readCh := make(chan byte, 1)
	errCh := make(chan error, 1)
	go func() {
		buf := make([]byte, 1)
		if _, readErr := clientConn.Read(buf); readErr != nil {
			errCh <- readErr
			return
		}
		readCh <- buf[0]
	}()

	_, err = serverConn.Write([]byte("y"))
	require.NoError(t, err)

	select {
	case b := <-readCh:
		require.Equal(t, byte('y'), b)
	case err := <-errCh:
		require.NoError(t, err)
	case <-time.After(time.Second):
		t.Fatal("timed out waiting for write after deadline reset")
	}
}

func TestServiceAcceptConnectionTracksClientHelloPresence(t *testing.T) {
	for _, tc := range []struct {
		name               string
		clientHelloPresent bool
		offeredCodecs      []string
		expectedCodec      string
	}{
		{
			name: "absent Hello",
		},
		{
			name:               "present Hello with JSON fallback",
			clientHelloPresent: true,
		},
		{
			name:               "present Hello with binary codec",
			clientHelloPresent: true,
			offeredCodecs:      []string{wire.UDPPacketCodecBinary},
			expectedCodec:      wire.UDPPacketCodecBinary,
		},
	} {
		t.Run(tc.name, func(t *testing.T) {
			serverConn, clientConn := net.Pipe()
			defer serverConn.Close()
			defer clientConn.Close()

			clientErrCh := make(chan error, 1)
			go func() {
				if err := wire.WriteMagic(clientConn); err != nil {
					clientErrCh <- err
					return
				}
				wireConn := wire.NewConn(clientConn)
				if tc.clientHelloPresent {
					hello, err := wire.NewClientHello(wire.BootstrapInfo{})
					if err != nil {
						clientErrCh <- err
						return
					}
					hello.Capabilities.Message.UDPPacketCodecs = tc.offeredCodecs
					if err := wireConn.WriteJSONFrame(wire.FrameTypeClientHello, hello); err != nil {
						clientErrCh <- err
						return
					}
					var serverHello wire.ServerHello
					if err := wireConn.ReadJSONFrame(wire.FrameTypeServerHello, &serverHello); err != nil {
						clientErrCh <- err
						return
					}
				}
				clientErrCh <- msg.NewV2ReadWriterWithConn(wireConn).WriteMsg(&msg.NewWorkConn{RunID: "shared-run"})
			}()

			acceptedConn, err := (&Service{}).acceptConnection(t.Context(), serverConn)
			require.NoError(t, err)
			require.NoError(t, <-clientErrCh)
			require.Equal(t, tc.clientHelloPresent, acceptedConn.clientHelloPresent)
			require.Equal(t, tc.expectedCodec, acceptedConn.udpPacketCodec)
			require.IsType(t, &msg.NewWorkConn{}, acceptedConn.firstMsg)
			require.NoError(t, acceptedConn.conn.Close())
		})
	}
}

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

	sharedMux := mux.NewMux(listener)
	httpListener := sharedMux.ListenHTTP(1)
	muxServeErr := make(chan error, 1)
	go func() {
		muxServeErr <- sharedMux.Serve()
	}()

	newProtocols := func(http1, unencryptedHTTP2 bool) *http.Protocols {
		protocols := new(http.Protocols)
		protocols.SetHTTP1(http1)
		protocols.SetUnencryptedHTTP2(unencryptedHTTP2)
		return protocols
	}

	const handlerProtocolHeader = "X-Test-Handler-Protocol"
	httpServer := &http.Server{
		Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			w.Header().Set(handlerProtocolHeader, r.Proto)
			w.WriteHeader(http.StatusNoContent)
		}),
		ReadHeaderTimeout: time.Second,
		Protocols:         newProtocols(true, true),
	}
	httpServeErr := make(chan error, 1)
	go func() {
		httpServeErr <- httpServer.Serve(httpListener)
	}()
	t.Cleanup(func() {
		require.NoError(t, httpServer.Close())
		require.ErrorIs(t, waitForResult(t, httpServeErr, "shared HTTP server to stop"), http.ErrServerClosed)
		require.NoError(t, sharedMux.Close())
		require.ErrorIs(t, waitForResult(t, muxServeErr, "shared mux to stop"), net.ErrClosed)
	})

	for _, tc := range []struct {
		name             string
		http1            bool
		unencryptedHTTP2 bool
		expectedProtocol string
	}{
		{name: "HTTP/1.1", http1: true, expectedProtocol: "HTTP/1.1"},
		{name: "HTTP/2 prior knowledge", unencryptedHTTP2: true, expectedProtocol: "HTTP/2.0"},
	} {
		t.Run(tc.name, func(t *testing.T) {
			transport := &http.Transport{
				Protocols: newProtocols(tc.http1, tc.unencryptedHTTP2),
			}
			defer transport.CloseIdleConnections()
			client := &http.Client{Transport: transport, Timeout: 3 * time.Second}
			request, err := http.NewRequestWithContext(t.Context(), http.MethodGet, "http://"+listener.Addr().String()+"/", nil)
			require.NoError(t, err)
			response, err := client.Do(request)
			require.NoError(t, err)
			require.Equal(t, http.StatusNoContent, response.StatusCode)
			require.Equal(t, tc.expectedProtocol, response.Proto)
			require.Equal(t, tc.expectedProtocol, response.Header.Get(handlerProtocolHeader))
			require.NoError(t, response.Body.Close())
		})
	}
}

func TestServiceControlHandoffSkipsStalePendingGeneration(t *testing.T) {
	svr := newControlTestService(t)
	metrics := newCountingServerMetrics()
	metrics.closeEnter = make(chan struct{})
	metrics.closeResume = make(chan struct{})

	ctlA, connA, err := registerLifecycleTestControl(svr)
	require.NoError(t, err)
	ctlA.serverMetrics = metrics
	require.NoError(t, svr.completeControlLogin(ctlA, func() error { return nil }))
	waitForSignal(t, connA.readStarted, "A reader to start")

	require.NoError(t, ctlA.Close())
	waitForSignal(t, metrics.closeEnter, "A finalization barrier")

	type registerResult struct {
		ctl  *Control
		conn *deadlineReadConn
		err  error
	}
	resultB := make(chan registerResult, 1)
	go func() {
		ctl, conn, registerErr := registerLifecycleTestControl(svr)
		resultB <- registerResult{ctl: ctl, conn: conn, err: registerErr}
	}()
	ctlB := waitForDifferentCurrentControl(t, svr.ctlManager, "shared-run", ctlA)
	ctlB.serverMetrics = metrics

	resultC := make(chan registerResult, 1)
	go func() {
		ctl, conn, registerErr := registerLifecycleTestControl(svr)
		resultC <- registerResult{ctl: ctl, conn: conn, err: registerErr}
	}()
	ctlC := waitForDifferentCurrentControl(t, svr.ctlManager, "shared-run", ctlB)
	ctlC.serverMetrics = metrics
	waitForControlDone(t, ctlB)

	select {
	case result := <-resultB:
		t.Fatalf("B returned before A finalized: %v", result.err)
	default:
	}
	select {
	case result := <-resultC:
		t.Fatalf("C returned before A finalized: %v", result.err)
	default:
	}

	close(metrics.closeResume)
	waitForControlDone(t, ctlA)

	b := <-resultB
	require.Same(t, ctlB, b.ctl)
	require.ErrorIs(t, b.err, errControlReplaced)
	require.False(t, svr.ctlManager.Remove(ctlB))
	require.NoError(t, ctlB.Close())

	c := <-resultC
	require.NoError(t, c.err)
	require.Same(t, ctlC, c.ctl)
	_, ok := svr.ctlManager.GetByID("shared-run")
	require.False(t, ok)
	require.Same(t, ctlC, currentControlForTest(svr.ctlManager, "shared-run"))

	info, ok := svr.clientRegistry.GetByKey("client")
	require.True(t, ok)
	require.True(t, info.Online)
	require.Equal(t, uint64(ctlC.ID()), info.ControlID)

	var staleWrites atomic.Int64
	err = svr.completeControlLogin(ctlB, func() error {
		staleWrites.Add(1)
		return nil
	})
	require.ErrorIs(t, err, errControlReplaced)
	require.Equal(t, int64(0), staleWrites.Load())

	require.NoError(t, svr.completeControlLogin(ctlC, func() error { return nil }))
	waitForSignal(t, c.conn.readStarted, "C reader to start")
	current, ok := svr.ctlManager.GetByID("shared-run")
	require.True(t, ok)
	require.Same(t, ctlC, current)
	require.Equal(t, int64(2), metrics.newClients())
	require.Equal(t, int64(1), metrics.closedClients())

	require.NoError(t, ctlC.Close())
	waitForControlDone(t, ctlC)
	require.Equal(t, int64(2), metrics.newClients())
	require.Equal(t, int64(2), metrics.closedClients())
	_, ok = svr.ctlManager.GetByID("shared-run")
	require.False(t, ok)
}

func TestServiceLoginResponseSynchronizationIsScopedToRun(t *testing.T) {
	svr := newControlTestService(t)
	metrics := newCountingServerMetrics()
	ctlA, connA, err := registerLifecycleTestControl(svr)
	require.NoError(t, err)
	ctlA.serverMetrics = metrics

	writeEntered := make(chan struct{})
	resumeWrite := make(chan struct{})
	var resumeWriteOnce sync.Once
	resume := func() {
		resumeWriteOnce.Do(func() { close(resumeWrite) })
	}
	t.Cleanup(resume)
	writeCount := atomic.Int64{}
	loginDone := make(chan error, 1)
	go func() {
		loginDone <- svr.completeControlLogin(ctlA, func() error {
			close(writeEntered)
			<-resumeWrite
			writeCount.Add(1)
			return nil
		})
	}()
	waitForSignal(t, writeEntered, "A LoginResp write")

	runMu := currentRunGateForTest(svr.ctlManager, "shared-run")
	require.NotNil(t, runMu)
	if !svr.ctlManager.mu.TryLock() {
		t.Fatal("ControlManager mutex was held while LoginResp write was in progress")
	}
	svr.ctlManager.mu.Unlock()

	ctlB, connB := newLifecycleTestControl(t, "shared-run", "client", metrics)
	gateAvailable := make(chan bool)
	addDone := make(chan error, 1)
	go func() {
		if runMu.TryLock() {
			runMu.Unlock()
			gateAvailable <- true
		} else {
			gateAvailable <- false
		}
		addErr := svr.ctlManager.Add(ctlB)
		addDone <- addErr
	}()
	available := waitForResult(t, gateAvailable, "same-run replacement gate probe")
	require.False(t, available, "same-run gate was available to replacement during LoginResp write")
	select {
	case addErr := <-addDone:
		t.Fatalf("same-run replacement completed during LoginResp write: %v", addErr)
	case <-time.After(20 * time.Millisecond):
	}
	require.Same(t, ctlA, currentControlForTest(svr.ctlManager, "shared-run"))

	otherMetrics := newCountingServerMetrics()
	otherCtl, otherConn := newLifecycleTestControl(t, "other-run", "other-client", otherMetrics)
	type unrelatedResult struct {
		addErr      error
		active      bool
		activateErr error
		loginErr    error
		current     *Control
		found       bool
	}
	unrelatedDone := make(chan unrelatedResult, 1)
	go func() {
		result := unrelatedResult{}
		result.addErr = svr.ctlManager.Add(otherCtl)
		if result.addErr == nil {
			result.active, result.activateErr = svr.ctlManager.Activate(otherCtl)
		}
		if result.activateErr == nil && result.active {
			result.loginErr = svr.completeControlLogin(otherCtl, func() error { return nil })
		}
		result.current, result.found = svr.ctlManager.GetByID("other-run")
		unrelatedDone <- result
	}()
	result := waitForResult(t, unrelatedDone, "unrelated run lifecycle")
	require.NoError(t, result.addErr)
	require.NoError(t, result.activateErr)
	require.True(t, result.active)
	require.NoError(t, result.loginErr)
	require.True(t, result.found)
	require.Same(t, otherCtl, result.current)
	waitForSignal(t, otherConn.readStarted, "unrelated control reader to start")
	require.Equal(t, int64(1), otherMetrics.newClients())

	resume()
	require.NoError(t, waitForResult(t, loginDone, "LoginResp completion"))
	require.NoError(t, waitForResult(t, addDone, "replacement"))
	waitForControlDone(t, ctlA)
	require.Same(t, ctlB, currentControlForTest(svr.ctlManager, "shared-run"))
	require.Equal(t, int64(1), writeCount.Load())
	require.Equal(t, int64(1), metrics.newClients())
	require.Equal(t, int64(1), metrics.closedClients())
	require.Equal(t, []string{"deadline", "close"}, connA.eventsSnapshot())

	require.False(t, svr.ctlManager.Remove(ctlA))
	require.NoError(t, ctlA.Close())
	require.True(t, svr.ctlManager.Remove(ctlB))
	require.NoError(t, ctlB.Close())
	require.Equal(t, []string{"deadline", "close"}, connB.eventsSnapshot())

	require.NoError(t, otherCtl.Close())
	waitForControlDone(t, otherCtl)
	require.Equal(t, int64(1), otherMetrics.newClients())
	require.Equal(t, int64(1), otherMetrics.closedClients())
}

func TestServiceVisitorAdmissionSerializesReplacement(t *testing.T) {
	svr := newControlTestService(t)
	ctlA, controlConn, err := registerLifecycleTestControl(svr)
	require.NoError(t, err)
	ctlA.sessionCtx.LoginMsg.User = "old-user"
	require.NoError(t, svr.completeControlLogin(ctlA, func() error { return nil }))
	waitForSignal(t, controlConn.readStarted, "A reader to start")

	admissionEntered := make(chan struct{})
	resumeAdmission := make(chan struct{})
	var resumeOnce sync.Once
	resume := func() {
		resumeOnce.Do(func() { close(resumeAdmission) })
	}
	t.Cleanup(resume)
	type admissionResult struct {
		admitted       bool
		user           string
		wireProtocol   string
		udpPacketCodec string
		err            error
	}
	admissionDone := make(chan admissionResult, 1)
	go func() {
		result := admissionResult{}
		result.admitted, result.err = svr.ctlManager.admitVisitorByRunID("shared-run", func(user, wireProtocol, udpPacketCodec string) error {
			result.user = user
			result.wireProtocol = wireProtocol
			result.udpPacketCodec = udpPacketCodec
			close(admissionEntered)
			<-resumeAdmission
			return nil
		})
		admissionDone <- result
	}()
	waitForSignal(t, admissionEntered, "visitor admission callback")
	runMu := currentRunGateForTest(svr.ctlManager, "shared-run")
	require.NotNil(t, runMu)

	type registerResult struct {
		ctl *Control
		err error
	}
	gateAvailable := make(chan bool)
	replacementDone := make(chan registerResult, 1)
	go func() {
		if runMu.TryLock() {
			runMu.Unlock()
			gateAvailable <- true
		} else {
			gateAvailable <- false
		}
		ctl, _, registerErr := registerLifecycleTestControl(svr)
		replacementDone <- registerResult{ctl: ctl, err: registerErr}
	}()
	available := waitForResult(t, gateAvailable, "visitor replacement gate probe")
	require.False(t, available, "same-run gate was available during visitor admission")
	select {
	case result := <-replacementDone:
		t.Fatalf("replacement completed during visitor admission: %v", result.err)
	case <-time.After(20 * time.Millisecond):
	}
	require.Same(t, ctlA, currentControlForTest(svr.ctlManager, "shared-run"))

	resume()
	admission := waitForResult(t, admissionDone, "visitor admission")
	require.NoError(t, admission.err)
	require.True(t, admission.admitted)
	require.Equal(t, "old-user", admission.user)
	require.Equal(t, wire.ProtocolV1, admission.wireProtocol)
	require.Empty(t, admission.udpPacketCodec)
	replacement := waitForResult(t, replacementDone, "replacement")
	require.NoError(t, replacement.err)
	ctlB := replacement.ctl
	require.Same(t, ctlB, currentControlForTest(svr.ctlManager, "shared-run"))
	waitForControlDone(t, ctlA)
	require.True(t, svr.ctlManager.Remove(ctlB))
	require.NoError(t, ctlB.Close())
}

func TestServiceWorkConnRoutingRequiresCurrentRunningControl(t *testing.T) {
	svr := newControlTestService(t)
	ctl, controlConn, err := registerLifecycleTestControl(svr)
	require.NoError(t, err)

	pendingConn := newCountingCloseConn()
	pendingMsgConn := msg.NewConn(pendingConn, msg.NewV1ReadWriter(pendingConn))
	err = registerWorkConnAsCaller(svr, pendingMsgConn, &msg.NewWorkConn{RunID: "shared-run"}, wire.ProtocolV1, false)
	require.Error(t, err)
	require.Equal(t, int64(1), pendingConn.closeCount.Load())
	require.Len(t, ctl.workConnCh, 0)

	require.NoError(t, svr.completeControlLogin(ctl, func() error { return nil }))
	waitForSignal(t, controlConn.readStarted, "control reader to start")
	current, ok := svr.ctlManager.GetByID("shared-run")
	require.True(t, ok)
	require.Same(t, ctl, current)
	require.Len(t, ctl.workConnCh, 0)

	runningConn := newCountingCloseConn()
	runningMsgConn := msg.NewConn(runningConn, msg.NewV1ReadWriter(runningConn))
	require.NoError(t, svr.RegisterWorkConn(runningMsgConn, &msg.NewWorkConn{RunID: "shared-run"}, wire.ProtocolV1, false))
	require.Len(t, ctl.workConnCh, 1)

	require.NoError(t, ctl.Close())
	waitForControlDone(t, ctl)
	require.Equal(t, int64(1), runningConn.closeCount.Load())
}

func TestServiceWorkConnRoutingRejectsWireProtocolMismatch(t *testing.T) {
	svr := newControlTestService(t)
	ctl, controlConn, err := registerLifecycleTestControl(svr)
	require.NoError(t, err)
	require.NoError(t, svr.completeControlLogin(ctl, func() error { return nil }))
	waitForSignal(t, controlConn.readStarted, "control reader to start")

	workConn := newCountingCloseConn()
	workMsgConn := msg.NewConn(workConn, msg.NewV2ReadWriter(workConn))
	err = svr.RegisterWorkConn(workMsgConn, &msg.NewWorkConn{RunID: "shared-run"}, wire.ProtocolV2, false)
	require.ErrorContains(t, err, "wire protocol mismatch")
	require.Len(t, ctl.workConnCh, 0)
	_ = workMsgConn.Close()
	require.NoError(t, ctl.Close())
}

func TestServiceWorkConnRoutingClientHelloPolicy(t *testing.T) {
	for _, tc := range []struct {
		name                   string
		controlUDPPacketCodec  string
		workClientHelloPresent bool
		errorSubstring         string
	}{
		{
			name: "JSON control allows work connection without Hello",
		},
		{
			name:                  "binary control allows work connection without Hello",
			controlUDPPacketCodec: wire.UDPPacketCodecBinary,
		},
		{
			name:                   "JSON control rejects work connection with Hello",
			workClientHelloPresent: true,
			errorSubstring:         "ClientHello is not allowed",
		},
		{
			name:                   "binary control rejects work connection with Hello",
			controlUDPPacketCodec:  wire.UDPPacketCodecBinary,
			workClientHelloPresent: true,
			errorSubstring:         "ClientHello is not allowed",
		},
	} {
		t.Run(tc.name, func(t *testing.T) {
			svr := newControlTestService(t)
			controlConn := newDeadlineReadConn()
			controlMsgConn := msg.NewConn(controlConn, msg.NewV2ReadWriter(controlConn))
			ctl, err := svr.RegisterControl(controlMsgConn, &msg.Login{
				RunID:    "shared-run",
				ClientID: "client",
				ClientSpec: msg.ClientSpec{
					AlwaysAuthPass: true,
				},
			}, true, wire.ProtocolV2, tc.controlUDPPacketCodec)
			require.NoError(t, err)
			require.NoError(t, svr.completeControlLogin(ctl, func() error { return nil }))
			waitForSignal(t, controlConn.readStarted, "control reader to start")

			workConn := newCountingCloseConn()
			workMsgConn := msg.NewConn(workConn, msg.NewV2ReadWriter(workConn))
			err = svr.RegisterWorkConn(
				workMsgConn,
				&msg.NewWorkConn{RunID: "shared-run"},
				wire.ProtocolV2,
				tc.workClientHelloPresent,
			)
			if tc.errorSubstring != "" {
				require.ErrorContains(t, err, tc.errorSubstring)
				require.Len(t, ctl.workConnCh, 0)
				require.NoError(t, workMsgConn.Close())
			} else {
				require.NoError(t, err)
				require.Len(t, ctl.workConnCh, 1)
			}

			require.NoError(t, ctl.Close())
			waitForControlDone(t, ctl)
			require.Equal(t, int64(1), workConn.closeCount.Load())
		})
	}
}

func TestServiceRegisterControlRejectsInvalidCodecSelection(t *testing.T) {
	for _, tc := range []struct {
		name           string
		wireProtocol   string
		udpPacketCodec string
		errorSubstring string
	}{
		{
			name:           "binary codec over v1",
			wireProtocol:   wire.ProtocolV1,
			udpPacketCodec: wire.UDPPacketCodecBinary,
			errorSubstring: "requires wire protocol v2",
		},
		{
			name:           "unknown v2 codec",
			wireProtocol:   wire.ProtocolV2,
			udpPacketCodec: "unknown",
			errorSubstring: "unsupported UDP packet codec",
		},
		{
			name:           "unknown wire protocol",
			wireProtocol:   "unknown",
			errorSubstring: "unsupported wire protocol",
		},
	} {
		t.Run(tc.name, func(t *testing.T) {
			svr := newControlTestService(t)
			conn := newDeadlineReadConn()
			msgConn := msg.NewConn(conn, msg.NewV1ReadWriter(conn))
			ctl, err := svr.RegisterControl(msgConn, &msg.Login{}, true, tc.wireProtocol, tc.udpPacketCodec)
			require.Nil(t, ctl)
			require.ErrorContains(t, err, tc.errorSubstring)
		})
	}
}

func TestServiceRegisterControlPoolCountBoundaries(t *testing.T) {
	for _, tc := range []struct {
		name      string
		poolCount int
		wantErr   bool
		wantPool  int
	}{
		{name: "less than channel offset", poolCount: -11, wantErr: true},
		{name: "channel offset", poolCount: -10, wantErr: true},
		{name: "negative", poolCount: -1, wantErr: true},
		{name: "zero", poolCount: 0, wantPool: 0},
		{name: "capped by server maximum", poolCount: 10, wantPool: 5},
		{name: "maximum int capped", poolCount: math.MaxInt, wantPool: 5},
	} {
		t.Run(tc.name, func(t *testing.T) {
			svr := newControlTestService(t)
			svr.cfg.Transport.MaxPoolCount = 5
			conn := newDeadlineReadConn()
			msgConn := msg.NewConn(conn, msg.NewV1ReadWriter(conn))
			const timestamp = int64(1)

			ctl, err := svr.RegisterControl(msgConn, &msg.Login{
				RunID:        "pool-count-run",
				ClientID:     "client",
				Timestamp:    timestamp,
				PrivilegeKey: util.GetAuthKey("", timestamp),
				PoolCount:    tc.poolCount,
			}, false, wire.ProtocolV1, "")
			if tc.wantErr {
				require.Nil(t, ctl)
				require.ErrorContains(t, err, "unexpected error when creating new controller")
				return
			}

			require.NoError(t, err)
			require.Equal(t, tc.wantPool, ctl.poolCount)
			require.NoError(t, ctl.Close())
			waitForControlDone(t, ctl)
		})
	}
}

func TestServiceWorkConnRoutingRejectsLostGeneration(t *testing.T) {
	for _, action := range []string{"replace", "close"} {
		t.Run(action, func(t *testing.T) {
			svr := newControlTestService(t)
			ctl, controlConn, err := registerLifecycleTestControl(svr)
			require.NoError(t, err)
			require.NoError(t, svr.completeControlLogin(ctl, func() error { return nil }))
			waitForSignal(t, controlConn.readStarted, "control reader to start")

			barrier := newWorkConnBarrierPlugin()
			svr.pluginManager.Register(barrier)
			workConn := newCountingCloseConn()
			workMsgConn := msg.NewConn(workConn, msg.NewV1ReadWriter(workConn))
			routeDone := make(chan error, 1)
			go func() {
				routeDone <- registerWorkConnAsCaller(svr, workMsgConn, &msg.NewWorkConn{RunID: "shared-run"}, wire.ProtocolV1, false)
			}()
			waitForSignal(t, barrier.entered, "work connection plugin barrier")

			var replacement *Control
			switch action {
			case "replace":
				replacement, _, err = registerLifecycleTestControl(svr)
				require.NoError(t, err)
			case "close":
				require.NoError(t, ctl.Close())
				waitForControlDone(t, ctl)
			}

			close(barrier.resume)
			require.Error(t, waitForResult(t, routeDone, "work connection route to finish"))
			require.Equal(t, int64(1), workConn.closeCount.Load())
			require.Len(t, ctl.workConnCh, 0)

			if replacement != nil {
				require.Len(t, replacement.workConnCh, 0)
				require.True(t, svr.ctlManager.Remove(replacement))
				require.NoError(t, replacement.Close())
				waitForControlDone(t, replacement)
			}
		})
	}
}

func TestServiceVisitorRoutingExcludesPendingUser(t *testing.T) {
	svr := newControlTestService(t)
	listener, err := svr.rc.VisitorManager.Listen("visitor", "secret", []string{"pending-user"})
	require.NoError(t, err)
	t.Cleanup(func() { _ = listener.Close() })

	controlConn := newDeadlineReadConn()
	controlMsgConn := msg.NewConn(controlConn, msg.NewV1ReadWriter(controlConn))
	ctl, err := svr.RegisterControl(controlMsgConn, &msg.Login{
		RunID:    "visitor-run",
		User:     "pending-user",
		ClientID: "visitor-client",
		ClientSpec: msg.ClientSpec{
			AlwaysAuthPass: true,
		},
	}, true, wire.ProtocolV1, "")
	require.NoError(t, err)

	timestamp := time.Now().Unix()
	visitorMsg := &msg.NewVisitorConn{
		RunID:     "visitor-run",
		ProxyName: "visitor",
		Timestamp: timestamp,
		SignKey:   util.GetAuthKey("secret", timestamp),
	}
	pendingConn := newCountingCloseConn()
	err = svr.RegisterVisitorConn(pendingConn, visitorMsg, wire.ProtocolV1)
	require.ErrorContains(t, err, "no client control found")
	require.NoError(t, pendingConn.Close())
	require.Equal(t, int64(1), pendingConn.closeCount.Load())

	require.NoError(t, svr.completeControlLogin(ctl, func() error { return nil }))
	waitForSignal(t, controlConn.readStarted, "control reader to start")
	runningConn := newCountingCloseConn()
	require.NoError(t, svr.RegisterVisitorConn(runningConn, visitorMsg, wire.ProtocolV1))
	accepted, err := listener.Accept()
	require.NoError(t, err)
	require.NoError(t, accepted.Close())
	require.Equal(t, int64(1), runningConn.closeCount.Load())

	require.NoError(t, ctl.Close())
	waitForControlDone(t, ctl)
}

func TestServiceVisitorRoutingCarriesControlPacketCodec(t *testing.T) {
	svr := newControlTestService(t)
	listener, err := svr.rc.VisitorManager.Listen("visitor", "secret", []string{"visitor-user"})
	require.NoError(t, err)
	t.Cleanup(func() { _ = listener.Close() })

	controlConn := newDeadlineReadConn()
	controlMsgConn := msg.NewConn(controlConn, msg.NewV2ReadWriter(controlConn))
	ctl, err := svr.RegisterControl(controlMsgConn, &msg.Login{
		RunID:    "visitor-binary-run",
		User:     "visitor-user",
		ClientID: "visitor-client",
		ClientSpec: msg.ClientSpec{
			AlwaysAuthPass: true,
		},
	}, true, wire.ProtocolV2, wire.UDPPacketCodecBinary)
	require.NoError(t, err)

	timestamp := time.Now().Unix()
	visitorMsg := &msg.NewVisitorConn{
		RunID:     "visitor-binary-run",
		ProxyName: "visitor",
		Timestamp: timestamp,
		SignKey:   util.GetAuthKey("secret", timestamp),
	}
	require.NoError(t, svr.completeControlLogin(ctl, func() error { return nil }))
	waitForSignal(t, controlConn.readStarted, "binary visitor control reader to start")

	runningConn := newCountingCloseConn()
	require.NoError(t, svr.RegisterVisitorConn(runningConn, visitorMsg, wire.ProtocolV2))
	accepted, err := listener.Accept()
	require.NoError(t, err)
	metadata, ok := accepted.(interface {
		WireProtocol() string
		UDPPacketCodec() string
	})
	require.True(t, ok)
	require.Equal(t, wire.ProtocolV2, metadata.WireProtocol())
	require.Equal(t, wire.UDPPacketCodecBinary, metadata.UDPPacketCodec())
	require.NoError(t, accepted.Close())
	require.Equal(t, int64(1), runningConn.closeCount.Load())

	mismatchConn := newCountingCloseConn()
	err = svr.RegisterVisitorConn(mismatchConn, visitorMsg, wire.ProtocolV1)
	require.ErrorContains(t, err, "visitor connection wire protocol mismatch")
	require.NoError(t, mismatchConn.Close())
	require.Equal(t, int64(1), mismatchConn.closeCount.Load())

	require.NoError(t, ctl.Close())
	waitForControlDone(t, ctl)
}

func TestServiceVisitorRoutingLegacyFallsBackToJSONPacketCodec(t *testing.T) {
	svr := newControlTestService(t)
	listener, err := svr.rc.VisitorManager.Listen("visitor", "secret", []string{""})
	require.NoError(t, err)
	t.Cleanup(func() { _ = listener.Close() })

	timestamp := time.Now().Unix()
	visitorMsg := &msg.NewVisitorConn{
		ProxyName: "visitor",
		Timestamp: timestamp,
		SignKey:   util.GetAuthKey("secret", timestamp),
	}
	visitorConn := newCountingCloseConn()
	require.NoError(t, svr.RegisterVisitorConn(visitorConn, visitorMsg, wire.ProtocolV2))
	accepted, err := listener.Accept()
	require.NoError(t, err)
	metadata, ok := accepted.(interface{ UDPPacketCodec() string })
	require.True(t, ok)
	require.Empty(t, metadata.UDPPacketCodec())
	require.NoError(t, accepted.Close())
	require.Equal(t, int64(1), visitorConn.closeCount.Load())
}

func newControlTestService(t *testing.T) *Service {
	t.Helper()
	cfg := &v1.ServerConfig{}
	cfg.Auth.Method = v1.AuthMethodToken
	authRuntime, err := auth.BuildServerAuth(&cfg.Auth)
	require.NoError(t, err)
	clientRegistry := registry.NewClientRegistry()
	return &Service{
		ctlManager:     NewControlManager(clientRegistry),
		clientRegistry: clientRegistry,
		pxyManager:     proxy.NewManager(),
		pluginManager:  plugin.NewManager(),
		rc: &controller.ResourceController{
			VisitorManager: visitor.NewManager(),
		},
		auth: authRuntime,
		cfg:  cfg,
	}
}

func registerLifecycleTestControl(svr *Service) (*Control, *deadlineReadConn, error) {
	conn := newDeadlineReadConn()
	msgConn := msg.NewConn(conn, msg.NewReadWriter(conn, wire.ProtocolV1))
	ctl, err := svr.RegisterControl(msgConn, &msg.Login{
		RunID:    "shared-run",
		ClientID: "client",
		ClientSpec: msg.ClientSpec{
			AlwaysAuthPass: true,
		},
	}, true, wire.ProtocolV1, "")
	return ctl, conn, err
}

func waitForDifferentCurrentControl(t *testing.T, manager *ControlManager, runID string, old *Control) *Control {
	t.Helper()
	deadline := time.Now().Add(3 * time.Second)
	for time.Now().Before(deadline) {
		if ctl := currentControlForTest(manager, runID); ctl != nil && ctl != old {
			return ctl
		}
		runtime.Gosched()
	}
	t.Fatalf("timed out waiting for a new current control after ID %d", old.ID())
	return nil
}

func registerWorkConnAsCaller(
	svr *Service,
	workConn *msg.Conn,
	newMsg *msg.NewWorkConn,
	wireProtocol string,
	clientHelloPresent bool,
) error {
	err := svr.RegisterWorkConn(workConn, newMsg, wireProtocol, clientHelloPresent)
	if err != nil {
		_ = workConn.Close()
	}
	return err
}

func waitForResult[T any](t *testing.T, ch <-chan T, description string) T {
	t.Helper()
	select {
	case result := <-ch:
		return result
	case <-time.After(3 * time.Second):
		t.Fatalf("timed out waiting for %s", description)
		var zero T
		return zero
	}
}

type workConnBarrierPlugin struct {
	entered chan struct{}
	resume  chan struct{}
}

func newWorkConnBarrierPlugin() *workConnBarrierPlugin {
	return &workConnBarrierPlugin{
		entered: make(chan struct{}),
		resume:  make(chan struct{}),
	}
}

func (*workConnBarrierPlugin) Name() string { return "work-conn-barrier" }

func (*workConnBarrierPlugin) IsSupport(op string) bool { return op == plugin.OpNewWorkConn }

func (p *workConnBarrierPlugin) Handle(
	context.Context,
	string,
	any,
) (*plugin.Response, any, error) {
	close(p.entered)
	<-p.resume
	return &plugin.Response{Unchange: true}, nil, nil
}

type countingCloseConn struct {
	closeCount atomic.Int64
}

func newCountingCloseConn() *countingCloseConn { return &countingCloseConn{} }

func (*countingCloseConn) Read([]byte) (int, error)         { return 0, net.ErrClosed }
func (*countingCloseConn) Write(p []byte) (int, error)      { return len(p), nil }
func (c *countingCloseConn) Close() error                   { c.closeCount.Add(1); return nil }
func (*countingCloseConn) LocalAddr() net.Addr              { return lifecycleTestAddr("local") }
func (*countingCloseConn) RemoteAddr() net.Addr             { return lifecycleTestAddr("remote") }
func (*countingCloseConn) SetDeadline(time.Time) error      { return nil }
func (*countingCloseConn) SetReadDeadline(time.Time) error  { return nil }
func (*countingCloseConn) SetWriteDeadline(time.Time) error { return nil }
