frp/client/proxy/proxy_manager.go

// 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.

package proxy

import (
	"context"
	"fmt"
	"net"
	"reflect"
	"sync"

	"github.com/samber/lo"

	"github.com/fatedier/frp/client/event"
	v1 "github.com/fatedier/frp/pkg/config/v1"
	"github.com/fatedier/frp/pkg/msg"
	"github.com/fatedier/frp/pkg/transport"
	"github.com/fatedier/frp/pkg/util/xlog"
	"github.com/fatedier/frp/pkg/vnet"
)

type Manager struct {
	proxies            map[string]*Wrapper
	msgTransporter     transport.MessageTransporter
	inWorkConnCallback func(*v1.ProxyBaseConfig, net.Conn, *msg.StartWorkConn) bool
	vnetController     *vnet.Controller

	closed bool
	mu     sync.RWMutex

	encryptionKey []byte
	clientCfg     *v1.ClientCommonConfig

	ctx            context.Context
	udpPacketCodec string
}

func NewManager(
	ctx context.Context,
	clientCfg *v1.ClientCommonConfig,
	encryptionKey []byte,
	msgTransporter transport.MessageTransporter,
	vnetController *vnet.Controller,
	udpPacketCodec string,
) *Manager {
	return &Manager{
		proxies:        make(map[string]*Wrapper),
		msgTransporter: msgTransporter,
		vnetController: vnetController,
		closed:         false,
		encryptionKey:  encryptionKey,
		clientCfg:      clientCfg,
		ctx:            ctx,
		udpPacketCodec: udpPacketCodec,
	}
}

func (pm *Manager) StartProxy(name string, remoteAddr string, serverRespErr string) error {
	pm.mu.RLock()
	pxy, ok := pm.proxies[name]
	pm.mu.RUnlock()
	if !ok {
		return fmt.Errorf("proxy [%s] not found", name)
	}

	err := pxy.SetRunningStatus(remoteAddr, serverRespErr)
	if err != nil {
		return err
	}
	return nil
}

func (pm *Manager) SetInWorkConnCallback(cb func(*v1.ProxyBaseConfig, net.Conn, *msg.StartWorkConn) bool) {
	pm.inWorkConnCallback = cb
}

func (pm *Manager) Close() {
	pm.mu.Lock()
	defer pm.mu.Unlock()
	for _, pxy := range pm.proxies {
		pxy.Stop()
	}
	pm.proxies = make(map[string]*Wrapper)
}

func (pm *Manager) HandleWorkConn(name string, workConn net.Conn, m *msg.StartWorkConn) {
	pm.mu.RLock()
	pw, ok := pm.proxies[name]
	pm.mu.RUnlock()
	if ok {
		pw.InWorkConn(workConn, m)
	} else {
		workConn.Close()
	}
}

func (pm *Manager) HandleEvent(payload any) error {
	var m msg.Message
	switch e := payload.(type) {
	case *event.StartProxyPayload:
		m = e.NewProxyMsg
	case *event.CloseProxyPayload:
		m = e.CloseProxyMsg
	default:
		return event.ErrPayloadType
	}

	return pm.msgTransporter.Send(m)
}

func (pm *Manager) GetAllProxyStatus() []*WorkingStatus {
	pm.mu.RLock()
	defer pm.mu.RUnlock()
	ps := make([]*WorkingStatus, 0, len(pm.proxies))
	for _, pxy := range pm.proxies {
		ps = append(ps, pxy.GetStatus())
	}
	return ps
}

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
}

// PreparePluginReload snapshots live plugin instances before applying config.
// Invoke the returned function only after config application succeeds and all
// configuration locks have been released. Replaced instances are skipped.
func (pm *Manager) PreparePluginReload() func() {
	type target struct {
		pw *Wrapper
		p  reloadablePlugin
	}
	pm.mu.RLock()
	targets := make([]target, 0, len(pm.proxies))
	for _, pw := range pm.proxies {
		pw.mu.RLock()
		if pw.Phase == ProxyPhaseRunning {
			if p := pw.reloadablePluginLocked(); p != nil {
				targets = append(targets, target{pw: pw, p: p})
			}
		}
		pw.mu.RUnlock()
	}
	pm.mu.RUnlock()
	return func() {
		for _, t := range targets {
			t.pw.mu.RLock()
			current := t.pw.Phase == ProxyPhaseRunning && t.pw.reloadablePluginLocked() == t.p
			t.pw.mu.RUnlock()
			if !current {
				continue
			}
			if err := t.p.Reload(); err != nil {
				t.pw.xl.Warnf("reload plugin %s: %v", t.p.Name(), err)
			}
		}
	}
}

func (pm *Manager) UpdateAll(proxyCfgs []v1.ProxyConfigurer) {
	xl := xlog.FromContextSafe(pm.ctx)
	proxyCfgsMap := lo.KeyBy(proxyCfgs, func(c v1.ProxyConfigurer) string {
		return c.GetBaseConfig().Name
	})
	pm.mu.Lock()
	defer pm.mu.Unlock()

	delPxyNames := make([]string, 0)
	for name, pxy := range pm.proxies {
		del := false
		cfg, ok := proxyCfgsMap[name]
		if !ok || !reflect.DeepEqual(pxy.Cfg, cfg) {
			del = true
		}

		if del {
			delPxyNames = append(delPxyNames, name)
			delete(pm.proxies, name)
			pxy.Stop()
		}
	}
	if len(delPxyNames) > 0 {
		xl.Infof("proxy removed: %s", delPxyNames)
	}

	addPxyNames := make([]string, 0)
	for _, cfg := range proxyCfgs {
		name := cfg.GetBaseConfig().Name
		if _, ok := pm.proxies[name]; !ok {
			pxy := NewWrapper(pm.ctx, cfg, pm.clientCfg, pm.encryptionKey, pm.HandleEvent, pm.msgTransporter, pm.vnetController, pm.udpPacketCodec)
			if pm.inWorkConnCallback != nil {
				pxy.SetInWorkConnCallback(pm.inWorkConnCallback)
			}
			pm.proxies[name] = pxy
			addPxyNames = append(addPxyNames, name)

			pxy.Start()
		}
	}
	if len(addPxyNames) > 0 {
		xl.Infof("proxy added: %s", addPxyNames)
	}
}