Mxmilu666/frp
1
1
Fork 1
mirror of https://github.com/fatedier/frp.git synced 2026-10-15 16:15:57 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(server): control work connection refill scheduling

Browse Source
This commit is contained in:
fatedier
2026-09-15 13:55:10 +08:00
parent 832df8dff6
commit 53d0a53a36
11 changed files with 3061 additions and 98 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pkg/msg/msg.go
+11 -1
View File
@@ -147,9 +147,19 @@ type NewWorkConn struct {
RunID string `json:"run_id,omitempty"`
PrivilegeKey string `json:"privilege_key,omitempty"`
Timestamp int64 `json:"timestamp,omitempty"`
WorkConnType string `json:"work_conn_type,omitempty"`
ControlID uint64 `json:"control_id,omitempty"`
}
type ReqWorkConn struct{}
const (
WorkConnTypeDemand = "demand"
WorkConnTypeReserve = "reserve"
)
type ReqWorkConn struct {
WorkConnType string `json:"work_conn_type,omitempty"`
ControlID uint64 `json:"control_id,omitempty"`
}
type StartWorkConn struct {
ProxyName string `json:"proxy_name,omitempty"`
pkg/msg/msg_test.go
+21
View File
@@ -19,6 +19,8 @@ import (
"testing"
"github.com/stretchr/testify/require"
"github.com/fatedier/frp/pkg/proto/wire"
)
func TestV1MessageTypeIDsAreStable(t *testing.T) {
@@ -53,3 +55,22 @@ func TestMessageTypeMapIsCompleteAndUnique(t *testing.T) {
msgTypes[msgType] = struct{}{}
}
}
func TestWorkConnRequestTypeRoundTrip(t *testing.T) {
in := &ReqWorkConn{WorkConnType: WorkConnTypeReserve, ControlID: 7}
out, err := DecodeV2MessageFrame(mustEncodeV2MessageFrame(t, in))
require.NoError(t, err)
require.Equal(t, in, out)
inNew := &NewWorkConn{RunID: "run", WorkConnType: WorkConnTypeDemand, ControlID: 7}
outNew, err := DecodeV2MessageFrame(mustEncodeV2MessageFrame(t, inNew))
require.NoError(t, err)
require.Equal(t, inNew, outNew)
}
func mustEncodeV2MessageFrame(t *testing.T, m Message) *wire.Frame {
t.Helper()
frame, err := EncodeV2MessageFrame(m)
require.NoError(t, err)
return frame
}
pkg/plugin/server/manager.go
+32 -5
View File
@@ -33,6 +33,10 @@ type Manager struct {
newUserConnPlugins []Plugin
}
// ErrNewWorkConnControlIDChanged identifies an attempt to change the control
// generation in a NewWorkConn plugin response.
var ErrNewWorkConnControlIDChanged = errors.New("plugin changed work connection control generation")
func NewManager() *Manager {
return &Manager{
loginPlugins: make([]Plugin, 0),
@@ -72,6 +76,7 @@ func handleMutableContent[T any](
op string,
content *T,
logMode pluginErrorLogMode,
afterChange func(*T) (*T, error),
) (*T, error) {
if len(plugins) == 0 {
return content, nil
@@ -100,6 +105,12 @@ func handleMutableContent[T any](
// Preserve the existing Plugin contract: changed content must be *T.
// Buggy Plugin implementations still panic here, by design.
content = retContent.(*T)
if afterChange != nil {
content, err = afterChange(content)
if err != nil {
return nil, err
}
}
}
}
return content, nil
@@ -127,11 +138,11 @@ func (m *Manager) Register(p Plugin) {
}
func (m *Manager) Login(content *LoginContent) (*LoginContent, error) {
return handleMutableContent(m.loginPlugins, OpLogin, content, pluginErrorLogWarn)
return handleMutableContent(m.loginPlugins, OpLogin, content, pluginErrorLogWarn, nil)
}
func (m *Manager) NewProxy(content *NewProxyContent) (*NewProxyContent, error) {
return handleMutableContent(m.newProxyPlugins, OpNewProxy, content, pluginErrorLogWarn)
return handleMutableContent(m.newProxyPlugins, OpNewProxy, content, pluginErrorLogWarn, nil)
}
func (m *Manager) CloseProxy(content *CloseProxyContent) error {
@@ -157,14 +168,30 @@ func (m *Manager) CloseProxy(content *CloseProxyContent) error {
}
func (m *Manager) Ping(content *PingContent) (*PingContent, error) {
return handleMutableContent(m.pingPlugins, OpPing, content, pluginErrorLogWarn)
return handleMutableContent(m.pingPlugins, OpPing, content, pluginErrorLogWarn, nil)
}
func (m *Manager) NewWorkConn(content *NewWorkConnContent) (*NewWorkConnContent, error) {
return handleMutableContent(m.newWorkConnPlugins, OpNewWorkConn, content, pluginErrorLogWarn)
if len(m.newWorkConnPlugins) == 0 {
return content, nil
}
controlID, workConnType := content.ControlID, content.WorkConnType
return handleMutableContent(m.newWorkConnPlugins, OpNewWorkConn, content, pluginErrorLogWarn,
func(changed *NewWorkConnContent) (*NewWorkConnContent, error) {
if changed.ControlID != 0 && changed.ControlID != controlID {
return nil, fmt.Errorf("%w from %d to %d", ErrNewWorkConnControlIDChanged, controlID, changed.ControlID)
}
// Older plugins may omit scheduling metadata. Restore it before the
// next plugin observes the replacement, leaving all other fields mutable.
// Copy the response so the manager does not mutate plugin-owned content.
restored := *changed
restored.ControlID = controlID
restored.WorkConnType = workConnType
return &restored, nil
})
}
func (m *Manager) NewUserConn(content *NewUserConnContent) (*NewUserConnContent, error) {
// Preserve the pre-refactor log level for NewUserConn plugin errors.
return handleMutableContent(m.newUserConnPlugins, OpNewUserConn, content, pluginErrorLogInfo)
return handleMutableContent(m.newUserConnPlugins, OpNewUserConn, content, pluginErrorLogInfo, nil)
}
pkg/plugin/server/manager_test.go
+128
View File
@@ -24,6 +24,7 @@ import (
"time"
goliblog "github.com/fatedier/golib/log"
"github.com/stretchr/testify/require"
"github.com/fatedier/frp/pkg/msg"
frplog "github.com/fatedier/frp/pkg/util/log"
@@ -259,6 +260,133 @@ func TestManagerMutableContentRejectStopsChain(t *testing.T) {
}
}
func TestManagerNewWorkConnPreservesMetadataBetweenPlugins(t *testing.T) {
for _, tc := range []struct {
name string
id uint64
kind string
}{
{name: "typed", id: 7, kind: msg.WorkConnTypeReserve},
{name: "legacy"},
} {
t.Run(tc.name, func(t *testing.T) {
m := NewManager()
observations := 0
for range 2 {
// Model an old plugin replacing the entire content, including
// intentionally omitted ordinary fields. Only metadata is restored.
partial := &NewWorkConnContent{NewWorkConn: msg.NewWorkConn{RunID: "mutated", Timestamp: 42}}
m.Register(testPlugin{
name: "partial", ops: map[string]bool{OpNewWorkConn: true},
handler: func(context.Context, string, any) (*Response, any, error) {
return &Response{Unchange: false}, partial, nil
},
})
m.Register(testPlugin{
name: "policy", ops: map[string]bool{OpNewWorkConn: true},
handler: func(_ context.Context, _ string, raw any) (*Response, any, error) {
content := raw.(NewWorkConnContent)
observations++
if content.ControlID != tc.id || content.WorkConnType != tc.kind {
return &Response{Reject: true, RejectReason: "lost scheduling metadata"}, nil, nil
}
require.Equal(t, "mutated", content.RunID)
require.Equal(t, int64(42), content.Timestamp)
require.Empty(t, content.User.User, "ordinary omitted fields must not be merged back")
require.Zero(t, partial.ControlID, "do not mutate the plugin's response")
// Unchange must ignore even an invalid returned generation.
content.ControlID = tc.id + 1
content.WorkConnType = "ignored"
return &Response{Unchange: true}, &content, nil
},
})
}
got, err := m.NewWorkConn(&NewWorkConnContent{
User: UserInfo{User: "original"},
NewWorkConn: msg.NewWorkConn{RunID: "original", ControlID: tc.id, WorkConnType: tc.kind},
})
require.NoError(t, err)
require.Equal(t, 2, observations)
require.Equal(t, tc.id, got.ControlID)
require.Equal(t, tc.kind, got.WorkConnType)
require.Equal(t, "mutated", got.RunID)
})
}
}
func TestManagerNewWorkConnMetadataValidationPreservesPluginContract(t *testing.T) {
for _, tc := range []struct {
name string
initialID uint64
returnedID uint64
response Response
pluginErr error
wantErr string
wantIDErr bool
}{
{name: "changed generation", initialID: 7, returnedID: 8, wantErr: "from 7 to 8", wantIDErr: true},
{name: "legacy cannot inject generation", returnedID: 8, wantErr: "from 0 to 8", wantIDErr: true},
{name: "same generation", initialID: 7, returnedID: 7},
{name: "omitted generation", initialID: 7},
{name: "unchange ignores generation rewrite", initialID: 7, returnedID: 8, response: Response{Unchange: true}},
{
name: "reject stops before validation", initialID: 7, returnedID: 8,
response: Response{Reject: true, RejectReason: "policy rejection"}, wantErr: "policy rejection",
},
{
name: "error stops before validation", initialID: 7, returnedID: 8,
pluginErr: errors.New("unavailable"), wantErr: "send NewWorkConn request to plugin error",
},
} {
t.Run(tc.name, func(t *testing.T) {
m := NewManager()
m.Register(testPlugin{
name: "mutate", ops: map[string]bool{OpNewWorkConn: true},
handler: func(context.Context, string, any) (*Response, any, error) {
return &tc.response, &NewWorkConnContent{NewWorkConn: msg.NewWorkConn{
ControlID: tc.returnedID, WorkConnType: msg.WorkConnTypeDemand,
}}, tc.pluginErr
},
})
called := false
m.Register(testPlugin{
name: "observe", ops: map[string]bool{OpNewWorkConn: true},
handler: func(_ context.Context, _ string, raw any) (*Response, any, error) {
called = true
content := raw.(NewWorkConnContent)
require.Equal(t, tc.initialID, content.ControlID)
require.Equal(t, msg.WorkConnTypeReserve, content.WorkConnType)
return &Response{Unchange: true}, nil, nil
},
})
got, err := m.NewWorkConn(&NewWorkConnContent{NewWorkConn: msg.NewWorkConn{
ControlID: tc.initialID, WorkConnType: msg.WorkConnTypeReserve,
}})
if tc.wantErr != "" {
require.ErrorContains(t, err, tc.wantErr)
require.Nil(t, got)
require.False(t, called)
} else {
require.NoError(t, err)
require.NotNil(t, got)
require.True(t, called)
}
require.Equal(t, tc.wantIDErr, errors.Is(err, ErrNewWorkConnControlIDChanged))
})
}
}
func TestManagerNewWorkConnChangedContentRequiresPointer(t *testing.T) {
m := NewManager()
m.Register(testPlugin{
name: "invalid", ops: map[string]bool{OpNewWorkConn: true},
handler: func(context.Context, string, any) (*Response, any, error) {
return &Response{Unchange: false}, NewWorkConnContent{}, nil
},
})
require.Panics(t, func() { _, _ = m.NewWorkConn(&NewWorkConnContent{}) })
}
func TestManagerMutableContentPluginErrorLogLevel(t *testing.T) {
tests := []struct {
name string