frp/server/proxy/udp.go

267 lines
7.0 KiB
Go
Raw Normal View History

2019-01-15 00:11:08 +08:00
// 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 proxy
import (
"context"
"fmt"
2020-09-07 14:57:23 +08:00
"io"
2019-01-15 00:11:08 +08:00
"net"
"reflect"
"strconv"
2019-01-15 00:11:08 +08:00
"time"
2022-08-29 01:02:53 +08:00
"github.com/fatedier/golib/errors"
2023-05-29 14:10:34 +08:00
libio "github.com/fatedier/golib/io"
2022-08-29 01:02:53 +08:00
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/proto/udp"
"github.com/fatedier/frp/pkg/util/limit"
2023-11-27 15:47:49 +08:00
netpkg "github.com/fatedier/frp/pkg/util/net"
"github.com/fatedier/frp/server/metrics"
2019-01-15 00:11:08 +08:00
)
func init() {
RegisterProxyFactory(reflect.TypeOf(&v1.UDPProxyConfig{}), NewUDPProxy)
}
2020-05-24 17:48:37 +08:00
type UDPProxy struct {
2019-01-31 16:49:23 +08:00
*BaseProxy
cfg *v1.UDPProxyConfig
2019-01-15 00:11:08 +08:00
realBindPort int
2019-01-15 00:11:08 +08:00
// udpConn is the listener of udp packages
udpConn *net.UDPConn
// there are always only one workConn at the same time
// get another one if it closed
workConn net.Conn
// sendCh is used for sending packages to workConn
2020-05-24 17:48:37 +08:00
sendCh chan *msg.UDPPacket
2019-01-15 00:11:08 +08:00
// readCh is used for reading packages from workConn
2020-05-24 17:48:37 +08:00
readCh chan *msg.UDPPacket
2019-01-15 00:11:08 +08:00
// checkCloseCh is used for watching if workConn is closed
checkCloseCh chan int
isClosed bool
}
func NewUDPProxy(baseProxy *BaseProxy) Proxy {
unwrapped, ok := baseProxy.GetConfigurer().(*v1.UDPProxyConfig)
if !ok {
return nil
}
baseProxy.usedPortsNum = 1
return &UDPProxy{
BaseProxy: baseProxy,
cfg: unwrapped,
}
}
2020-05-24 17:48:37 +08:00
func (pxy *UDPProxy) Run() (remoteAddr string, err error) {
2019-10-12 20:13:12 +08:00
xl := pxy.xl
pxy.realBindPort, err = pxy.rc.UDPPortManager.Acquire(pxy.name, pxy.cfg.RemotePort)
2019-01-15 00:11:08 +08:00
if err != nil {
2021-03-31 16:57:39 +08:00
return "", fmt.Errorf("acquire port %d error: %v", pxy.cfg.RemotePort, err)
2019-01-15 00:11:08 +08:00
}
defer func() {
if err != nil {
pxy.rc.UDPPortManager.Release(pxy.realBindPort)
2019-01-15 00:11:08 +08:00
}
}()
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)))
2019-01-15 00:11:08 +08:00
if errRet != nil {
err = errRet
return
}
udpConn, errRet := net.ListenUDP("udp", addr)
if errRet != nil {
err = errRet
2019-10-12 20:13:12 +08:00
xl.Warn("listen udp port error: %v", err)
2019-01-15 00:11:08 +08:00
return
}
2019-10-12 20:13:12 +08:00
xl.Info("udp proxy listen port [%d]", pxy.cfg.RemotePort)
2019-01-15 00:11:08 +08:00
pxy.udpConn = udpConn
2020-05-24 17:48:37 +08:00
pxy.sendCh = make(chan *msg.UDPPacket, 1024)
pxy.readCh = make(chan *msg.UDPPacket, 1024)
2019-01-15 00:11:08 +08:00
pxy.checkCloseCh = make(chan int)
// read message from workConn, if it returns any error, notify proxy to start a new workConn
workConnReaderFn := func(conn net.Conn) {
for {
var (
rawMsg msg.Message
errRet error
)
2019-10-12 20:13:12 +08:00
xl.Trace("loop waiting message from udp workConn")
2019-01-15 00:11:08 +08:00
// client will send heartbeat in workConn for keeping alive
2022-08-29 01:02:53 +08:00
_ = conn.SetReadDeadline(time.Now().Add(time.Duration(60) * time.Second))
2019-01-15 00:11:08 +08:00
if rawMsg, errRet = msg.ReadMsg(conn); errRet != nil {
2019-10-12 20:13:12 +08:00
xl.Warn("read from workConn for udp error: %v", errRet)
2022-08-29 01:02:53 +08:00
_ = conn.Close()
2019-01-15 00:11:08 +08:00
// notify proxy to start a new work connection
// ignore error here, it means the proxy is closed
2022-08-29 01:02:53 +08:00
_ = errors.PanicToError(func() {
2019-01-15 00:11:08 +08:00
pxy.checkCloseCh <- 1
})
return
}
2022-08-29 01:02:53 +08:00
if err := conn.SetReadDeadline(time.Time{}); err != nil {
xl.Warn("set read deadline error: %v", err)
}
2019-01-15 00:11:08 +08:00
switch m := rawMsg.(type) {
case *msg.Ping:
2019-10-12 20:13:12 +08:00
xl.Trace("udp work conn get ping message")
2019-01-15 00:11:08 +08:00
continue
2020-05-24 17:48:37 +08:00
case *msg.UDPPacket:
2019-01-15 00:11:08 +08:00
if errRet := errors.PanicToError(func() {
2019-10-12 20:13:12 +08:00
xl.Trace("get udp message from workConn: %s", m.Content)
2019-01-15 00:11:08 +08:00
pxy.readCh <- m
metrics.Server.AddTrafficOut(
pxy.GetName(),
pxy.GetConfigurer().GetBaseConfig().Type,
int64(len(m.Content)),
)
2019-01-15 00:11:08 +08:00
}); errRet != nil {
conn.Close()
2019-10-12 20:13:12 +08:00
xl.Info("reader goroutine for udp work connection closed")
2019-01-15 00:11:08 +08:00
return
}
}
}
}
// send message to workConn
workConnSenderFn := func(conn net.Conn, ctx context.Context) {
var errRet error
for {
select {
case udpMsg, ok := <-pxy.sendCh:
if !ok {
2019-10-12 20:13:12 +08:00
xl.Info("sender goroutine for udp work connection closed")
2019-01-15 00:11:08 +08:00
return
}
if errRet = msg.WriteMsg(conn, udpMsg); errRet != nil {
2019-10-12 20:13:12 +08:00
xl.Info("sender goroutine for udp work connection closed: %v", errRet)
2019-01-15 00:11:08 +08:00
conn.Close()
return
}
2020-05-24 17:48:37 +08:00
xl.Trace("send message to udp workConn: %s", udpMsg.Content)
metrics.Server.AddTrafficIn(
pxy.GetName(),
pxy.GetConfigurer().GetBaseConfig().Type,
2020-05-24 17:48:37 +08:00
int64(len(udpMsg.Content)),
)
continue
2019-01-15 00:11:08 +08:00
case <-ctx.Done():
2019-10-12 20:13:12 +08:00
xl.Info("sender goroutine for udp work connection closed")
2019-01-15 00:11:08 +08:00
return
}
}
}
go func() {
// Sleep a while for waiting control send the NewProxyResp to client.
time.Sleep(500 * time.Millisecond)
for {
2019-03-29 19:01:18 +08:00
workConn, err := pxy.GetWorkConnFromPool(nil, nil)
2019-01-15 00:11:08 +08:00
if err != nil {
time.Sleep(1 * time.Second)
// check if proxy is closed
select {
case _, ok := <-pxy.checkCloseCh:
if !ok {
return
}
default:
}
continue
}
2020-09-07 14:57:23 +08:00
// close the old workConn and replace it with a new one
2019-01-15 00:11:08 +08:00
if pxy.workConn != nil {
pxy.workConn.Close()
}
2020-09-07 14:57:23 +08:00
var rwc io.ReadWriteCloser = workConn
if pxy.cfg.Transport.UseEncryption {
rwc, err = libio.WithEncryption(rwc, []byte(pxy.serverCfg.Auth.Token))
2020-09-07 14:57:23 +08:00
if err != nil {
xl.Error("create encryption stream error: %v", err)
workConn.Close()
continue
}
}
if pxy.cfg.Transport.UseCompression {
2023-05-29 14:10:34 +08:00
rwc = libio.WithCompression(rwc)
2020-09-07 14:57:23 +08:00
}
if pxy.GetLimiter() != nil {
2023-05-29 14:10:34 +08:00
rwc = libio.WrapReadWriteCloser(limit.NewReader(rwc, pxy.GetLimiter()), limit.NewWriter(rwc, pxy.GetLimiter()), func() error {
return rwc.Close()
})
}
2023-11-27 15:47:49 +08:00
pxy.workConn = netpkg.WrapReadWriteCloserToConn(rwc, workConn)
2019-01-15 00:11:08 +08:00
ctx, cancel := context.WithCancel(context.Background())
2020-09-07 14:57:23 +08:00
go workConnReaderFn(pxy.workConn)
go workConnSenderFn(pxy.workConn, ctx)
2019-01-15 00:11:08 +08:00
_, ok := <-pxy.checkCloseCh
cancel()
if !ok {
return
}
}
}()
// Read from user connections and send wrapped udp message to sendCh (forwarded by workConn).
// Client will transfor udp message to local udp service and waiting for response for a while.
// Response will be wrapped to be forwarded by work connection to server.
// Close readCh and sendCh at the end.
go func() {
2020-05-24 17:48:37 +08:00
udp.ForwardUserConn(udpConn, pxy.readCh, pxy.sendCh, int(pxy.serverCfg.UDPPacketSize))
2019-01-15 00:11:08 +08:00
pxy.Close()
}()
return remoteAddr, nil
}
2020-05-24 17:48:37 +08:00
func (pxy *UDPProxy) Close() {
2019-01-15 00:11:08 +08:00
pxy.mu.Lock()
defer pxy.mu.Unlock()
if !pxy.isClosed {
pxy.isClosed = true
pxy.BaseProxy.Close()
if pxy.workConn != nil {
pxy.workConn.Close()
}
pxy.udpConn.Close()
// all channels only closed here
close(pxy.checkCloseCh)
close(pxy.readCh)
close(pxy.sendCh)
}
pxy.rc.UDPPortManager.Release(pxy.realBindPort)
2019-01-15 00:11:08 +08:00
}