forked from Mxmilu666/frp
fatedier
f32adb3bbb
When a client logs in again with the same run ID, frps replaces the old control and closes its control stream, but the yamux session that carried it stays open until keepalive reclaims it. The replaced control's work connections and streams share that session, so they are left behind for up to tcpMuxKeepaliveInterval plus the ping budget. A mux session carries exactly one control connection, so the control now records its session and closes it from the single teardown point shared by replacement, fatal exit and shutdown. QUIC and non-multiplexed transports keep the previous path. Closing the session ends every stream that still belongs to the replaced control. Those proxies are already being torn down, so the cost is that in-flight connections of the old control are interrupted immediately instead of at the keepalive deadline. A replacement control is not affected, and the ownership rules for a still-online control do not change. Tests: replacement and fatal exit close the session while the replacement control keeps working, and concurrent logins with the same run ID leave exactly one live session.
977 lines
32 KiB
Go
977 lines
32 KiB
Go
// 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"
|
|
"fmt"
|
|
"math"
|
|
"net"
|
|
"net/http"
|
|
"runtime"
|
|
"strings"
|
|
"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/config/v1/validation"
|
|
"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, nil)
|
|
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, nil)
|
|
require.Nil(t, ctl)
|
|
require.ErrorContains(t, err, tc.errorSubstring)
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestServiceRegisterControlRejectsInvalidRunID(t *testing.T) {
|
|
for _, runID := range []string{
|
|
"run\nforged",
|
|
strings.Repeat("a", validation.MaxRunIDLength+1),
|
|
} {
|
|
t.Run(fmt.Sprintf("run_id_%d", len(runID)), func(t *testing.T) {
|
|
svr := newControlTestService(t)
|
|
conn := newDeadlineReadConn()
|
|
msgConn := msg.NewConn(conn, msg.NewV1ReadWriter(conn))
|
|
ctl, err := svr.RegisterControl(msgConn, &msg.Login{RunID: runID}, true, wire.ProtocolV1, "", nil)
|
|
require.Nil(t, ctl)
|
|
require.ErrorContains(t, err, "invalid run id")
|
|
})
|
|
}
|
|
}
|
|
|
|
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, "", nil)
|
|
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, "", nil)
|
|
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, nil)
|
|
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, "", nil)
|
|
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 }
|