Revival

pkg/plugin/server/manager.go

@@ -1,17 +1,3 @@
-// Copyright 2019 fatedier, fatedier@gmail.com
-//
-// 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 (
@@ -24,238 +10,126 @@ import (
 	"github.com/fatedier/frp/pkg/util/xlog"
 )
 
-type Manager struct {
+type CompositeOperationPluginGateway struct {
-	loginPlugins       []Plugin
+	router *PluginOpsRouter
-	newProxyPlugins    []Plugin
-	closeProxyPlugins  []Plugin
-	pingPlugins        []Plugin
-	newWorkConnPlugins []Plugin
-	newUserConnPlugins []Plugin
 }
 
-func NewManager() *Manager {
+func NewManager() *CompositeOperationPluginGateway {
-	return &Manager{
+	return &CompositeOperationPluginGateway{
-		loginPlugins:       make([]Plugin, 0),
+		router: NewPluginOpsRouter(),
-		newProxyPlugins:    make([]Plugin, 0),
-		closeProxyPlugins:  make([]Plugin, 0),
-		pingPlugins:        make([]Plugin, 0),
-		newWorkConnPlugins: make([]Plugin, 0),
-		newUserConnPlugins: make([]Plugin, 0),
 	}
 }
 
-func (m *Manager) Register(p Plugin) {
+type PluginOpsRouter struct {
-	if p.IsSupport(OpLogin) {
+	operations map[string][]Plugin
-		m.loginPlugins = append(m.loginPlugins, p)
+}
-	}
+
-	if p.IsSupport(OpNewProxy) {
+func NewPluginOpsRouter() *PluginOpsRouter {
-		m.newProxyPlugins = append(m.newProxyPlugins, p)
+	return &PluginOpsRouter{
-	}
+		operations: make(map[string][]Plugin),
-	if p.IsSupport(OpCloseProxy) {
-		m.closeProxyPlugins = append(m.closeProxyPlugins, p)
-	}
-	if p.IsSupport(OpPing) {
-		m.pingPlugins = append(m.pingPlugins, p)
-	}
-	if p.IsSupport(OpNewWorkConn) {
-		m.newWorkConnPlugins = append(m.newWorkConnPlugins, p)
-	}
-	if p.IsSupport(OpNewUserConn) {
-		m.newUserConnPlugins = append(m.newUserConnPlugins, p)
 	}
 }
 
-func (m *Manager) Login(content *LoginContent) (*LoginContent, error) {
+func (r *PluginOpsRouter) AddPlugin(op string, p Plugin) {
-	if len(m.loginPlugins) == 0 {
+	r.operations[op] = append(r.operations[op], p)
-		return content, nil
+}
-	}
 
-	var (
+func (m *CompositeOperationPluginGateway) Register(p Plugin) {
-		res = &Response{
+	for _, op := range []string{OpLogin, OpNewProxy, OpCloseProxy, OpPing, OpNewWorkConn, OpNewUserConn} {
-			Reject:   false,
+		if p.IsSupport(op) {
-			Unchange: true,
+			m.router.AddPlugin(op, p)
 		}
-		retContent any
+	}
-		err        error
+}
-	)
+
+func buildCtx() (context.Context, string) {
 	reqid, _ := util.RandID()
 	xl := xlog.New().AppendPrefix("reqid: " + reqid)
 	ctx := xlog.NewContext(context.Background(), xl)
 	ctx = NewReqidContext(ctx, reqid)
-
+	return ctx, reqid
-	for _, p := range m.loginPlugins {
-		res, retContent, err = p.Handle(ctx, OpLogin, *content)
-		if err != nil {
-			xl.Warnf("send Login request to plugin [%s] error: %v", p.Name(), err)
-			return nil, errors.New("send Login request to plugin error")
-		}
-		if res.Reject {
-			return nil, fmt.Errorf("%s", res.RejectReason)
-		}
-		if !res.Unchange {
-			content = retContent.(*LoginContent)
-		}
-	}
-	return content, nil
 }
 
-func (m *Manager) NewProxy(content *NewProxyContent) (*NewProxyContent, error) {
+func (m *CompositeOperationPluginGateway) doRequest(op string, input any) (any, error) {
-	if len(m.newProxyPlugins) == 0 {
+	ctx, _ := buildCtx()
-		return content, nil
+	plugins := m.router.operations[op]
+	if len(plugins) == 0 {
+		return input, nil
 	}
 
-	var (
+	response := &Response{
-		res = &Response{
+		Reject:   false,
-			Reject:   false,
+		Unchange: true,
-			Unchange: true,
+	}
-		}
-		retContent any
-		err        error
-	)
-	reqid, _ := util.RandID()
-	xl := xlog.New().AppendPrefix("reqid: " + reqid)
-	ctx := xlog.NewContext(context.Background(), xl)
-	ctx = NewReqidContext(ctx, reqid)
 
-	for _, p := range m.newProxyPlugins {
+	var retContent any
-		res, retContent, err = p.Handle(ctx, OpNewProxy, *content)
+	var err error
+	for _, p := range plugins {
+		response, retContent, err = p.Handle(ctx, op, input)
 		if err != nil {
-			xl.Warnf("send NewProxy request to plugin [%s] error: %v", p.Name(), err)
+			return nil, fmt.Errorf("plugin [%s] failed: %v", p.Name(), err)
-			return nil, errors.New("send NewProxy request to plugin error")
 		}
-		if res.Reject {
+		if response.Reject {
-			return nil, fmt.Errorf("%s", res.RejectReason)
+			return nil, fmt.Errorf("%s", response.RejectReason)
 		}
-		if !res.Unchange {
+		if !response.Unchange {
-			content = retContent.(*NewProxyContent)
+			input = retContent
 		}
 	}
-	return content, nil
+	return input, nil
 }
 
-func (m *Manager) CloseProxy(content *CloseProxyContent) error {
+func (m *CompositeOperationPluginGateway) Login(c *LoginContent) (*LoginContent, error) {
-	if len(m.closeProxyPlugins) == 0 {
+	out, err := m.doRequest(OpLogin, *c)
-		return nil
+	if err != nil {
+		return nil, err
 	}
+	return out.(*LoginContent), nil
+}
 
-	errs := make([]string, 0)
+func (m *CompositeOperationPluginGateway) NewProxy(c *NewProxyContent) (*NewProxyContent, error) {
-	reqid, _ := util.RandID()
+	out, err := m.doRequest(OpNewProxy, *c)
-	xl := xlog.New().AppendPrefix("reqid: " + reqid)
+	if err != nil {
-	ctx := xlog.NewContext(context.Background(), xl)
+		return nil, err
-	ctx = NewReqidContext(ctx, reqid)
+	}
+	return out.(*NewProxyContent), nil
+}
 
-	for _, p := range m.closeProxyPlugins {
+func (m *CompositeOperationPluginGateway) Ping(c *PingContent) (*PingContent, error) {
-		_, _, err := p.Handle(ctx, OpCloseProxy, *content)
+	out, err := m.doRequest(OpPing, *c)
+	if err != nil {
+		return nil, err
+	}
+	return out.(*PingContent), nil
+}
+
+func (m *CompositeOperationPluginGateway) NewWorkConn(c *NewWorkConnContent) (*NewWorkConnContent, error) {
+	out, err := m.doRequest(OpNewWorkConn, *c)
+	if err != nil {
+		return nil, err
+	}
+	return out.(*NewWorkConnContent), nil
+}
+
+func (m *CompositeOperationPluginGateway) NewUserConn(c *NewUserConnContent) (*NewUserConnContent, error) {
+	out, err := m.doRequest(OpNewUserConn, *c)
+	if err != nil {
+		return nil, err
+	}
+	return out.(*NewUserConnContent), nil
+}
+
+func (m *CompositeOperationPluginGateway) CloseProxy(c *CloseProxyContent) error {
+	ctx, _ := buildCtx()
+	errs := []string{}
+	for _, p := range m.router.operations[OpCloseProxy] {
+		_, _, err := p.Handle(ctx, OpCloseProxy, *c)
 		if err != nil {
-			xl.Warnf("send CloseProxy request to plugin [%s] error: %v", p.Name(), err)
 			errs = append(errs, fmt.Sprintf("[%s]: %v", p.Name(), err))
 		}
 	}
-
 	if len(errs) > 0 {
-		return fmt.Errorf("send CloseProxy request to plugin errors: %s", strings.Join(errs, "; "))
+		return fmt.Errorf("plugin CloseProxy errors: %s", strings.Join(errs, "; "))
 	}
 	return nil
 }
-
-func (m *Manager) Ping(content *PingContent) (*PingContent, error) {
-	if len(m.pingPlugins) == 0 {
-		return content, nil
-	}
-
-	var (
-		res = &Response{
-			Reject:   false,
-			Unchange: true,
-		}
-		retContent any
-		err        error
-	)
-	reqid, _ := util.RandID()
-	xl := xlog.New().AppendPrefix("reqid: " + reqid)
-	ctx := xlog.NewContext(context.Background(), xl)
-	ctx = NewReqidContext(ctx, reqid)
-
-	for _, p := range m.pingPlugins {
-		res, retContent, err = p.Handle(ctx, OpPing, *content)
-		if err != nil {
-			xl.Warnf("send Ping request to plugin [%s] error: %v", p.Name(), err)
-			return nil, errors.New("send Ping request to plugin error")
-		}
-		if res.Reject {
-			return nil, fmt.Errorf("%s", res.RejectReason)
-		}
-		if !res.Unchange {
-			content = retContent.(*PingContent)
-		}
-	}
-	return content, nil
-}
-
-func (m *Manager) NewWorkConn(content *NewWorkConnContent) (*NewWorkConnContent, error) {
-	if len(m.newWorkConnPlugins) == 0 {
-		return content, nil
-	}
-
-	var (
-		res = &Response{
-			Reject:   false,
-			Unchange: true,
-		}
-		retContent any
-		err        error
-	)
-	reqid, _ := util.RandID()
-	xl := xlog.New().AppendPrefix("reqid: " + reqid)
-	ctx := xlog.NewContext(context.Background(), xl)
-	ctx = NewReqidContext(ctx, reqid)
-
-	for _, p := range m.newWorkConnPlugins {
-		res, retContent, err = p.Handle(ctx, OpNewWorkConn, *content)
-		if err != nil {
-			xl.Warnf("send NewWorkConn request to plugin [%s] error: %v", p.Name(), err)
-			return nil, errors.New("send NewWorkConn request to plugin error")
-		}
-		if res.Reject {
-			return nil, fmt.Errorf("%s", res.RejectReason)
-		}
-		if !res.Unchange {
-			content = retContent.(*NewWorkConnContent)
-		}
-	}
-	return content, nil
-}
-
-func (m *Manager) NewUserConn(content *NewUserConnContent) (*NewUserConnContent, error) {
-	if len(m.newUserConnPlugins) == 0 {
-		return content, nil
-	}
-
-	var (
-		res = &Response{
-			Reject:   false,
-			Unchange: true,
-		}
-		retContent any
-		err        error
-	)
-	reqid, _ := util.RandID()
-	xl := xlog.New().AppendPrefix("reqid: " + reqid)
-	ctx := xlog.NewContext(context.Background(), xl)
-	ctx = NewReqidContext(ctx, reqid)
-
-	for _, p := range m.newUserConnPlugins {
-		res, retContent, err = p.Handle(ctx, OpNewUserConn, *content)
-		if err != nil {
-			xl.Infof("send NewUserConn request to plugin [%s] error: %v", p.Name(), err)
-			return nil, errors.New("send NewUserConn request to plugin error")
-		}
-		if res.Reject {
-			return nil, fmt.Errorf("%s", res.RejectReason)
-		}
-		if !res.Unchange {
-			content = retContent.(*NewUserConnContent)
-		}
-	}
-	return content, nil
-}