mirror of
https://github.com/fatedier/frp.git
synced 2026-10-19 19:13:32 +08:00
fix(client): serialize config updates with control handoff (#5577)
This commit is contained in:
fatedier
@@ -0,0 +1,232 @@
|
||||
// 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 client
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"net"
|
||||
"reflect"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/fatedier/frp/client/proxy"
|
||||
"github.com/fatedier/frp/pkg/config/source"
|
||||
v1 "github.com/fatedier/frp/pkg/config/v1"
|
||||
"github.com/fatedier/frp/pkg/msg"
|
||||
)
|
||||
|
||||
// Block proxy construction after the service has read its config snapshot, but
|
||||
// before Control.Run returns and the replacement control is published.
|
||||
type handoffProxyConfig struct {
|
||||
*v1.TCPProxyConfig
|
||||
started chan struct{}
|
||||
resume chan struct{}
|
||||
}
|
||||
|
||||
func init() {
|
||||
proxy.RegisterProxyFactory(reflect.TypeFor[*handoffProxyConfig](), func(base *proxy.BaseProxy, cfg v1.ProxyConfigurer) proxy.Proxy {
|
||||
blocking := cfg.(*handoffProxyConfig)
|
||||
close(blocking.started)
|
||||
<-blocking.resume
|
||||
return proxy.NewGeneralTCPProxy(base, blocking.TCPProxyConfig)
|
||||
})
|
||||
}
|
||||
|
||||
func newReloadTestService(t *testing.T, beforeLogin func()) *Service {
|
||||
t.Helper()
|
||||
svr, err := NewService(ServiceOptions{
|
||||
Common: &v1.ClientCommonConfig{
|
||||
Transport: v1.ClientTransportConfig{HeartbeatInterval: -1, HeartbeatTimeout: -1},
|
||||
},
|
||||
ConfigSourceAggregator: source.NewAggregator(source.NewConfigSource()),
|
||||
ConnectorCreator: func(context.Context, *v1.ClientCommonConfig) Connector {
|
||||
clientConn, serverConn := net.Pipe()
|
||||
peerDone := make(chan struct{})
|
||||
t.Cleanup(func() {
|
||||
_ = clientConn.Close()
|
||||
_ = serverConn.Close()
|
||||
waitLifecycleSignal(t, peerDone)
|
||||
})
|
||||
go func() {
|
||||
defer close(peerDone)
|
||||
rw := msg.NewV1ReadWriter(serverConn)
|
||||
var login msg.Login
|
||||
if err := rw.ReadMsgInto(&login); err != nil {
|
||||
t.Errorf("read login: %v", err)
|
||||
return
|
||||
}
|
||||
if beforeLogin != nil {
|
||||
beforeLogin()
|
||||
}
|
||||
if err := rw.WriteMsg(&msg.LoginResp{RunID: "reload-test"}); err != nil {
|
||||
t.Errorf("write login response: %v", err)
|
||||
return
|
||||
}
|
||||
_, _ = io.Copy(io.Discard, serverConn)
|
||||
}()
|
||||
return &testConnector{conn: clientConn}
|
||||
},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
svr.ctx, svr.cancel = context.WithCancelCause(context.Background())
|
||||
t.Cleanup(func() {
|
||||
svr.cancel(nil)
|
||||
if svr.ctl != nil {
|
||||
_ = svr.ctl.Close()
|
||||
waitLifecycleSignal(t, svr.ctl.Done())
|
||||
}
|
||||
})
|
||||
return svr
|
||||
}
|
||||
|
||||
func TestServiceReloadDuringControlHandoff(t *testing.T) {
|
||||
for _, updateSource := range []bool{false, true} {
|
||||
name := "direct update"
|
||||
if updateSource {
|
||||
name = "config source reload"
|
||||
}
|
||||
t.Run(name, func(t *testing.T) {
|
||||
svr := newReloadTestService(t, nil)
|
||||
svr.loopLoginUntilSuccess(time.Second, true)
|
||||
previous := svr.ctl
|
||||
require.NotNil(t, previous)
|
||||
require.NoError(t, previous.Close())
|
||||
waitLifecycleSignal(t, previous.Done())
|
||||
|
||||
oldProxy := &handoffProxyConfig{
|
||||
TCPProxyConfig: &v1.TCPProxyConfig{
|
||||
ProxyBaseConfig: v1.ProxyBaseConfig{Name: "proxy", Type: "tcp"},
|
||||
RemotePort: 18001,
|
||||
},
|
||||
started: make(chan struct{}),
|
||||
resume: make(chan struct{}),
|
||||
}
|
||||
oldVisitor := &v1.STCPVisitorConfig{VisitorBaseConfig: v1.VisitorBaseConfig{
|
||||
Name: "visitor", Type: "stcp", BindPort: -1, ServerName: "old-target",
|
||||
}}
|
||||
// Seed the desired config for the next connection without restarting
|
||||
// proxies on the already closed control.
|
||||
svr.proxyCfgs = []v1.ProxyConfigurer{oldProxy}
|
||||
svr.visitorCfgs = []v1.VisitorConfigurer{oldVisitor}
|
||||
loginDone, updateDone := make(chan struct{}), make(chan struct{})
|
||||
resume := sync.OnceFunc(func() { close(oldProxy.resume) })
|
||||
t.Cleanup(func() {
|
||||
resume()
|
||||
waitLifecycleSignal(t, loginDone)
|
||||
})
|
||||
go func() {
|
||||
svr.loopLoginUntilSuccess(time.Second, true)
|
||||
close(loginDone)
|
||||
}()
|
||||
waitLifecycleSignal(t, oldProxy.started)
|
||||
svr.ctlMu.RLock()
|
||||
current := svr.ctl
|
||||
svr.ctlMu.RUnlock()
|
||||
require.Same(t, previous, current, "replacement must still be unpublished")
|
||||
// Check before starting the updater so it cannot be the lock owner.
|
||||
unlocked := svr.reloadMu.TryLock()
|
||||
if unlocked {
|
||||
svr.reloadMu.Unlock()
|
||||
}
|
||||
|
||||
newProxy := &v1.TCPProxyConfig{
|
||||
ProxyBaseConfig: v1.ProxyBaseConfig{Name: "proxy", Type: "tcp"},
|
||||
RemotePort: 18002,
|
||||
}
|
||||
newVisitor := &v1.STCPVisitorConfig{VisitorBaseConfig: v1.VisitorBaseConfig{
|
||||
Name: "visitor", Type: "stcp", BindPort: -1, ServerName: "new-target",
|
||||
}}
|
||||
var updateErr error
|
||||
go func() {
|
||||
proxies, visitors := []v1.ProxyConfigurer{newProxy}, []v1.VisitorConfigurer{newVisitor}
|
||||
if updateSource {
|
||||
updateErr = svr.UpdateConfigSource(svr.common, proxies, visitors)
|
||||
} else {
|
||||
updateErr = svr.UpdateAllConfigurer(proxies, visitors)
|
||||
}
|
||||
close(updateDone)
|
||||
}()
|
||||
t.Cleanup(func() {
|
||||
resume()
|
||||
waitLifecycleSignal(t, updateDone)
|
||||
})
|
||||
|
||||
// The gate fixes the handoff window. On the old implementation the
|
||||
// update completes against the retired control before this gate opens.
|
||||
select {
|
||||
case <-updateDone:
|
||||
t.Error("config update completed before control handoff")
|
||||
case <-time.After(20 * time.Millisecond):
|
||||
}
|
||||
require.False(t, unlocked, "control startup must hold the reload transaction until publication")
|
||||
resume()
|
||||
waitLifecycleSignal(t, loginDone)
|
||||
waitLifecycleSignal(t, updateDone)
|
||||
require.NoError(t, updateErr)
|
||||
require.NotSame(t, previous, svr.ctl)
|
||||
status, ok := svr.getProxyStatus("proxy")
|
||||
require.True(t, ok)
|
||||
require.Equal(t, 18002, status.Cfg.(*v1.TCPProxyConfig).RemotePort)
|
||||
visitor, ok := svr.getVisitorCfg("visitor")
|
||||
require.True(t, ok)
|
||||
require.Equal(t, "new-target", visitor.GetBaseConfig().ServerName)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestServiceReloadWhileLoginPending(t *testing.T) {
|
||||
loginStarted, allowLogin, loginDone := make(chan struct{}), make(chan struct{}), make(chan struct{})
|
||||
resume := sync.OnceFunc(func() { close(allowLogin) })
|
||||
svr := newReloadTestService(t, func() {
|
||||
close(loginStarted)
|
||||
<-allowLogin
|
||||
})
|
||||
t.Cleanup(func() {
|
||||
resume()
|
||||
waitLifecycleSignal(t, loginDone)
|
||||
})
|
||||
go func() {
|
||||
svr.loopLoginUntilSuccess(time.Second, true)
|
||||
close(loginDone)
|
||||
}()
|
||||
waitLifecycleSignal(t, loginStarted)
|
||||
|
||||
proxyCfg := &v1.TCPProxyConfig{
|
||||
ProxyBaseConfig: v1.ProxyBaseConfig{Name: "during-login", Type: "tcp"},
|
||||
RemotePort: 18003,
|
||||
}
|
||||
updateDone := make(chan struct{})
|
||||
var updateErr error
|
||||
go func() {
|
||||
updateErr = svr.UpdateAllConfigurer([]v1.ProxyConfigurer{proxyCfg}, nil)
|
||||
close(updateDone)
|
||||
}()
|
||||
t.Cleanup(func() {
|
||||
resume()
|
||||
waitLifecycleSignal(t, updateDone)
|
||||
})
|
||||
// Config updates must complete even while the server has not replied to login.
|
||||
waitLifecycleSignal(t, updateDone)
|
||||
require.NoError(t, updateErr)
|
||||
resume()
|
||||
waitLifecycleSignal(t, loginDone)
|
||||
status, ok := svr.getProxyStatus("during-login")
|
||||
require.True(t, ok)
|
||||
require.Equal(t, 18003, status.Cfg.(*v1.TCPProxyConfig).RemotePort)
|
||||
}
|
||||
Reference in New Issue
Block a user