lee/frp
1
0
Fork 0
mirror of https://github.com/fatedier/frp.git synced 2025-12-11 02:19:06 +00:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

server/proxy: simplify the code (#3488)

Browse Source
This commit is contained in:
fatedier
2023-06-16 00:14:19 +08:00
committed by GitHub
parent 9ba6a06470
commit e1cef053be
10 changed files with 235 additions and 173 deletions
Download Patch File Download Diff File Expand all files Collapse all files

36
server/proxy/udp.go
View File

@@ -19,12 +19,12 @@ import (
"fmt"
"io"
"net"
"reflect"
"strconv"
"time"
"github.com/fatedier/golib/errors"
libio "github.com/fatedier/golib/io"
"golang.org/x/time/rate"
"github.com/fatedier/frp/pkg/config"
"github.com/fatedier/frp/pkg/msg"
@@ -34,11 +34,15 @@ import (
"github.com/fatedier/frp/server/metrics"
)
func init() {
RegisterProxyFactory(reflect.TypeOf(&config.UDPProxyConf{}), NewUDPProxy)
}
type UDPProxy struct {
*BaseProxy
cfg *config.UDPProxyConf
realPort int
realBindPort int
// udpConn is the listener of udp packages
udpConn *net.UDPConn
@@ -59,21 +63,33 @@ type UDPProxy struct {
isClosed bool
}
func NewUDPProxy(baseProxy *BaseProxy, cfg config.ProxyConf) Proxy {
unwrapped, ok := cfg.(*config.UDPProxyConf)
if !ok {
return nil
}
baseProxy.usedPortsNum = 1
return &UDPProxy{
BaseProxy: baseProxy,
cfg: unwrapped,
}
}
func (pxy *UDPProxy) Run() (remoteAddr string, err error) {
xl := pxy.xl
pxy.realPort, err = pxy.rc.UDPPortManager.Acquire(pxy.name, pxy.cfg.RemotePort)
pxy.realBindPort, err = pxy.rc.UDPPortManager.Acquire(pxy.name, pxy.cfg.RemotePort)
if err != nil {
return "", fmt.Errorf("acquire port %d error: %v", pxy.cfg.RemotePort, err)
}
defer func() {
if err != nil {
pxy.rc.UDPPortManager.Release(pxy.realPort)
pxy.rc.UDPPortManager.Release(pxy.realBindPort)
}
}()
remoteAddr = fmt.Sprintf(":%d", pxy.realPort)
pxy.cfg.RemotePort = pxy.realPort
addr, errRet := net.ResolveUDPAddr("udp", net.JoinHostPort(pxy.serverCfg.ProxyBindAddr, strconv.Itoa(pxy.realPort)))
remoteAddr = fmt.Sprintf(":%d", pxy.realBindPort)
pxy.cfg.RemotePort = pxy.realBindPort
addr, errRet := net.ResolveUDPAddr("udp", net.JoinHostPort(pxy.serverCfg.ProxyBindAddr, strconv.Itoa(pxy.realBindPort)))
if errRet != nil {
err = errRet
return
@@ -233,10 +249,6 @@ func (pxy *UDPProxy) GetConf() config.ProxyConf {
return pxy.cfg
}
func (pxy *UDPProxy) GetLimiter() *rate.Limiter {
return pxy.limiter
}
func (pxy *UDPProxy) Close() {
pxy.mu.Lock()
defer pxy.mu.Unlock()
@@ -254,5 +266,5 @@ func (pxy *UDPProxy) Close() {
close(pxy.readCh)
close(pxy.sendCh)
}
pxy.rc.UDPPortManager.Release(pxy.realPort)
pxy.rc.UDPPortManager.Release(pxy.realBindPort)
}