frp/client/proxy/proxy_manager.go

176 lines
4.0 KiB
Go
Raw Normal View History

// Copyright 2023 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.
2018-12-09 22:06:22 +08:00
package proxy
2018-01-17 01:09:33 +08:00
import (
2019-10-12 20:13:12 +08:00
"context"
2018-01-17 01:09:33 +08:00
"fmt"
2019-10-12 20:13:12 +08:00
"net"
2023-05-30 20:25:22 +08:00
"reflect"
2018-01-17 01:09:33 +08:00
"sync"
"github.com/samber/lo"
2018-12-09 22:06:22 +08:00
"github.com/fatedier/frp/client/event"
v1 "github.com/fatedier/frp/pkg/config/v1"
2020-09-23 13:49:14 +08:00
"github.com/fatedier/frp/pkg/msg"
"github.com/fatedier/frp/pkg/transport"
2020-09-23 13:49:14 +08:00
"github.com/fatedier/frp/pkg/util/xlog"
2018-01-17 01:09:33 +08:00
)
2020-05-24 17:48:37 +08:00
type Manager struct {
2023-11-21 11:19:35 +08:00
proxies map[string]*Wrapper
msgTransporter transport.MessageTransporter
inWorkConnCallback func(*v1.ProxyBaseConfig, net.Conn, *msg.StartWorkConn) bool
2018-01-17 01:09:33 +08:00
2018-12-07 17:05:36 +08:00
closed bool
mu sync.RWMutex
clientCfg *v1.ClientCommonConfig
2019-10-12 20:13:12 +08:00
ctx context.Context
2018-01-17 01:09:33 +08:00
}
func NewManager(
ctx context.Context,
clientCfg *v1.ClientCommonConfig,
msgTransporter transport.MessageTransporter,
) *Manager {
2020-05-24 17:48:37 +08:00
return &Manager{
proxies: make(map[string]*Wrapper),
msgTransporter: msgTransporter,
closed: false,
clientCfg: clientCfg,
ctx: ctx,
2018-01-17 01:09:33 +08:00
}
}
2020-05-24 17:48:37 +08:00
func (pm *Manager) StartProxy(name string, remoteAddr string, serverRespErr string) error {
2018-12-07 17:05:36 +08:00
pm.mu.RLock()
2018-01-17 01:09:33 +08:00
pxy, ok := pm.proxies[name]
2018-12-07 17:05:36 +08:00
pm.mu.RUnlock()
2018-01-17 01:09:33 +08:00
if !ok {
2018-12-07 17:05:36 +08:00
return fmt.Errorf("proxy [%s] not found", name)
2018-01-17 01:09:33 +08:00
}
2018-12-07 17:05:36 +08:00
err := pxy.SetRunningStatus(remoteAddr, serverRespErr)
if err != nil {
return err
2018-01-17 01:09:33 +08:00
}
return nil
}
2023-11-21 11:19:35 +08:00
func (pm *Manager) SetInWorkConnCallback(cb func(*v1.ProxyBaseConfig, net.Conn, *msg.StartWorkConn) bool) {
pm.inWorkConnCallback = cb
}
2020-05-24 17:48:37 +08:00
func (pm *Manager) Close() {
2018-12-11 15:06:54 +08:00
pm.mu.Lock()
defer pm.mu.Unlock()
2018-01-17 01:09:33 +08:00
for _, pxy := range pm.proxies {
2018-12-07 17:05:36 +08:00
pxy.Stop()
2018-01-17 01:09:33 +08:00
}
2020-05-24 17:48:37 +08:00
pm.proxies = make(map[string]*Wrapper)
2018-01-17 01:09:33 +08:00
}
2020-05-24 17:48:37 +08:00
func (pm *Manager) HandleWorkConn(name string, workConn net.Conn, m *msg.StartWorkConn) {
2018-01-17 01:09:33 +08:00
pm.mu.RLock()
2018-12-07 17:05:36 +08:00
pw, ok := pm.proxies[name]
pm.mu.RUnlock()
if ok {
2019-03-29 19:01:18 +08:00
pw.InWorkConn(workConn, m)
2018-12-07 17:05:36 +08:00
} else {
workConn.Close()
2018-01-17 01:09:33 +08:00
}
2018-12-07 17:05:36 +08:00
}
func (pm *Manager) HandleEvent(payload interface{}) error {
2018-12-07 17:05:36 +08:00
var m msg.Message
2018-12-09 22:06:22 +08:00
switch e := payload.(type) {
case *event.StartProxyPayload:
m = e.NewProxyMsg
case *event.CloseProxyPayload:
m = e.CloseProxyMsg
2018-12-07 17:05:36 +08:00
default:
2018-12-09 22:06:22 +08:00
return event.ErrPayloadType
2018-12-07 17:05:36 +08:00
}
return pm.msgTransporter.Send(m)
2018-12-07 17:05:36 +08:00
}
2018-01-17 01:09:33 +08:00
2020-05-24 17:48:37 +08:00
func (pm *Manager) GetAllProxyStatus() []*WorkingStatus {
ps := make([]*WorkingStatus, 0)
2018-12-07 17:05:36 +08:00
pm.mu.RLock()
defer pm.mu.RUnlock()
2018-01-17 01:09:33 +08:00
for _, pxy := range pm.proxies {
2018-12-07 17:05:36 +08:00
ps = append(ps, pxy.GetStatus())
2018-01-17 01:09:33 +08:00
}
2018-12-07 17:05:36 +08:00
return ps
2018-01-17 01:09:33 +08:00
}
2023-11-27 15:47:49 +08:00
func (pm *Manager) GetProxyStatus(name string) (*WorkingStatus, bool) {
pm.mu.RLock()
defer pm.mu.RUnlock()
if pxy, ok := pm.proxies[name]; ok {
return pxy.GetStatus(), true
}
return nil, false
}
func (pm *Manager) UpdateAll(proxyCfgs []v1.ProxyConfigurer) {
2019-10-12 20:13:12 +08:00
xl := xlog.FromContextSafe(pm.ctx)
2023-11-27 15:47:49 +08:00
proxyCfgsMap := lo.KeyBy(proxyCfgs, func(c v1.ProxyConfigurer) string {
return c.GetBaseConfig().Name
})
2018-01-17 01:09:33 +08:00
pm.mu.Lock()
2018-12-07 17:05:36 +08:00
defer pm.mu.Unlock()
2018-01-17 01:09:33 +08:00
delPxyNames := make([]string, 0)
for name, pxy := range pm.proxies {
del := false
2023-11-27 15:47:49 +08:00
cfg, ok := proxyCfgsMap[name]
2023-05-30 20:25:22 +08:00
if !ok || !reflect.DeepEqual(pxy.Cfg, cfg) {
2022-08-29 01:02:53 +08:00
del = true
2018-01-17 01:09:33 +08:00
}
if del {
delPxyNames = append(delPxyNames, name)
delete(pm.proxies, name)
2018-12-07 17:05:36 +08:00
pxy.Stop()
2018-01-17 01:09:33 +08:00
}
}
2018-12-07 17:05:36 +08:00
if len(delPxyNames) > 0 {
2023-05-30 20:25:22 +08:00
xl.Info("proxy removed: %s", delPxyNames)
2018-12-07 17:05:36 +08:00
}
2018-01-17 01:09:33 +08:00
addPxyNames := make([]string, 0)
2023-11-27 15:47:49 +08:00
for _, cfg := range proxyCfgs {
name := cfg.GetBaseConfig().Name
2018-01-17 01:09:33 +08:00
if _, ok := pm.proxies[name]; !ok {
pxy := NewWrapper(pm.ctx, cfg, pm.clientCfg, pm.HandleEvent, pm.msgTransporter)
2023-11-21 11:19:35 +08:00
if pm.inWorkConnCallback != nil {
pxy.SetInWorkConnCallback(pm.inWorkConnCallback)
}
2018-01-17 01:09:33 +08:00
pm.proxies[name] = pxy
addPxyNames = append(addPxyNames, name)
2018-11-06 18:35:05 +08:00
2018-12-07 17:05:36 +08:00
pxy.Start()
}
2018-11-06 18:35:05 +08:00
}
2018-12-07 17:05:36 +08:00
if len(addPxyNames) > 0 {
2023-05-30 20:25:22 +08:00
xl.Info("proxy added: %s", addPxyNames)
2018-11-06 18:35:05 +08:00
}
}