mirror of https://github.com/fatedier/frp.git
refactor the code related to xtcp (#3449)
fatedier
2 years ago
committed by
GitHub
parent
9f029e3248
commit
c71efde303
@ -0,0 +1,190 @@
|
||||
// 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 (
|
||||
"io"
|
||||
"net"
|
||||
"strconv"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/fatedier/golib/errors"
|
||||
frpIo "github.com/fatedier/golib/io"
|
||||
|
||||
"github.com/fatedier/frp/pkg/config"
|
||||
"github.com/fatedier/frp/pkg/msg"
|
||||
"github.com/fatedier/frp/pkg/proto/udp"
|
||||
"github.com/fatedier/frp/pkg/util/limit"
|
||||
frpNet "github.com/fatedier/frp/pkg/util/net"
|
||||
)
|
||||
|
||||
type SUDPProxy struct {
|
||||
*BaseProxy
|
||||
|
||||
cfg *config.SUDPProxyConf
|
||||
|
||||
localAddr *net.UDPAddr
|
||||
|
||||
closeCh chan struct{}
|
||||
}
|
||||
|
||||
func (pxy *SUDPProxy) Run() (err error) {
|
||||
pxy.localAddr, err = net.ResolveUDPAddr("udp", net.JoinHostPort(pxy.cfg.LocalIP, strconv.Itoa(pxy.cfg.LocalPort)))
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
func (pxy *SUDPProxy) Close() {
|
||||
pxy.mu.Lock()
|
||||
defer pxy.mu.Unlock()
|
||||
select {
|
||||
case <-pxy.closeCh:
|
||||
return
|
||||
default:
|
||||
close(pxy.closeCh)
|
||||
}
|
||||
}
|
||||
|
||||
func (pxy *SUDPProxy) InWorkConn(conn net.Conn, m *msg.StartWorkConn) {
|
||||
xl := pxy.xl
|
||||
xl.Info("incoming a new work connection for sudp proxy, %s", conn.RemoteAddr().String())
|
||||
|
||||
var rwc io.ReadWriteCloser = conn
|
||||
var err error
|
||||
if pxy.limiter != nil {
|
||||
rwc = frpIo.WrapReadWriteCloser(limit.NewReader(conn, pxy.limiter), limit.NewWriter(conn, pxy.limiter), func() error {
|
||||
return conn.Close()
|
||||
})
|
||||
}
|
||||
if pxy.cfg.UseEncryption {
|
||||
rwc, err = frpIo.WithEncryption(rwc, []byte(pxy.clientCfg.Token))
|
||||
if err != nil {
|
||||
conn.Close()
|
||||
xl.Error("create encryption stream error: %v", err)
|
||||
return
|
||||
}
|
||||
}
|
||||
if pxy.cfg.UseCompression {
|
||||
rwc = frpIo.WithCompression(rwc)
|
||||
}
|
||||
conn = frpNet.WrapReadWriteCloserToConn(rwc, conn)
|
||||
|
||||
workConn := conn
|
||||
readCh := make(chan *msg.UDPPacket, 1024)
|
||||
sendCh := make(chan msg.Message, 1024)
|
||||
isClose := false
|
||||
|
||||
mu := &sync.Mutex{}
|
||||
|
||||
closeFn := func() {
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
if isClose {
|
||||
return
|
||||
}
|
||||
|
||||
isClose = true
|
||||
if workConn != nil {
|
||||
workConn.Close()
|
||||
}
|
||||
close(readCh)
|
||||
close(sendCh)
|
||||
}
|
||||
|
||||
// udp service <- frpc <- frps <- frpc visitor <- user
|
||||
workConnReaderFn := func(conn net.Conn, readCh chan *msg.UDPPacket) {
|
||||
defer closeFn()
|
||||
|
||||
for {
|
||||
// first to check sudp proxy is closed or not
|
||||
select {
|
||||
case <-pxy.closeCh:
|
||||
xl.Trace("frpc sudp proxy is closed")
|
||||
return
|
||||
default:
|
||||
}
|
||||
|
||||
var udpMsg msg.UDPPacket
|
||||
if errRet := msg.ReadMsgInto(conn, &udpMsg); errRet != nil {
|
||||
xl.Warn("read from workConn for sudp error: %v", errRet)
|
||||
return
|
||||
}
|
||||
|
||||
if errRet := errors.PanicToError(func() {
|
||||
readCh <- &udpMsg
|
||||
}); errRet != nil {
|
||||
xl.Warn("reader goroutine for sudp work connection closed: %v", errRet)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// udp service -> frpc -> frps -> frpc visitor -> user
|
||||
workConnSenderFn := func(conn net.Conn, sendCh chan msg.Message) {
|
||||
defer func() {
|
||||
closeFn()
|
||||
xl.Info("writer goroutine for sudp work connection closed")
|
||||
}()
|
||||
|
||||
var errRet error
|
||||
for rawMsg := range sendCh {
|
||||
switch m := rawMsg.(type) {
|
||||
case *msg.UDPPacket:
|
||||
xl.Trace("frpc send udp package to frpc visitor, [udp local: %v, remote: %v], [tcp work conn local: %v, remote: %v]",
|
||||
m.LocalAddr.String(), m.RemoteAddr.String(), conn.LocalAddr().String(), conn.RemoteAddr().String())
|
||||
case *msg.Ping:
|
||||
xl.Trace("frpc send ping message to frpc visitor")
|
||||
}
|
||||
|
||||
if errRet = msg.WriteMsg(conn, rawMsg); errRet != nil {
|
||||
xl.Error("sudp work write error: %v", errRet)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
heartbeatFn := func(sendCh chan msg.Message) {
|
||||
ticker := time.NewTicker(30 * time.Second)
|
||||
defer func() {
|
||||
ticker.Stop()
|
||||
closeFn()
|
||||
}()
|
||||
|
||||
var errRet error
|
||||
for {
|
||||
select {
|
||||
case <-ticker.C:
|
||||
if errRet = errors.PanicToError(func() {
|
||||
sendCh <- &msg.Ping{}
|
||||
}); errRet != nil {
|
||||
xl.Warn("heartbeat goroutine for sudp work connection closed")
|
||||
return
|
||||
}
|
||||
case <-pxy.closeCh:
|
||||
xl.Trace("frpc sudp proxy is closed")
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
go workConnSenderFn(workConn, sendCh)
|
||||
go workConnReaderFn(workConn, readCh)
|
||||
go heartbeatFn(sendCh)
|
||||
|
||||
udp.Forwarder(pxy.localAddr, readCh, sendCh, int(pxy.clientCfg.UDPPacketSize))
|
||||
}
|
||||
@ -0,0 +1,157 @@
|
||||
// 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 (
|
||||
"io"
|
||||
"net"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"github.com/fatedier/golib/errors"
|
||||
frpIo "github.com/fatedier/golib/io"
|
||||
|
||||
"github.com/fatedier/frp/pkg/config"
|
||||
"github.com/fatedier/frp/pkg/msg"
|
||||
"github.com/fatedier/frp/pkg/proto/udp"
|
||||
"github.com/fatedier/frp/pkg/util/limit"
|
||||
frpNet "github.com/fatedier/frp/pkg/util/net"
|
||||
)
|
||||
|
||||
// UDP
|
||||
type UDPProxy struct {
|
||||
*BaseProxy
|
||||
|
||||
cfg *config.UDPProxyConf
|
||||
|
||||
localAddr *net.UDPAddr
|
||||
readCh chan *msg.UDPPacket
|
||||
|
||||
// include msg.UDPPacket and msg.Ping
|
||||
sendCh chan msg.Message
|
||||
workConn net.Conn
|
||||
}
|
||||
|
||||
func (pxy *UDPProxy) Run() (err error) {
|
||||
pxy.localAddr, err = net.ResolveUDPAddr("udp", net.JoinHostPort(pxy.cfg.LocalIP, strconv.Itoa(pxy.cfg.LocalPort)))
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
func (pxy *UDPProxy) Close() {
|
||||
pxy.mu.Lock()
|
||||
defer pxy.mu.Unlock()
|
||||
|
||||
if !pxy.closed {
|
||||
pxy.closed = true
|
||||
if pxy.workConn != nil {
|
||||
pxy.workConn.Close()
|
||||
}
|
||||
if pxy.readCh != nil {
|
||||
close(pxy.readCh)
|
||||
}
|
||||
if pxy.sendCh != nil {
|
||||
close(pxy.sendCh)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (pxy *UDPProxy) InWorkConn(conn net.Conn, m *msg.StartWorkConn) {
|
||||
xl := pxy.xl
|
||||
xl.Info("incoming a new work connection for udp proxy, %s", conn.RemoteAddr().String())
|
||||
// close resources releated with old workConn
|
||||
pxy.Close()
|
||||
|
||||
var rwc io.ReadWriteCloser = conn
|
||||
var err error
|
||||
if pxy.limiter != nil {
|
||||
rwc = frpIo.WrapReadWriteCloser(limit.NewReader(conn, pxy.limiter), limit.NewWriter(conn, pxy.limiter), func() error {
|
||||
return conn.Close()
|
||||
})
|
||||
}
|
||||
if pxy.cfg.UseEncryption {
|
||||
rwc, err = frpIo.WithEncryption(rwc, []byte(pxy.clientCfg.Token))
|
||||
if err != nil {
|
||||
conn.Close()
|
||||
xl.Error("create encryption stream error: %v", err)
|
||||
return
|
||||
}
|
||||
}
|
||||
if pxy.cfg.UseCompression {
|
||||
rwc = frpIo.WithCompression(rwc)
|
||||
}
|
||||
conn = frpNet.WrapReadWriteCloserToConn(rwc, conn)
|
||||
|
||||
pxy.mu.Lock()
|
||||
pxy.workConn = conn
|
||||
pxy.readCh = make(chan *msg.UDPPacket, 1024)
|
||||
pxy.sendCh = make(chan msg.Message, 1024)
|
||||
pxy.closed = false
|
||||
pxy.mu.Unlock()
|
||||
|
||||
workConnReaderFn := func(conn net.Conn, readCh chan *msg.UDPPacket) {
|
||||
for {
|
||||
var udpMsg msg.UDPPacket
|
||||
if errRet := msg.ReadMsgInto(conn, &udpMsg); errRet != nil {
|
||||
xl.Warn("read from workConn for udp error: %v", errRet)
|
||||
return
|
||||
}
|
||||
if errRet := errors.PanicToError(func() {
|
||||
xl.Trace("get udp package from workConn: %s", udpMsg.Content)
|
||||
readCh <- &udpMsg
|
||||
}); errRet != nil {
|
||||
xl.Info("reader goroutine for udp work connection closed: %v", errRet)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
workConnSenderFn := func(conn net.Conn, sendCh chan msg.Message) {
|
||||
defer func() {
|
||||
xl.Info("writer goroutine for udp work connection closed")
|
||||
}()
|
||||
var errRet error
|
||||
for rawMsg := range sendCh {
|
||||
switch m := rawMsg.(type) {
|
||||
case *msg.UDPPacket:
|
||||
xl.Trace("send udp package to workConn: %s", m.Content)
|
||||
case *msg.Ping:
|
||||
xl.Trace("send ping message to udp workConn")
|
||||
}
|
||||
if errRet = msg.WriteMsg(conn, rawMsg); errRet != nil {
|
||||
xl.Error("udp work write error: %v", errRet)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
heartbeatFn := func(sendCh chan msg.Message) {
|
||||
var errRet error
|
||||
for {
|
||||
time.Sleep(time.Duration(30) * time.Second)
|
||||
if errRet = errors.PanicToError(func() {
|
||||
sendCh <- &msg.Ping{}
|
||||
}); errRet != nil {
|
||||
xl.Trace("heartbeat goroutine for udp work connection closed")
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
go workConnSenderFn(pxy.workConn, pxy.sendCh)
|
||||
go workConnReaderFn(pxy.workConn, pxy.readCh)
|
||||
go heartbeatFn(pxy.sendCh)
|
||||
udp.Forwarder(pxy.localAddr, pxy.readCh, pxy.sendCh, int(pxy.clientCfg.UDPPacketSize))
|
||||
}
|
||||
@ -0,0 +1,200 @@
|
||||
// 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 (
|
||||
"io"
|
||||
"net"
|
||||
"time"
|
||||
|
||||
fmux "github.com/hashicorp/yamux"
|
||||
"github.com/quic-go/quic-go"
|
||||
|
||||
"github.com/fatedier/frp/pkg/config"
|
||||
"github.com/fatedier/frp/pkg/msg"
|
||||
"github.com/fatedier/frp/pkg/nathole"
|
||||
plugin "github.com/fatedier/frp/pkg/plugin/client"
|
||||
"github.com/fatedier/frp/pkg/transport"
|
||||
frpNet "github.com/fatedier/frp/pkg/util/net"
|
||||
)
|
||||
|
||||
// XTCP
|
||||
type XTCPProxy struct {
|
||||
*BaseProxy
|
||||
|
||||
cfg *config.XTCPProxyConf
|
||||
proxyPlugin plugin.Plugin
|
||||
}
|
||||
|
||||
func (pxy *XTCPProxy) Run() (err error) {
|
||||
if pxy.cfg.Plugin != "" {
|
||||
pxy.proxyPlugin, err = plugin.Create(pxy.cfg.Plugin, pxy.cfg.PluginParams)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
func (pxy *XTCPProxy) Close() {
|
||||
if pxy.proxyPlugin != nil {
|
||||
pxy.proxyPlugin.Close()
|
||||
}
|
||||
}
|
||||
|
||||
func (pxy *XTCPProxy) InWorkConn(conn net.Conn, startWorkConnMsg *msg.StartWorkConn) {
|
||||
xl := pxy.xl
|
||||
defer conn.Close()
|
||||
var natHoleSidMsg msg.NatHoleSid
|
||||
err := msg.ReadMsgInto(conn, &natHoleSidMsg)
|
||||
if err != nil {
|
||||
xl.Error("xtcp read from workConn error: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
prepareResult, err := nathole.Prepare([]string{pxy.clientCfg.NatHoleSTUNServer})
|
||||
if err != nil {
|
||||
xl.Warn("nathole prepare error: %v", err)
|
||||
return
|
||||
}
|
||||
xl.Info("nathole prepare success, nat type: %s, behavior: %s, addresses: %v, assistedAddresses: %v",
|
||||
prepareResult.NatType, prepareResult.Behavior, prepareResult.Addrs, prepareResult.AssistedAddrs)
|
||||
defer prepareResult.ListenConn.Close()
|
||||
|
||||
// send NatHoleClient msg to server
|
||||
transactionID := nathole.NewTransactionID()
|
||||
natHoleClientMsg := &msg.NatHoleClient{
|
||||
TransactionID: transactionID,
|
||||
ProxyName: pxy.cfg.ProxyName,
|
||||
Sid: natHoleSidMsg.Sid,
|
||||
MappedAddrs: prepareResult.Addrs,
|
||||
AssistedAddrs: prepareResult.AssistedAddrs,
|
||||
}
|
||||
|
||||
natHoleRespMsg, err := nathole.ExchangeInfo(pxy.ctx, pxy.msgTransporter, transactionID, natHoleClientMsg, 5*time.Second)
|
||||
if err != nil {
|
||||
xl.Warn("nathole exchange info error: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
xl.Info("get natHoleRespMsg, sid [%s], protocol [%s], candidate address %v, assisted address %v, detectBehavior: %+v",
|
||||
natHoleRespMsg.Sid, natHoleRespMsg.Protocol, natHoleRespMsg.CandidateAddrs,
|
||||
natHoleRespMsg.AssistedAddrs, natHoleRespMsg.DetectBehavior)
|
||||
|
||||
listenConn := prepareResult.ListenConn
|
||||
newListenConn, raddr, err := nathole.MakeHole(pxy.ctx, listenConn, natHoleRespMsg, []byte(pxy.cfg.Sk))
|
||||
if err != nil {
|
||||
listenConn.Close()
|
||||
xl.Warn("make hole error: %v", err)
|
||||
_ = pxy.msgTransporter.Send(&msg.NatHoleReport{
|
||||
Sid: natHoleRespMsg.Sid,
|
||||
Success: false,
|
||||
})
|
||||
return
|
||||
}
|
||||
listenConn = newListenConn
|
||||
xl.Info("establishing nat hole connection successful, sid [%s], remoteAddr [%s]", natHoleRespMsg.Sid, raddr)
|
||||
|
||||
_ = pxy.msgTransporter.Send(&msg.NatHoleReport{
|
||||
Sid: natHoleRespMsg.Sid,
|
||||
Success: true,
|
||||
})
|
||||
|
||||
if natHoleRespMsg.Protocol == "kcp" {
|
||||
pxy.listenByKCP(listenConn, raddr, startWorkConnMsg)
|
||||
return
|
||||
}
|
||||
|
||||
// default is quic
|
||||
pxy.listenByQUIC(listenConn, raddr, startWorkConnMsg)
|
||||
}
|
||||
|
||||
func (pxy *XTCPProxy) listenByKCP(listenConn *net.UDPConn, raddr *net.UDPAddr, startWorkConnMsg *msg.StartWorkConn) {
|
||||
xl := pxy.xl
|
||||
listenConn.Close()
|
||||
laddr, _ := net.ResolveUDPAddr("udp", listenConn.LocalAddr().String())
|
||||
lConn, err := net.DialUDP("udp", laddr, raddr)
|
||||
if err != nil {
|
||||
xl.Warn("dial udp error: %v", err)
|
||||
return
|
||||
}
|
||||
defer lConn.Close()
|
||||
|
||||
remote, err := frpNet.NewKCPConnFromUDP(lConn, true, raddr.String())
|
||||
if err != nil {
|
||||
xl.Warn("create kcp connection from udp connection error: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
fmuxCfg := fmux.DefaultConfig()
|
||||
fmuxCfg.KeepAliveInterval = 10 * time.Second
|
||||
fmuxCfg.MaxStreamWindowSize = 2 * 1024 * 1024
|
||||
fmuxCfg.LogOutput = io.Discard
|
||||
session, err := fmux.Server(remote, fmuxCfg)
|
||||
if err != nil {
|
||||
xl.Error("create mux session error: %v", err)
|
||||
return
|
||||
}
|
||||
defer session.Close()
|
||||
|
||||
for {
|
||||
muxConn, err := session.Accept()
|
||||
if err != nil {
|
||||
xl.Error("accept connection error: %v", err)
|
||||
return
|
||||
}
|
||||
go HandleTCPWorkConnection(pxy.ctx, &pxy.cfg.LocalSvrConf, pxy.proxyPlugin, pxy.cfg.GetBaseInfo(), pxy.limiter,
|
||||
muxConn, []byte(pxy.cfg.Sk), startWorkConnMsg)
|
||||
}
|
||||
}
|
||||
|
||||
func (pxy *XTCPProxy) listenByQUIC(listenConn *net.UDPConn, _ *net.UDPAddr, startWorkConnMsg *msg.StartWorkConn) {
|
||||
xl := pxy.xl
|
||||
defer listenConn.Close()
|
||||
|
||||
tlsConfig, err := transport.NewServerTLSConfig("", "", "")
|
||||
if err != nil {
|
||||
xl.Warn("create tls config error: %v", err)
|
||||
return
|
||||
}
|
||||
tlsConfig.NextProtos = []string{"frp"}
|
||||
quicListener, err := quic.Listen(listenConn, tlsConfig,
|
||||
&quic.Config{
|
||||
MaxIdleTimeout: time.Duration(pxy.clientCfg.QUICMaxIdleTimeout) * time.Second,
|
||||
MaxIncomingStreams: int64(pxy.clientCfg.QUICMaxIncomingStreams),
|
||||
KeepAlivePeriod: time.Duration(pxy.clientCfg.QUICKeepalivePeriod) * time.Second,
|
||||
},
|
||||
)
|
||||
if err != nil {
|
||||
xl.Warn("dial quic error: %v", err)
|
||||
return
|
||||
}
|
||||
// only accept one connection from raddr
|
||||
c, err := quicListener.Accept(pxy.ctx)
|
||||
if err != nil {
|
||||
xl.Error("quic accept connection error: %v", err)
|
||||
return
|
||||
}
|
||||
for {
|
||||
stream, err := c.AcceptStream(pxy.ctx)
|
||||
if err != nil {
|
||||
xl.Debug("quic accept stream error: %v", err)
|
||||
_ = c.CloseWithError(0, "")
|
||||
return
|
||||
}
|
||||
go HandleTCPWorkConnection(pxy.ctx, &pxy.cfg.LocalSvrConf, pxy.proxyPlugin, pxy.cfg.GetBaseInfo(), pxy.limiter,
|
||||
frpNet.QuicStreamToNetConn(stream, c), []byte(pxy.cfg.Sk), startWorkConnMsg)
|
||||
}
|
||||
}
|
||||
@ -1,575 +0,0 @@
|
||||
// Copyright 2017 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 client
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"strconv"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/fatedier/golib/errors"
|
||||
frpIo "github.com/fatedier/golib/io"
|
||||
"github.com/fatedier/golib/pool"
|
||||
fmux "github.com/hashicorp/yamux"
|
||||
|
||||
"github.com/fatedier/frp/pkg/config"
|
||||
"github.com/fatedier/frp/pkg/msg"
|
||||
"github.com/fatedier/frp/pkg/proto/udp"
|
||||
frpNet "github.com/fatedier/frp/pkg/util/net"
|
||||
"github.com/fatedier/frp/pkg/util/util"
|
||||
"github.com/fatedier/frp/pkg/util/xlog"
|
||||
)
|
||||
|
||||
// Visitor is used for forward traffics from local port to remote service.
|
||||
type Visitor interface {
|
||||
Run() error
|
||||
Close()
|
||||
}
|
||||
|
||||
func NewVisitor(ctx context.Context, ctl *Control, cfg config.VisitorConf) (visitor Visitor) {
|
||||
xl := xlog.FromContextSafe(ctx).Spawn().AppendPrefix(cfg.GetBaseInfo().ProxyName)
|
||||
baseVisitor := BaseVisitor{
|
||||
ctl: ctl,
|
||||
ctx: xlog.NewContext(ctx, xl),
|
||||
}
|
||||
switch cfg := cfg.(type) {
|
||||
case *config.STCPVisitorConf:
|
||||
visitor = &STCPVisitor{
|
||||
BaseVisitor: &baseVisitor,
|
||||
cfg: cfg,
|
||||
}
|
||||
case *config.XTCPVisitorConf:
|
||||
visitor = &XTCPVisitor{
|
||||
BaseVisitor: &baseVisitor,
|
||||
cfg: cfg,
|
||||
}
|
||||
case *config.SUDPVisitorConf:
|
||||
visitor = &SUDPVisitor{
|
||||
BaseVisitor: &baseVisitor,
|
||||
cfg: cfg,
|
||||
checkCloseCh: make(chan struct{}),
|
||||
}
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
type BaseVisitor struct {
|
||||
ctl *Control
|
||||
l net.Listener
|
||||
|
||||
mu sync.RWMutex
|
||||
ctx context.Context
|
||||
}
|
||||
|
||||
type STCPVisitor struct {
|
||||
*BaseVisitor
|
||||
|
||||
cfg *config.STCPVisitorConf
|
||||
}
|
||||
|
||||
func (sv *STCPVisitor) Run() (err error) {
|
||||
sv.l, err = net.Listen("tcp", net.JoinHostPort(sv.cfg.BindAddr, strconv.Itoa(sv.cfg.BindPort)))
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
go sv.worker()
|
||||
return
|
||||
}
|
||||
|
||||
func (sv *STCPVisitor) Close() {
|
||||
sv.l.Close()
|
||||
}
|
||||
|
||||
func (sv *STCPVisitor) worker() {
|
||||
xl := xlog.FromContextSafe(sv.ctx)
|
||||
for {
|
||||
conn, err := sv.l.Accept()
|
||||
if err != nil {
|
||||
xl.Warn("stcp local listener closed")
|
||||
return
|
||||
}
|
||||
|
||||
go sv.handleConn(conn)
|
||||
}
|
||||
}
|
||||
|
||||
func (sv *STCPVisitor) handleConn(userConn net.Conn) {
|
||||
xl := xlog.FromContextSafe(sv.ctx)
|
||||
defer userConn.Close()
|
||||
|
||||
xl.Debug("get a new stcp user connection")
|
||||
visitorConn, err := sv.ctl.connectServer()
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer visitorConn.Close()
|
||||
|
||||
now := time.Now().Unix()
|
||||
newVisitorConnMsg := &msg.NewVisitorConn{
|
||||
ProxyName: sv.cfg.ServerName,
|
||||
SignKey: util.GetAuthKey(sv.cfg.Sk, now),
|
||||
Timestamp: now,
|
||||
UseEncryption: sv.cfg.UseEncryption,
|
||||
UseCompression: sv.cfg.UseCompression,
|
||||
}
|
||||
err = msg.WriteMsg(visitorConn, newVisitorConnMsg)
|
||||
if err != nil {
|
||||
xl.Warn("send newVisitorConnMsg to server error: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
var newVisitorConnRespMsg msg.NewVisitorConnResp
|
||||
_ = visitorConn.SetReadDeadline(time.Now().Add(10 * time.Second))
|
||||
err = msg.ReadMsgInto(visitorConn, &newVisitorConnRespMsg)
|
||||
if err != nil {
|
||||
xl.Warn("get newVisitorConnRespMsg error: %v", err)
|
||||
return
|
||||
}
|
||||
_ = visitorConn.SetReadDeadline(time.Time{})
|
||||
|
||||
if newVisitorConnRespMsg.Error != "" {
|
||||
xl.Warn("start new visitor connection error: %s", newVisitorConnRespMsg.Error)
|
||||
return
|
||||
}
|
||||
|
||||
var remote io.ReadWriteCloser
|
||||
remote = visitorConn
|
||||
if sv.cfg.UseEncryption {
|
||||
remote, err = frpIo.WithEncryption(remote, []byte(sv.cfg.Sk))
|
||||
if err != nil {
|
||||
xl.Error("create encryption stream error: %v", err)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
if sv.cfg.UseCompression {
|
||||
remote = frpIo.WithCompression(remote)
|
||||
}
|
||||
|
||||
frpIo.Join(userConn, remote)
|
||||
}
|
||||
|
||||
type XTCPVisitor struct {
|
||||
*BaseVisitor
|
||||
|
||||
cfg *config.XTCPVisitorConf
|
||||
}
|
||||
|
||||
func (sv *XTCPVisitor) Run() (err error) {
|
||||
sv.l, err = net.Listen("tcp", net.JoinHostPort(sv.cfg.BindAddr, strconv.Itoa(sv.cfg.BindPort)))
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
go sv.worker()
|
||||
return
|
||||
}
|
||||
|
||||
func (sv *XTCPVisitor) Close() {
|
||||
sv.l.Close()
|
||||
}
|
||||
|
||||
func (sv *XTCPVisitor) worker() {
|
||||
xl := xlog.FromContextSafe(sv.ctx)
|
||||
for {
|
||||
conn, err := sv.l.Accept()
|
||||
if err != nil {
|
||||
xl.Warn("xtcp local listener closed")
|
||||
return
|
||||
}
|
||||
|
||||
go sv.handleConn(conn)
|
||||
}
|
||||
}
|
||||
|
||||
func (sv *XTCPVisitor) handleConn(userConn net.Conn) {
|
||||
xl := xlog.FromContextSafe(sv.ctx)
|
||||
defer userConn.Close()
|
||||
|
||||
xl.Debug("get a new xtcp user connection")
|
||||
if sv.ctl.serverUDPPort == 0 {
|
||||
xl.Error("xtcp is not supported by server")
|
||||
return
|
||||
}
|
||||
|
||||
serverAddr := sv.ctl.clientCfg.NatHoleServerAddr
|
||||
if serverAddr == "" {
|
||||
serverAddr = sv.ctl.clientCfg.ServerAddr
|
||||
}
|
||||
raddr, err := net.ResolveUDPAddr("udp",
|
||||
net.JoinHostPort(serverAddr, strconv.Itoa(sv.ctl.serverUDPPort)))
|
||||
if err != nil {
|
||||
xl.Error("resolve server UDP addr error")
|
||||
return
|
||||
}
|
||||
|
||||
visitorConn, err := net.DialUDP("udp", nil, raddr)
|
||||
if err != nil {
|
||||
xl.Warn("dial server udp addr error: %v", err)
|
||||
return
|
||||
}
|
||||
defer visitorConn.Close()
|
||||
|
||||
now := time.Now().Unix()
|
||||
natHoleVisitorMsg := &msg.NatHoleVisitor{
|
||||
ProxyName: sv.cfg.ServerName,
|
||||
SignKey: util.GetAuthKey(sv.cfg.Sk, now),
|
||||
Timestamp: now,
|
||||
}
|
||||
err = msg.WriteMsg(visitorConn, natHoleVisitorMsg)
|
||||
if err != nil {
|
||||
xl.Warn("send natHoleVisitorMsg to server error: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
// Wait for client address at most 10 seconds.
|
||||
var natHoleRespMsg msg.NatHoleResp
|
||||
_ = visitorConn.SetReadDeadline(time.Now().Add(10 * time.Second))
|
||||
buf := pool.GetBuf(1024)
|
||||
n, err := visitorConn.Read(buf)
|
||||
if err != nil {
|
||||
xl.Warn("get natHoleRespMsg error: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
err = msg.ReadMsgInto(bytes.NewReader(buf[:n]), &natHoleRespMsg)
|
||||
if err != nil {
|
||||
xl.Warn("get natHoleRespMsg error: %v", err)
|
||||
return
|
||||
}
|
||||
_ = visitorConn.SetReadDeadline(time.Time{})
|
||||
pool.PutBuf(buf)
|
||||
|
||||
if natHoleRespMsg.Error != "" {
|
||||
xl.Error("natHoleRespMsg get error info: %s", natHoleRespMsg.Error)
|
||||
return
|
||||
}
|
||||
|
||||
xl.Trace("get natHoleRespMsg, sid [%s], client address [%s], visitor address [%s]", natHoleRespMsg.Sid, natHoleRespMsg.ClientAddr, natHoleRespMsg.VisitorAddr)
|
||||
|
||||
// Close visitorConn, so we can use it's local address.
|
||||
visitorConn.Close()
|
||||
|
||||
// send sid message to client
|
||||
laddr, _ := net.ResolveUDPAddr("udp", visitorConn.LocalAddr().String())
|
||||
daddr, err := net.ResolveUDPAddr("udp", natHoleRespMsg.ClientAddr)
|
||||
if err != nil {
|
||||
xl.Error("resolve client udp address error: %v", err)
|
||||
return
|
||||
}
|
||||
lConn, err := net.DialUDP("udp", laddr, daddr)
|
||||
if err != nil {
|
||||
xl.Error("dial client udp address error: %v", err)
|
||||
return
|
||||
}
|
||||
defer lConn.Close()
|
||||
|
||||
if _, err := lConn.Write([]byte(natHoleRespMsg.Sid)); err != nil {
|
||||
xl.Error("write sid error: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
// read ack sid from client
|
||||
sidBuf := pool.GetBuf(1024)
|
||||
_ = lConn.SetReadDeadline(time.Now().Add(8 * time.Second))
|
||||
n, err = lConn.Read(sidBuf)
|
||||
if err != nil {
|
||||
xl.Warn("get sid from client error: %v", err)
|
||||
return
|
||||
}
|
||||
_ = lConn.SetReadDeadline(time.Time{})
|
||||
if string(sidBuf[:n]) != natHoleRespMsg.Sid {
|
||||
xl.Warn("incorrect sid from client")
|
||||
return
|
||||
}
|
||||
pool.PutBuf(sidBuf)
|
||||
|
||||
xl.Info("nat hole connection make success, sid [%s]", natHoleRespMsg.Sid)
|
||||
|
||||
// wrap kcp connection
|
||||
var remote io.ReadWriteCloser
|
||||
remote, err = frpNet.NewKCPConnFromUDP(lConn, true, natHoleRespMsg.ClientAddr)
|
||||
if err != nil {
|
||||
xl.Error("create kcp connection from udp connection error: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
fmuxCfg := fmux.DefaultConfig()
|
||||
fmuxCfg.KeepAliveInterval = 5 * time.Second
|
||||
fmuxCfg.LogOutput = io.Discard
|
||||
sess, err := fmux.Client(remote, fmuxCfg)
|
||||
if err != nil {
|
||||
xl.Error("create yamux session error: %v", err)
|
||||
return
|
||||
}
|
||||
defer sess.Close()
|
||||
muxConn, err := sess.Open()
|
||||
if err != nil {
|
||||
xl.Error("open yamux stream error: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
var muxConnRWCloser io.ReadWriteCloser = muxConn
|
||||
if sv.cfg.UseEncryption {
|
||||
muxConnRWCloser, err = frpIo.WithEncryption(muxConnRWCloser, []byte(sv.cfg.Sk))
|
||||
if err != nil {
|
||||
xl.Error("create encryption stream error: %v", err)
|
||||
return
|
||||
}
|
||||
}
|
||||
if sv.cfg.UseCompression {
|
||||
muxConnRWCloser = frpIo.WithCompression(muxConnRWCloser)
|
||||
}
|
||||
|
||||
_, _, errs := frpIo.Join(userConn, muxConnRWCloser)
|
||||
xl.Debug("join connections closed")
|
||||
if len(errs) > 0 {
|
||||
xl.Trace("join connections errors: %v", errs)
|
||||
}
|
||||
}
|
||||
|
||||
type SUDPVisitor struct {
|
||||
*BaseVisitor
|
||||
|
||||
checkCloseCh chan struct{}
|
||||
// udpConn is the listener of udp packet
|
||||
udpConn *net.UDPConn
|
||||
readCh chan *msg.UDPPacket
|
||||
sendCh chan *msg.UDPPacket
|
||||
|
||||
cfg *config.SUDPVisitorConf
|
||||
}
|
||||
|
||||
// SUDP Run start listen a udp port
|
||||
func (sv *SUDPVisitor) Run() (err error) {
|
||||
xl := xlog.FromContextSafe(sv.ctx)
|
||||
|
||||
addr, err := net.ResolveUDPAddr("udp", net.JoinHostPort(sv.cfg.BindAddr, strconv.Itoa(sv.cfg.BindPort)))
|
||||
if err != nil {
|
||||
return fmt.Errorf("sudp ResolveUDPAddr error: %v", err)
|
||||
}
|
||||
|
||||
sv.udpConn, err = net.ListenUDP("udp", addr)
|
||||
if err != nil {
|
||||
return fmt.Errorf("listen udp port %s error: %v", addr.String(), err)
|
||||
}
|
||||
|
||||
sv.sendCh = make(chan *msg.UDPPacket, 1024)
|
||||
sv.readCh = make(chan *msg.UDPPacket, 1024)
|
||||
|
||||
xl.Info("sudp start to work, listen on %s", addr)
|
||||
|
||||
go sv.dispatcher()
|
||||
go udp.ForwardUserConn(sv.udpConn, sv.readCh, sv.sendCh, int(sv.ctl.clientCfg.UDPPacketSize))
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
func (sv *SUDPVisitor) dispatcher() {
|
||||
xl := xlog.FromContextSafe(sv.ctx)
|
||||
|
||||
var (
|
||||
visitorConn net.Conn
|
||||
err error
|
||||
|
||||
firstPacket *msg.UDPPacket
|
||||
)
|
||||
|
||||
for {
|
||||
select {
|
||||
case firstPacket = <-sv.sendCh:
|
||||
if firstPacket == nil {
|
||||
xl.Info("frpc sudp visitor proxy is closed")
|
||||
return
|
||||
}
|
||||
case <-sv.checkCloseCh:
|
||||
xl.Info("frpc sudp visitor proxy is closed")
|
||||
return
|
||||
}
|
||||
|
||||
visitorConn, err = sv.getNewVisitorConn()
|
||||
if err != nil {
|
||||
xl.Warn("newVisitorConn to frps error: %v, try to reconnect", err)
|
||||
continue
|
||||
}
|
||||
|
||||
// visitorConn always be closed when worker done.
|
||||
sv.worker(visitorConn, firstPacket)
|
||||
|
||||
select {
|
||||
case <-sv.checkCloseCh:
|
||||
return
|
||||
default:
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (sv *SUDPVisitor) worker(workConn net.Conn, firstPacket *msg.UDPPacket) {
|
||||
xl := xlog.FromContextSafe(sv.ctx)
|
||||
xl.Debug("starting sudp proxy worker")
|
||||
|
||||
wg := &sync.WaitGroup{}
|
||||
wg.Add(2)
|
||||
closeCh := make(chan struct{})
|
||||
|
||||
// udp service -> frpc -> frps -> frpc visitor -> user
|
||||
workConnReaderFn := func(conn net.Conn) {
|
||||
defer func() {
|
||||
conn.Close()
|
||||
close(closeCh)
|
||||
wg.Done()
|
||||
}()
|
||||
|
||||
for {
|
||||
var (
|
||||
rawMsg msg.Message
|
||||
errRet error
|
||||
)
|
||||
|
||||
// frpc will send heartbeat in workConn to frpc visitor for keeping alive
|
||||
_ = conn.SetReadDeadline(time.Now().Add(60 * time.Second))
|
||||
if rawMsg, errRet = msg.ReadMsg(conn); errRet != nil {
|
||||
xl.Warn("read from workconn for user udp conn error: %v", errRet)
|
||||
return
|
||||
}
|
||||
|
||||
_ = conn.SetReadDeadline(time.Time{})
|
||||
switch m := rawMsg.(type) {
|
||||
case *msg.Ping:
|
||||
xl.Debug("frpc visitor get ping message from frpc")
|
||||
continue
|
||||
case *msg.UDPPacket:
|
||||
if errRet := errors.PanicToError(func() {
|
||||
sv.readCh <- m
|
||||
xl.Trace("frpc visitor get udp packet from workConn: %s", m.Content)
|
||||
}); errRet != nil {
|
||||
xl.Info("reader goroutine for udp work connection closed")
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// udp service <- frpc <- frps <- frpc visitor <- user
|
||||
workConnSenderFn := func(conn net.Conn) {
|
||||
defer func() {
|
||||
conn.Close()
|
||||
wg.Done()
|
||||
}()
|
||||
|
||||
var errRet error
|
||||
if firstPacket != nil {
|
||||
if errRet = msg.WriteMsg(conn, firstPacket); errRet != nil {
|
||||
xl.Warn("sender goroutine for udp work connection closed: %v", errRet)
|
||||
return
|
||||
}
|
||||
xl.Trace("send udp package to workConn: %s", firstPacket.Content)
|
||||
}
|
||||
|
||||
for {
|
||||
select {
|
||||
case udpMsg, ok := <-sv.sendCh:
|
||||
if !ok {
|
||||
xl.Info("sender goroutine for udp work connection closed")
|
||||
return
|
||||
}
|
||||
|
||||
if errRet = msg.WriteMsg(conn, udpMsg); errRet != nil {
|
||||
xl.Warn("sender goroutine for udp work connection closed: %v", errRet)
|
||||
return
|
||||
}
|
||||
xl.Trace("send udp package to workConn: %s", udpMsg.Content)
|
||||
case <-closeCh:
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
go workConnReaderFn(workConn)
|
||||
go workConnSenderFn(workConn)
|
||||
|
||||
wg.Wait()
|
||||
xl.Info("sudp worker is closed")
|
||||
}
|
||||
|
||||
func (sv *SUDPVisitor) getNewVisitorConn() (net.Conn, error) {
|
||||
xl := xlog.FromContextSafe(sv.ctx)
|
||||
visitorConn, err := sv.ctl.connectServer()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("frpc connect frps error: %v", err)
|
||||
}
|
||||
|
||||
now := time.Now().Unix()
|
||||
newVisitorConnMsg := &msg.NewVisitorConn{
|
||||
ProxyName: sv.cfg.ServerName,
|
||||
SignKey: util.GetAuthKey(sv.cfg.Sk, now),
|
||||
Timestamp: now,
|
||||
UseEncryption: sv.cfg.UseEncryption,
|
||||
UseCompression: sv.cfg.UseCompression,
|
||||
}
|
||||
err = msg.WriteMsg(visitorConn, newVisitorConnMsg)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("frpc send newVisitorConnMsg to frps error: %v", err)
|
||||
}
|
||||
|
||||
var newVisitorConnRespMsg msg.NewVisitorConnResp
|
||||
_ = visitorConn.SetReadDeadline(time.Now().Add(10 * time.Second))
|
||||
err = msg.ReadMsgInto(visitorConn, &newVisitorConnRespMsg)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("frpc read newVisitorConnRespMsg error: %v", err)
|
||||
}
|
||||
_ = visitorConn.SetReadDeadline(time.Time{})
|
||||
|
||||
if newVisitorConnRespMsg.Error != "" {
|
||||
return nil, fmt.Errorf("start new visitor connection error: %s", newVisitorConnRespMsg.Error)
|
||||
}
|
||||
|
||||
var remote io.ReadWriteCloser
|
||||
remote = visitorConn
|
||||
if sv.cfg.UseEncryption {
|
||||
remote, err = frpIo.WithEncryption(remote, []byte(sv.cfg.Sk))
|
||||
if err != nil {
|
||||
xl.Error("create encryption stream error: %v", err)
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
if sv.cfg.UseCompression {
|
||||
remote = frpIo.WithCompression(remote)
|
||||
}
|
||||
return frpNet.WrapReadWriteCloserToConn(remote, visitorConn), nil
|
||||
}
|
||||
|
||||
func (sv *SUDPVisitor) Close() {
|
||||
sv.mu.Lock()
|
||||
defer sv.mu.Unlock()
|
||||
|
||||
select {
|
||||
case <-sv.checkCloseCh:
|
||||
return
|
||||
default:
|
||||
close(sv.checkCloseCh)
|
||||
}
|
||||
if sv.udpConn != nil {
|
||||
sv.udpConn.Close()
|
||||
}
|
||||
close(sv.readCh)
|
||||
close(sv.sendCh)
|
||||
}
|
||||
@ -0,0 +1,118 @@
|
||||
// Copyright 2017 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 visitor
|
||||
|
||||
import (
|
||||
"io"
|
||||
"net"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
frpIo "github.com/fatedier/golib/io"
|
||||
|
||||
"github.com/fatedier/frp/pkg/config"
|
||||
"github.com/fatedier/frp/pkg/msg"
|
||||
"github.com/fatedier/frp/pkg/util/util"
|
||||
"github.com/fatedier/frp/pkg/util/xlog"
|
||||
)
|
||||
|
||||
type STCPVisitor struct {
|
||||
*BaseVisitor
|
||||
|
||||
cfg *config.STCPVisitorConf
|
||||
}
|
||||
|
||||
func (sv *STCPVisitor) Run() (err error) {
|
||||
sv.l, err = net.Listen("tcp", net.JoinHostPort(sv.cfg.BindAddr, strconv.Itoa(sv.cfg.BindPort)))
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
go sv.worker()
|
||||
return
|
||||
}
|
||||
|
||||
func (sv *STCPVisitor) Close() {
|
||||
sv.l.Close()
|
||||
}
|
||||
|
||||
func (sv *STCPVisitor) worker() {
|
||||
xl := xlog.FromContextSafe(sv.ctx)
|
||||
for {
|
||||
conn, err := sv.l.Accept()
|
||||
if err != nil {
|
||||
xl.Warn("stcp local listener closed")
|
||||
return
|
||||
}
|
||||
|
||||
go sv.handleConn(conn)
|
||||
}
|
||||
}
|
||||
|
||||
func (sv *STCPVisitor) handleConn(userConn net.Conn) {
|
||||
xl := xlog.FromContextSafe(sv.ctx)
|
||||
defer userConn.Close()
|
||||
|
||||
xl.Debug("get a new stcp user connection")
|
||||
visitorConn, err := sv.connectServer()
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer visitorConn.Close()
|
||||
|
||||
now := time.Now().Unix()
|
||||
newVisitorConnMsg := &msg.NewVisitorConn{
|
||||
ProxyName: sv.cfg.ServerName,
|
||||
SignKey: util.GetAuthKey(sv.cfg.Sk, now),
|
||||
Timestamp: now,
|
||||
UseEncryption: sv.cfg.UseEncryption,
|
||||
UseCompression: sv.cfg.UseCompression,
|
||||
}
|
||||
err = msg.WriteMsg(visitorConn, newVisitorConnMsg)
|
||||
if err != nil {
|
||||
xl.Warn("send newVisitorConnMsg to server error: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
var newVisitorConnRespMsg msg.NewVisitorConnResp
|
||||
_ = visitorConn.SetReadDeadline(time.Now().Add(10 * time.Second))
|
||||
err = msg.ReadMsgInto(visitorConn, &newVisitorConnRespMsg)
|
||||
if err != nil {
|
||||
xl.Warn("get newVisitorConnRespMsg error: %v", err)
|
||||
return
|
||||
}
|
||||
_ = visitorConn.SetReadDeadline(time.Time{})
|
||||
|
||||
if newVisitorConnRespMsg.Error != "" {
|
||||
xl.Warn("start new visitor connection error: %s", newVisitorConnRespMsg.Error)
|
||||
return
|
||||
}
|
||||
|
||||
var remote io.ReadWriteCloser
|
||||
remote = visitorConn
|
||||
if sv.cfg.UseEncryption {
|
||||
remote, err = frpIo.WithEncryption(remote, []byte(sv.cfg.Sk))
|
||||
if err != nil {
|
||||
xl.Error("create encryption stream error: %v", err)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
if sv.cfg.UseCompression {
|
||||
remote = frpIo.WithCompression(remote)
|
||||
}
|
||||
|
||||
frpIo.Join(userConn, remote)
|
||||
}
|
||||
@ -0,0 +1,262 @@
|
||||
// Copyright 2017 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 visitor
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"strconv"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/fatedier/golib/errors"
|
||||
frpIo "github.com/fatedier/golib/io"
|
||||
|
||||
"github.com/fatedier/frp/pkg/config"
|
||||
"github.com/fatedier/frp/pkg/msg"
|
||||
"github.com/fatedier/frp/pkg/proto/udp"
|
||||
frpNet "github.com/fatedier/frp/pkg/util/net"
|
||||
"github.com/fatedier/frp/pkg/util/util"
|
||||
"github.com/fatedier/frp/pkg/util/xlog"
|
||||
)
|
||||
|
||||
type SUDPVisitor struct {
|
||||
*BaseVisitor
|
||||
|
||||
checkCloseCh chan struct{}
|
||||
// udpConn is the listener of udp packet
|
||||
udpConn *net.UDPConn
|
||||
readCh chan *msg.UDPPacket
|
||||
sendCh chan *msg.UDPPacket
|
||||
|
||||
cfg *config.SUDPVisitorConf
|
||||
}
|
||||
|
||||
// SUDP Run start listen a udp port
|
||||
func (sv *SUDPVisitor) Run() (err error) {
|
||||
xl := xlog.FromContextSafe(sv.ctx)
|
||||
|
||||
addr, err := net.ResolveUDPAddr("udp", net.JoinHostPort(sv.cfg.BindAddr, strconv.Itoa(sv.cfg.BindPort)))
|
||||
if err != nil {
|
||||
return fmt.Errorf("sudp ResolveUDPAddr error: %v", err)
|
||||
}
|
||||
|
||||
sv.udpConn, err = net.ListenUDP("udp", addr)
|
||||
if err != nil {
|
||||
return fmt.Errorf("listen udp port %s error: %v", addr.String(), err)
|
||||
}
|
||||
|
||||
sv.sendCh = make(chan *msg.UDPPacket, 1024)
|
||||
sv.readCh = make(chan *msg.UDPPacket, 1024)
|
||||
|
||||
xl.Info("sudp start to work, listen on %s", addr)
|
||||
|
||||
go sv.dispatcher()
|
||||
go udp.ForwardUserConn(sv.udpConn, sv.readCh, sv.sendCh, int(sv.clientCfg.UDPPacketSize))
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
func (sv *SUDPVisitor) dispatcher() {
|
||||
xl := xlog.FromContextSafe(sv.ctx)
|
||||
|
||||
var (
|
||||
visitorConn net.Conn
|
||||
err error
|
||||
|
||||
firstPacket *msg.UDPPacket
|
||||
)
|
||||
|
||||
for {
|
||||
select {
|
||||
case firstPacket = <-sv.sendCh:
|
||||
if firstPacket == nil {
|
||||
xl.Info("frpc sudp visitor proxy is closed")
|
||||
return
|
||||
}
|
||||
case <-sv.checkCloseCh:
|
||||
xl.Info("frpc sudp visitor proxy is closed")
|
||||
return
|
||||
}
|
||||
|
||||
visitorConn, err = sv.getNewVisitorConn()
|
||||
if err != nil {
|
||||
xl.Warn("newVisitorConn to frps error: %v, try to reconnect", err)
|
||||
continue
|
||||
}
|
||||
|
||||
// visitorConn always be closed when worker done.
|
||||
sv.worker(visitorConn, firstPacket)
|
||||
|
||||
select {
|
||||
case <-sv.checkCloseCh:
|
||||
return
|
||||
default:
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (sv *SUDPVisitor) worker(workConn net.Conn, firstPacket *msg.UDPPacket) {
|
||||
xl := xlog.FromContextSafe(sv.ctx)
|
||||
xl.Debug("starting sudp proxy worker")
|
||||
|
||||
wg := &sync.WaitGroup{}
|
||||
wg.Add(2)
|
||||
closeCh := make(chan struct{})
|
||||
|
||||
// udp service -> frpc -> frps -> frpc visitor -> user
|
||||
workConnReaderFn := func(conn net.Conn) {
|
||||
defer func() {
|
||||
conn.Close()
|
||||
close(closeCh)
|
||||
wg.Done()
|
||||
}()
|
||||
|
||||
for {
|
||||
var (
|
||||
rawMsg msg.Message
|
||||
errRet error
|
||||
)
|
||||
|
||||
// frpc will send heartbeat in workConn to frpc visitor for keeping alive
|
||||
_ = conn.SetReadDeadline(time.Now().Add(60 * time.Second))
|
||||
if rawMsg, errRet = msg.ReadMsg(conn); errRet != nil {
|
||||
xl.Warn("read from workconn for user udp conn error: %v", errRet)
|
||||
return
|
||||
}
|
||||
|
||||
_ = conn.SetReadDeadline(time.Time{})
|
||||
switch m := rawMsg.(type) {
|
||||
case *msg.Ping:
|
||||
xl.Debug("frpc visitor get ping message from frpc")
|
||||
continue
|
||||
case *msg.UDPPacket:
|
||||
if errRet := errors.PanicToError(func() {
|
||||
sv.readCh <- m
|
||||
xl.Trace("frpc visitor get udp packet from workConn: %s", m.Content)
|
||||
}); errRet != nil {
|
||||
xl.Info("reader goroutine for udp work connection closed")
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// udp service <- frpc <- frps <- frpc visitor <- user
|
||||
workConnSenderFn := func(conn net.Conn) {
|
||||
defer func() {
|
||||
conn.Close()
|
||||
wg.Done()
|
||||
}()
|
||||
|
||||
var errRet error
|
||||
if firstPacket != nil {
|
||||
if errRet = msg.WriteMsg(conn, firstPacket); errRet != nil {
|
||||
xl.Warn("sender goroutine for udp work connection closed: %v", errRet)
|
||||
return
|
||||
}
|
||||
xl.Trace("send udp package to workConn: %s", firstPacket.Content)
|
||||
}
|
||||
|
||||
for {
|
||||
select {
|
||||
case udpMsg, ok := <-sv.sendCh:
|
||||
if !ok {
|
||||
xl.Info("sender goroutine for udp work connection closed")
|
||||
return
|
||||
}
|
||||
|
||||
if errRet = msg.WriteMsg(conn, udpMsg); errRet != nil {
|
||||
xl.Warn("sender goroutine for udp work connection closed: %v", errRet)
|
||||
return
|
||||
}
|
||||
xl.Trace("send udp package to workConn: %s", udpMsg.Content)
|
||||
case <-closeCh:
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
go workConnReaderFn(workConn)
|
||||
go workConnSenderFn(workConn)
|
||||
|
||||
wg.Wait()
|
||||
xl.Info("sudp worker is closed")
|
||||
}
|
||||
|
||||
func (sv *SUDPVisitor) getNewVisitorConn() (net.Conn, error) {
|
||||
xl := xlog.FromContextSafe(sv.ctx)
|
||||
visitorConn, err := sv.connectServer()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("frpc connect frps error: %v", err)
|
||||
}
|
||||
|
||||
now := time.Now().Unix()
|
||||
newVisitorConnMsg := &msg.NewVisitorConn{
|
||||
ProxyName: sv.cfg.ServerName,
|
||||
SignKey: util.GetAuthKey(sv.cfg.Sk, now),
|
||||
Timestamp: now,
|
||||
UseEncryption: sv.cfg.UseEncryption,
|
||||
UseCompression: sv.cfg.UseCompression,
|
||||
}
|
||||
err = msg.WriteMsg(visitorConn, newVisitorConnMsg)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("frpc send newVisitorConnMsg to frps error: %v", err)
|
||||
}
|
||||
|
||||
var newVisitorConnRespMsg msg.NewVisitorConnResp
|
||||
_ = visitorConn.SetReadDeadline(time.Now().Add(10 * time.Second))
|
||||
err = msg.ReadMsgInto(visitorConn, &newVisitorConnRespMsg)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("frpc read newVisitorConnRespMsg error: %v", err)
|
||||
}
|
||||
_ = visitorConn.SetReadDeadline(time.Time{})
|
||||
|
||||
if newVisitorConnRespMsg.Error != "" {
|
||||
return nil, fmt.Errorf("start new visitor connection error: %s", newVisitorConnRespMsg.Error)
|
||||
}
|
||||
|
||||
var remote io.ReadWriteCloser
|
||||
remote = visitorConn
|
||||
if sv.cfg.UseEncryption {
|
||||
remote, err = frpIo.WithEncryption(remote, []byte(sv.cfg.Sk))
|
||||
if err != nil {
|
||||
xl.Error("create encryption stream error: %v", err)
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
if sv.cfg.UseCompression {
|
||||
remote = frpIo.WithCompression(remote)
|
||||
}
|
||||
return frpNet.WrapReadWriteCloserToConn(remote, visitorConn), nil
|
||||
}
|
||||
|
||||
func (sv *SUDPVisitor) Close() {
|
||||
sv.mu.Lock()
|
||||
defer sv.mu.Unlock()
|
||||
|
||||
select {
|
||||
case <-sv.checkCloseCh:
|
||||
return
|
||||
default:
|
||||
close(sv.checkCloseCh)
|
||||
}
|
||||
if sv.udpConn != nil {
|
||||
sv.udpConn.Close()
|
||||
}
|
||||
close(sv.readCh)
|
||||
close(sv.sendCh)
|
||||
}
|
||||
@ -0,0 +1,77 @@
|
||||
// Copyright 2017 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 visitor
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net"
|
||||
"sync"
|
||||
|
||||
"github.com/fatedier/frp/pkg/config"
|
||||
"github.com/fatedier/frp/pkg/transport"
|
||||
"github.com/fatedier/frp/pkg/util/xlog"
|
||||
)
|
||||
|
||||
// Visitor is used for forward traffics from local port tot remote service.
|
||||
type Visitor interface {
|
||||
Run() error
|
||||
Close()
|
||||
}
|
||||
|
||||
func NewVisitor(
|
||||
ctx context.Context,
|
||||
cfg config.VisitorConf,
|
||||
clientCfg config.ClientCommonConf,
|
||||
connectServer func() (net.Conn, error),
|
||||
msgTransporter transport.MessageTransporter,
|
||||
) (visitor Visitor) {
|
||||
xl := xlog.FromContextSafe(ctx).Spawn().AppendPrefix(cfg.GetBaseInfo().ProxyName)
|
||||
baseVisitor := BaseVisitor{
|
||||
clientCfg: clientCfg,
|
||||
connectServer: connectServer,
|
||||
msgTransporter: msgTransporter,
|
||||
ctx: xlog.NewContext(ctx, xl),
|
||||
}
|
||||
switch cfg := cfg.(type) {
|
||||
case *config.STCPVisitorConf:
|
||||
visitor = &STCPVisitor{
|
||||
BaseVisitor: &baseVisitor,
|
||||
cfg: cfg,
|
||||
}
|
||||
case *config.XTCPVisitorConf:
|
||||
visitor = &XTCPVisitor{
|
||||
BaseVisitor: &baseVisitor,
|
||||
cfg: cfg,
|
||||
startTunnelCh: make(chan struct{}),
|
||||
}
|
||||
case *config.SUDPVisitorConf:
|
||||
visitor = &SUDPVisitor{
|
||||
BaseVisitor: &baseVisitor,
|
||||
cfg: cfg,
|
||||
checkCloseCh: make(chan struct{}),
|
||||
}
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
type BaseVisitor struct {
|
||||
clientCfg config.ClientCommonConf
|
||||
connectServer func() (net.Conn, error)
|
||||
msgTransporter transport.MessageTransporter
|
||||
l net.Listener
|
||||
|
||||
mu sync.RWMutex
|
||||
ctx context.Context
|
||||
}
|
||||
@ -0,0 +1,410 @@
|
||||
// Copyright 2017 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 visitor
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"strconv"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
frpIo "github.com/fatedier/golib/io"
|
||||
fmux "github.com/hashicorp/yamux"
|
||||
quic "github.com/quic-go/quic-go"
|
||||
"golang.org/x/time/rate"
|
||||
|
||||
"github.com/fatedier/frp/pkg/config"
|
||||
"github.com/fatedier/frp/pkg/msg"
|
||||
"github.com/fatedier/frp/pkg/nathole"
|
||||
"github.com/fatedier/frp/pkg/transport"
|
||||
frpNet "github.com/fatedier/frp/pkg/util/net"
|
||||
"github.com/fatedier/frp/pkg/util/util"
|
||||
"github.com/fatedier/frp/pkg/util/xlog"
|
||||
)
|
||||
|
||||
var ErrNoTunnelSession = errors.New("no tunnel session")
|
||||
|
||||
type XTCPVisitor struct {
|
||||
*BaseVisitor
|
||||
session TunnelSession
|
||||
startTunnelCh chan struct{}
|
||||
retryLimiter *rate.Limiter
|
||||
cancel context.CancelFunc
|
||||
|
||||
cfg *config.XTCPVisitorConf
|
||||
}
|
||||
|
||||
func (sv *XTCPVisitor) Run() (err error) {
|
||||
sv.ctx, sv.cancel = context.WithCancel(sv.ctx)
|
||||
|
||||
if sv.cfg.Protocol == "kcp" {
|
||||
sv.session = NewKCPTunnelSession()
|
||||
} else {
|
||||
sv.session = NewQUICTunnelSession(&sv.clientCfg)
|
||||
}
|
||||
|
||||
sv.l, err = net.Listen("tcp", net.JoinHostPort(sv.cfg.BindAddr, strconv.Itoa(sv.cfg.BindPort)))
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
go sv.worker()
|
||||
go sv.processTunnelStartEvents()
|
||||
if sv.cfg.KeepTunnelOpen {
|
||||
sv.retryLimiter = rate.NewLimiter(rate.Every(time.Hour/time.Duration(sv.cfg.MaxRetriesAnHour)), sv.cfg.MaxRetriesAnHour)
|
||||
go sv.keepTunnelOpenWorker()
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
func (sv *XTCPVisitor) Close() {
|
||||
sv.l.Close()
|
||||
sv.cancel()
|
||||
if sv.session != nil {
|
||||
sv.session.Close()
|
||||
}
|
||||
}
|
||||
|
||||
func (sv *XTCPVisitor) worker() {
|
||||
xl := xlog.FromContextSafe(sv.ctx)
|
||||
for {
|
||||
conn, err := sv.l.Accept()
|
||||
if err != nil {
|
||||
xl.Warn("xtcp local listener closed")
|
||||
return
|
||||
}
|
||||
|
||||
go sv.handleConn(conn)
|
||||
}
|
||||
}
|
||||
|
||||
func (sv *XTCPVisitor) processTunnelStartEvents() {
|
||||
for {
|
||||
select {
|
||||
case <-sv.ctx.Done():
|
||||
return
|
||||
case <-sv.startTunnelCh:
|
||||
start := time.Now()
|
||||
sv.makeNatHole()
|
||||
duration := time.Since(start)
|
||||
// avoid too frequently
|
||||
if duration < 10*time.Second {
|
||||
time.Sleep(10*time.Second - duration)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (sv *XTCPVisitor) keepTunnelOpenWorker() {
|
||||
xl := xlog.FromContextSafe(sv.ctx)
|
||||
ticker := time.NewTicker(time.Duration(sv.cfg.MinRetryInterval) * time.Second)
|
||||
defer ticker.Stop()
|
||||
|
||||
sv.startTunnelCh <- struct{}{}
|
||||
for {
|
||||
select {
|
||||
case <-sv.ctx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
xl.Debug("keepTunnelOpenWorker try to check tunnel...")
|
||||
conn, err := sv.getTunnelConn()
|
||||
if err != nil {
|
||||
xl.Warn("keepTunnelOpenWorker get tunnel connection error: %v", err)
|
||||
_ = sv.retryLimiter.Wait(sv.ctx)
|
||||
continue
|
||||
}
|
||||
xl.Debug("keepTunnelOpenWorker check success")
|
||||
if conn != nil {
|
||||
conn.Close()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (sv *XTCPVisitor) handleConn(userConn net.Conn) {
|
||||
xl := xlog.FromContextSafe(sv.ctx)
|
||||
defer userConn.Close()
|
||||
|
||||
xl.Debug("get a new xtcp user connection")
|
||||
|
||||
// Open a tunnel connection to the server. If there is already a successful hole-punching connection,
|
||||
// it will be reused. Otherwise, it will block and wait for a successful hole-punching connection until timeout.
|
||||
tunnelConn, err := sv.openTunnel()
|
||||
if err != nil {
|
||||
xl.Error("open tunnel error: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
var muxConnRWCloser io.ReadWriteCloser = tunnelConn
|
||||
if sv.cfg.UseEncryption {
|
||||
muxConnRWCloser, err = frpIo.WithEncryption(muxConnRWCloser, []byte(sv.cfg.Sk))
|
||||
if err != nil {
|
||||
xl.Error("create encryption stream error: %v", err)
|
||||
return
|
||||
}
|
||||
}
|
||||
if sv.cfg.UseCompression {
|
||||
muxConnRWCloser = frpIo.WithCompression(muxConnRWCloser)
|
||||
}
|
||||
|
||||
_, _, errs := frpIo.Join(userConn, muxConnRWCloser)
|
||||
xl.Debug("join connections closed")
|
||||
if len(errs) > 0 {
|
||||
xl.Trace("join connections errors: %v", errs)
|
||||
}
|
||||
}
|
||||
|
||||
// openTunnel will open a tunnel connection to the target server.
|
||||
func (sv *XTCPVisitor) openTunnel() (conn net.Conn, err error) {
|
||||
xl := xlog.FromContextSafe(sv.ctx)
|
||||
ticker := time.NewTicker(500 * time.Millisecond)
|
||||
defer ticker.Stop()
|
||||
|
||||
timeoutC := time.After(20 * time.Second)
|
||||
immediateTrigger := make(chan struct{}, 1)
|
||||
defer close(immediateTrigger)
|
||||
immediateTrigger <- struct{}{}
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-sv.ctx.Done():
|
||||
return nil, sv.ctx.Err()
|
||||
case <-immediateTrigger:
|
||||
conn, err = sv.getTunnelConn()
|
||||
case <-ticker.C:
|
||||
conn, err = sv.getTunnelConn()
|
||||
case <-timeoutC:
|
||||
return nil, fmt.Errorf("open tunnel timeout")
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
if err != ErrNoTunnelSession {
|
||||
xl.Warn("get tunnel connection error: %v", err)
|
||||
}
|
||||
continue
|
||||
}
|
||||
return conn, nil
|
||||
}
|
||||
}
|
||||
|
||||
func (sv *XTCPVisitor) getTunnelConn() (net.Conn, error) {
|
||||
conn, err := sv.session.OpenConn(sv.ctx)
|
||||
if err == nil {
|
||||
return conn, nil
|
||||
}
|
||||
sv.session.Close()
|
||||
|
||||
select {
|
||||
case sv.startTunnelCh <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// 0. PreCheck
|
||||
// 1. Prepare
|
||||
// 2. ExchangeInfo
|
||||
// 3. MakeNATHole
|
||||
// 4. Create a tunnel session using an underlying UDP connection.
|
||||
func (sv *XTCPVisitor) makeNatHole() {
|
||||
xl := xlog.FromContextSafe(sv.ctx)
|
||||
if err := nathole.PreCheck(sv.ctx, sv.msgTransporter, sv.cfg.ServerName, 5*time.Second); err != nil {
|
||||
xl.Warn("nathole precheck error: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
prepareResult, err := nathole.Prepare([]string{sv.clientCfg.NatHoleSTUNServer})
|
||||
if err != nil {
|
||||
xl.Warn("nathole prepare error: %v", err)
|
||||
return
|
||||
}
|
||||
xl.Info("nathole prepare success, nat type: %s, behavior: %s, addresses: %v, assistedAddresses: %v",
|
||||
prepareResult.NatType, prepareResult.Behavior, prepareResult.Addrs, prepareResult.AssistedAddrs)
|
||||
|
||||
listenConn := prepareResult.ListenConn
|
||||
|
||||
// send NatHoleVisitor to server
|
||||
now := time.Now().Unix()
|
||||
transactionID := nathole.NewTransactionID()
|
||||
natHoleVisitorMsg := &msg.NatHoleVisitor{
|
||||
TransactionID: transactionID,
|
||||
ProxyName: sv.cfg.ServerName,
|
||||
Protocol: sv.cfg.Protocol,
|
||||
SignKey: util.GetAuthKey(sv.cfg.Sk, now),
|
||||
Timestamp: now,
|
||||
MappedAddrs: prepareResult.Addrs,
|
||||
AssistedAddrs: prepareResult.AssistedAddrs,
|
||||
}
|
||||
|
||||
natHoleRespMsg, err := nathole.ExchangeInfo(sv.ctx, sv.msgTransporter, transactionID, natHoleVisitorMsg, 5*time.Second)
|
||||
if err != nil {
|
||||
listenConn.Close()
|
||||
xl.Warn("nathole exchange info error: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
xl.Info("get natHoleRespMsg, sid [%s], protocol [%s], candidate address %v, assisted address %v, detectBehavior: %+v",
|
||||
natHoleRespMsg.Sid, natHoleRespMsg.Protocol, natHoleRespMsg.CandidateAddrs,
|
||||
natHoleRespMsg.AssistedAddrs, natHoleRespMsg.DetectBehavior)
|
||||
|
||||
newListenConn, raddr, err := nathole.MakeHole(sv.ctx, listenConn, natHoleRespMsg, []byte(sv.cfg.Sk))
|
||||
if err != nil {
|
||||
listenConn.Close()
|
||||
xl.Warn("make hole error: %v", err)
|
||||
return
|
||||
}
|
||||
listenConn = newListenConn
|
||||
xl.Info("establishing nat hole connection successful, sid [%s], remoteAddr [%s]", natHoleRespMsg.Sid, raddr)
|
||||
|
||||
if err := sv.session.Init(listenConn, raddr); err != nil {
|
||||
listenConn.Close()
|
||||
xl.Warn("init tunnel session error: %v", err)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
type TunnelSession interface {
|
||||
Init(listenConn *net.UDPConn, raddr *net.UDPAddr) error
|
||||
OpenConn(context.Context) (net.Conn, error)
|
||||
Close()
|
||||
}
|
||||
|
||||
type KCPTunnelSession struct {
|
||||
session *fmux.Session
|
||||
lConn *net.UDPConn
|
||||
mu sync.RWMutex
|
||||
}
|
||||
|
||||
func NewKCPTunnelSession() TunnelSession {
|
||||
return &KCPTunnelSession{}
|
||||
}
|
||||
|
||||
func (ks *KCPTunnelSession) Init(listenConn *net.UDPConn, raddr *net.UDPAddr) error {
|
||||
listenConn.Close()
|
||||
laddr, _ := net.ResolveUDPAddr("udp", listenConn.LocalAddr().String())
|
||||
lConn, err := net.DialUDP("udp", laddr, raddr)
|
||||
if err != nil {
|
||||
return fmt.Errorf("dial udp error: %v", err)
|
||||
}
|
||||
remote, err := frpNet.NewKCPConnFromUDP(lConn, true, raddr.String())
|
||||
if err != nil {
|
||||
return fmt.Errorf("create kcp connection from udp connection error: %v", err)
|
||||
}
|
||||
|
||||
fmuxCfg := fmux.DefaultConfig()
|
||||
fmuxCfg.KeepAliveInterval = 10 * time.Second
|
||||
fmuxCfg.MaxStreamWindowSize = 2 * 1024 * 1024
|
||||
fmuxCfg.LogOutput = io.Discard
|
||||
session, err := fmux.Client(remote, fmuxCfg)
|
||||
if err != nil {
|
||||
remote.Close()
|
||||
return fmt.Errorf("initial client session error: %v", err)
|
||||
}
|
||||
ks.mu.Lock()
|
||||
ks.session = session
|
||||
ks.lConn = lConn
|
||||
ks.mu.Unlock()
|
||||
return nil
|
||||
}
|
||||
|
||||
func (ks *KCPTunnelSession) OpenConn(ctx context.Context) (net.Conn, error) {
|
||||
ks.mu.RLock()
|
||||
defer ks.mu.RUnlock()
|
||||
session := ks.session
|
||||
if session == nil {
|
||||
return nil, ErrNoTunnelSession
|
||||
}
|
||||
return session.Open()
|
||||
}
|
||||
|
||||
func (ks *KCPTunnelSession) Close() {
|
||||
ks.mu.Lock()
|
||||
defer ks.mu.Unlock()
|
||||
if ks.session != nil {
|
||||
_ = ks.session.Close()
|
||||
ks.session = nil
|
||||
}
|
||||
if ks.lConn != nil {
|
||||
_ = ks.lConn.Close()
|
||||
ks.lConn = nil
|
||||
}
|
||||
}
|
||||
|
||||
type QUICTunnelSession struct {
|
||||
session quic.Connection
|
||||
listenConn *net.UDPConn
|
||||
mu sync.RWMutex
|
||||
|
||||
clientCfg *config.ClientCommonConf
|
||||
}
|
||||
|
||||
func NewQUICTunnelSession(clientCfg *config.ClientCommonConf) TunnelSession {
|
||||
return &QUICTunnelSession{
|
||||
clientCfg: clientCfg,
|
||||
}
|
||||
}
|
||||
|
||||
func (qs *QUICTunnelSession) Init(listenConn *net.UDPConn, raddr *net.UDPAddr) error {
|
||||
tlsConfig, err := transport.NewClientTLSConfig("", "", "", raddr.String())
|
||||
if err != nil {
|
||||
return fmt.Errorf("create tls config error: %v", err)
|
||||
}
|
||||
tlsConfig.NextProtos = []string{"frp"}
|
||||
quicConn, err := quic.Dial(listenConn, raddr, raddr.String(), tlsConfig,
|
||||
&quic.Config{
|
||||
MaxIdleTimeout: time.Duration(qs.clientCfg.QUICMaxIdleTimeout) * time.Second,
|
||||
MaxIncomingStreams: int64(qs.clientCfg.QUICMaxIncomingStreams),
|
||||
KeepAlivePeriod: time.Duration(qs.clientCfg.QUICKeepalivePeriod) * time.Second,
|
||||
})
|
||||
if err != nil {
|
||||
return fmt.Errorf("dial quic error: %v", err)
|
||||
}
|
||||
qs.mu.Lock()
|
||||
qs.session = quicConn
|
||||
qs.listenConn = listenConn
|
||||
qs.mu.Unlock()
|
||||
return nil
|
||||
}
|
||||
|
||||
func (qs *QUICTunnelSession) OpenConn(ctx context.Context) (net.Conn, error) {
|
||||
qs.mu.RLock()
|
||||
defer qs.mu.RUnlock()
|
||||
session := qs.session
|
||||
if session == nil {
|
||||
return nil, ErrNoTunnelSession
|
||||
}
|
||||
stream, err := session.OpenStreamSync(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return frpNet.QuicStreamToNetConn(stream, session), nil
|
||||
}
|
||||
|
||||
func (qs *QUICTunnelSession) Close() {
|
||||
qs.mu.Lock()
|
||||
defer qs.mu.Unlock()
|
||||
if qs.session != nil {
|
||||
_ = qs.session.CloseWithError(0, "")
|
||||
qs.session = nil
|
||||
}
|
||||
if qs.listenConn != nil {
|
||||
_ = qs.listenConn.Close()
|
||||
qs.listenConn = nil
|
||||
}
|
||||
}
|
||||
@ -0,0 +1,328 @@
|
||||
// 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 nathole
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/samber/lo"
|
||||
)
|
||||
|
||||
var (
|
||||
// mode 0, both EasyNAT, PublicNetwork is always receiver
|
||||
// sender | receiver, ttl 7
|
||||
// receiver, ttl 7 | sender
|
||||
// sender | receiver, ttl 4
|
||||
// receiver, ttl 4 | sender
|
||||
// sender | receiver
|
||||
// receiver | sender
|
||||
// sender, sendDelayMs 5000 | receiver
|
||||
// sender, sendDelayMs 10000 | receiver
|
||||
// receiver | sender, sendDelayMs 5000
|
||||
// receiver | sender, sendDelayMs 10000
|
||||
mode0Behaviors = []lo.Tuple2[RecommandBehavior, RecommandBehavior]{
|
||||
lo.T2(RecommandBehavior{Role: DetectRoleSender}, RecommandBehavior{Role: DetectRoleReceiver, TTL: 7}),
|
||||
lo.T2(RecommandBehavior{Role: DetectRoleReceiver, TTL: 7}, RecommandBehavior{Role: DetectRoleSender}),
|
||||
lo.T2(RecommandBehavior{Role: DetectRoleSender}, RecommandBehavior{Role: DetectRoleReceiver, TTL: 4}),
|
||||
lo.T2(RecommandBehavior{Role: DetectRoleReceiver, TTL: 4}, RecommandBehavior{Role: DetectRoleSender}),
|
||||
lo.T2(RecommandBehavior{Role: DetectRoleSender}, RecommandBehavior{Role: DetectRoleReceiver}),
|
||||
lo.T2(RecommandBehavior{Role: DetectRoleReceiver}, RecommandBehavior{Role: DetectRoleSender}),
|
||||
lo.T2(RecommandBehavior{Role: DetectRoleSender, SendDelayMs: 5000}, RecommandBehavior{Role: DetectRoleReceiver}),
|
||||
lo.T2(RecommandBehavior{Role: DetectRoleSender, SendDelayMs: 10000}, RecommandBehavior{Role: DetectRoleReceiver}),
|
||||
lo.T2(RecommandBehavior{Role: DetectRoleReceiver}, RecommandBehavior{Role: DetectRoleSender, SendDelayMs: 5000}),
|
||||
lo.T2(RecommandBehavior{Role: DetectRoleReceiver}, RecommandBehavior{Role: DetectRoleSender, SendDelayMs: 10000}),
|
||||
}
|
||||
|
||||
// mode 1, HardNAT is sender, EasyNAT is receiver, port changes is regular
|
||||
// sender | receiver, ttl 7, portsRangeNumber max 10
|
||||
// sender, sendDelayMs 2000 | receiver, ttl 7, portsRangeNumber max 10
|
||||
// sender | receiver, ttl 4, portsRangeNumber max 10
|
||||
// sender, sendDelayMs 2000 | receiver, ttl 4, portsRangeNumber max 10
|
||||
// sender | receiver, portsRangeNumber max 10
|
||||
// sender, sendDelayMs 2000 | receiver, portsRangeNumber max 10
|
||||
mode1Behaviors = []lo.Tuple2[RecommandBehavior, RecommandBehavior]{
|
||||
lo.T2(RecommandBehavior{Role: DetectRoleSender}, RecommandBehavior{Role: DetectRoleReceiver, TTL: 7, PortsRangeNumber: 10}),
|
||||
lo.T2(RecommandBehavior{Role: DetectRoleSender, SendDelayMs: 2000}, RecommandBehavior{Role: DetectRoleReceiver, TTL: 7, PortsRangeNumber: 10}),
|
||||
lo.T2(RecommandBehavior{Role: DetectRoleSender}, RecommandBehavior{Role: DetectRoleReceiver, TTL: 4, PortsRangeNumber: 10}),
|
||||
lo.T2(RecommandBehavior{Role: DetectRoleSender, SendDelayMs: 2000}, RecommandBehavior{Role: DetectRoleReceiver, TTL: 4, PortsRangeNumber: 10}),
|
||||
lo.T2(RecommandBehavior{Role: DetectRoleSender}, RecommandBehavior{Role: DetectRoleReceiver, PortsRangeNumber: 10}),
|
||||
lo.T2(RecommandBehavior{Role: DetectRoleSender, SendDelayMs: 2000}, RecommandBehavior{Role: DetectRoleReceiver, PortsRangeNumber: 10}),
|
||||
}
|
||||
|
||||
// mode 2, HardNAT is receiver, EasyNAT is sender
|
||||
// sender, portsRandomNumber 1000, sendDelayMs 2000 | receiver, listen 256 ports, ttl 7
|
||||
// sender, portsRandomNumber 1000, sendDelayMs 2000 | receiver, listen 256 ports, ttl 4
|
||||
// sender, portsRandomNumber 1000, sendDelayMs 2000 | receiver, listen 256 ports
|
||||
mode2Behaviors = []lo.Tuple2[RecommandBehavior, RecommandBehavior]{
|
||||
lo.T2(
|
||||
RecommandBehavior{Role: DetectRoleSender, PortsRandomNumber: 1000, SendDelayMs: 2000},
|
||||
RecommandBehavior{Role: DetectRoleReceiver, ListenRandomPorts: 256, TTL: 7},
|
||||
),
|
||||
lo.T2(
|
||||
RecommandBehavior{Role: DetectRoleSender, PortsRandomNumber: 1000, SendDelayMs: 2000},
|
||||
RecommandBehavior{Role: DetectRoleReceiver, ListenRandomPorts: 256, TTL: 4},
|
||||
),
|
||||
lo.T2(
|
||||
RecommandBehavior{Role: DetectRoleSender, PortsRandomNumber: 1000, SendDelayMs: 2000},
|
||||
RecommandBehavior{Role: DetectRoleReceiver, ListenRandomPorts: 256},
|
||||
),
|
||||
}
|
||||
|
||||
// mode 3, For HardNAT & HardNAT, both changes in the ports are regular
|
||||
// sender, portsRangeNumber 10 | receiver, ttl 7, portsRangeNumber 10
|
||||
// sender, portsRangeNumber 10 | receiver, ttl 4, portsRangeNumber 10
|
||||
// sender, portsRangeNumber 10 | receiver, portsRangeNumber 10
|
||||
// receiver, ttl 7, portsRangeNumber 10 | sender, portsRangeNumber 10
|
||||
// receiver, ttl 4, portsRangeNumber 10 | sender, portsRangeNumber 10
|
||||
// receiver, portsRangeNumber 10 | sender, portsRangeNumber 10
|
||||
mode3Behaviors = []lo.Tuple2[RecommandBehavior, RecommandBehavior]{
|
||||
lo.T2(RecommandBehavior{Role: DetectRoleSender, PortsRangeNumber: 10}, RecommandBehavior{Role: DetectRoleReceiver, TTL: 7, PortsRangeNumber: 10}),
|
||||
lo.T2(RecommandBehavior{Role: DetectRoleSender, PortsRangeNumber: 10}, RecommandBehavior{Role: DetectRoleReceiver, TTL: 4, PortsRangeNumber: 10}),
|
||||
lo.T2(RecommandBehavior{Role: DetectRoleSender, PortsRangeNumber: 10}, RecommandBehavior{Role: DetectRoleReceiver, PortsRangeNumber: 10}),
|
||||
lo.T2(RecommandBehavior{Role: DetectRoleReceiver, TTL: 7, PortsRangeNumber: 10}, RecommandBehavior{Role: DetectRoleSender, PortsRangeNumber: 10}),
|
||||
lo.T2(RecommandBehavior{Role: DetectRoleReceiver, TTL: 4, PortsRangeNumber: 10}, RecommandBehavior{Role: DetectRoleSender, PortsRangeNumber: 10}),
|
||||
lo.T2(RecommandBehavior{Role: DetectRoleReceiver, PortsRangeNumber: 10}, RecommandBehavior{Role: DetectRoleSender, PortsRangeNumber: 10}),
|
||||
}
|
||||
|
||||
// mode 4, Regular ports changes are usually the sender.
|
||||
// sender, portsRandomNumber 1000, sendDelayMs: 2000 | receiver, listen 256 ports, ttl 7, portsRangeNumber 10
|
||||
// sender, portsRandomNumber 1000, sendDelayMs: 2000 | receiver, listen 256 ports, ttl 4, portsRangeNumber 10
|
||||
// sender, portsRandomNumber 1000, SendDelayMs: 2000 | receiver, listen 256 ports, portsRangeNumber 10
|
||||
mode4Behaviors = []lo.Tuple2[RecommandBehavior, RecommandBehavior]{
|
||||
lo.T2(
|
||||
RecommandBehavior{Role: DetectRoleSender, PortsRandomNumber: 1000, SendDelayMs: 2000},
|
||||
RecommandBehavior{Role: DetectRoleReceiver, ListenRandomPorts: 256, TTL: 7, PortsRangeNumber: 10},
|
||||
),
|
||||
lo.T2(
|
||||
RecommandBehavior{Role: DetectRoleSender, PortsRandomNumber: 1000, SendDelayMs: 2000},
|
||||
RecommandBehavior{Role: DetectRoleReceiver, ListenRandomPorts: 256, TTL: 4, PortsRangeNumber: 10},
|
||||
),
|
||||
lo.T2(
|
||||
RecommandBehavior{Role: DetectRoleSender, PortsRandomNumber: 1000, SendDelayMs: 2000},
|
||||
RecommandBehavior{Role: DetectRoleReceiver, ListenRandomPorts: 256, PortsRangeNumber: 10},
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
func getBehaviorByMode(mode int) []lo.Tuple2[RecommandBehavior, RecommandBehavior] {
|
||||
switch mode {
|
||||
case 0:
|
||||
return mode0Behaviors
|
||||
case 1:
|
||||
return mode1Behaviors
|
||||
case 2:
|
||||
return mode2Behaviors
|
||||
case 3:
|
||||
return mode3Behaviors
|
||||
case 4:
|
||||
return mode4Behaviors
|
||||
}
|
||||
// default
|
||||
return mode0Behaviors
|
||||
}
|
||||
|
||||
func getBehaviorByModeAndIndex(mode int, index int) (RecommandBehavior, RecommandBehavior) {
|
||||
behaviors := getBehaviorByMode(mode)
|
||||
if index >= len(behaviors) {
|
||||
return RecommandBehavior{}, RecommandBehavior{}
|
||||
}
|
||||
return behaviors[index].A, behaviors[index].B
|
||||
}
|
||||
|
||||
func getBehaviorScoresByMode(mode int, defaultScore int) []*BehaviorScore {
|
||||
return getBehaviorScoresByMode2(mode, defaultScore, defaultScore)
|
||||
}
|
||||
|
||||
func getBehaviorScoresByMode2(mode int, senderScore, receiverScore int) []*BehaviorScore {
|
||||
behaviors := getBehaviorByMode(mode)
|
||||
scores := make([]*BehaviorScore, 0, len(behaviors))
|
||||
for i := 0; i < len(behaviors); i++ {
|
||||
score := receiverScore
|
||||
if behaviors[i].A.Role == DetectRoleSender {
|
||||
score = senderScore
|
||||
}
|
||||
scores = append(scores, &BehaviorScore{Mode: mode, Index: i, Score: score})
|
||||
}
|
||||
return scores
|
||||
}
|
||||
|
||||
type RecommandBehavior struct {
|
||||
Role string
|
||||
TTL int
|
||||
SendDelayMs int
|
||||
PortsRangeNumber int
|
||||
PortsRandomNumber int
|
||||
ListenRandomPorts int
|
||||
}
|
||||
|
||||
type MakeHoleRecords struct {
|
||||
mu sync.Mutex
|
||||
scores []*BehaviorScore
|
||||
LastUpdateTime time.Time
|
||||
}
|
||||
|
||||
func NewMakeHoleRecords(c, v *NatFeature) *MakeHoleRecords {
|
||||
scores := []*BehaviorScore{}
|
||||
easyCount, hardCount, portsChangedRegularCount := ClassifyFeatureCount([]*NatFeature{c, v})
|
||||
appendMode0 := func() {
|
||||
switch {
|
||||
case c.PublicNetwork:
|
||||
scores = append(scores, getBehaviorScoresByMode2(DetectMode0, 0, 1)...)
|
||||
case v.PublicNetwork:
|
||||
scores = append(scores, getBehaviorScoresByMode2(DetectMode0, 1, 0)...)
|
||||
default:
|
||||
scores = append(scores, getBehaviorScoresByMode(DetectMode0, 0)...)
|
||||
}
|
||||
}
|
||||
|
||||
switch {
|
||||
case easyCount == 2:
|
||||
appendMode0()
|
||||
case hardCount == 1 && portsChangedRegularCount == 1:
|
||||
scores = append(scores, getBehaviorScoresByMode(DetectMode1, 0)...)
|
||||
scores = append(scores, getBehaviorScoresByMode(DetectMode2, 0)...)
|
||||
appendMode0()
|
||||
case hardCount == 1 && portsChangedRegularCount == 0:
|
||||
scores = append(scores, getBehaviorScoresByMode(DetectMode2, 0)...)
|
||||
scores = append(scores, getBehaviorScoresByMode(DetectMode1, 0)...)
|
||||
appendMode0()
|
||||
case hardCount == 2 && portsChangedRegularCount == 2:
|
||||
scores = append(scores, getBehaviorScoresByMode(DetectMode3, 0)...)
|
||||
scores = append(scores, getBehaviorScoresByMode(DetectMode4, 0)...)
|
||||
case hardCount == 2 && portsChangedRegularCount == 1:
|
||||
scores = append(scores, getBehaviorScoresByMode(DetectMode4, 0)...)
|
||||
default:
|
||||
// hard to make hole, just trying it out.
|
||||
scores = append(scores, getBehaviorScoresByMode(DetectMode0, 1)...)
|
||||
scores = append(scores, getBehaviorScoresByMode(DetectMode1, 1)...)
|
||||
scores = append(scores, getBehaviorScoresByMode(DetectMode3, 1)...)
|
||||
}
|
||||
return &MakeHoleRecords{scores: scores, LastUpdateTime: time.Now()}
|
||||
}
|
||||
|
||||
func (mhr *MakeHoleRecords) ReportSuccess(mode int, index int) {
|
||||
mhr.mu.Lock()
|
||||
defer mhr.mu.Unlock()
|
||||
mhr.LastUpdateTime = time.Now()
|
||||
for i := range mhr.scores {
|
||||
score := mhr.scores[i]
|
||||
if score.Mode != mode || score.Index != index {
|
||||
continue
|
||||
}
|
||||
|
||||
score.Score += 2
|
||||
score.Score = lo.Min([]int{score.Score, 10})
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
func (mhr *MakeHoleRecords) Recommand() (mode, index int) {
|
||||
mhr.mu.Lock()
|
||||
defer mhr.mu.Unlock()
|
||||
|
||||
maxScore := lo.MaxBy(mhr.scores, func(item, max *BehaviorScore) bool {
|
||||
return item.Score > max.Score
|
||||
})
|
||||
if maxScore == nil {
|
||||
return 0, 0
|
||||
}
|
||||
maxScore.Score--
|
||||
mhr.LastUpdateTime = time.Now()
|
||||
return maxScore.Mode, maxScore.Index
|
||||
}
|
||||
|
||||
type BehaviorScore struct {
|
||||
Mode int
|
||||
Index int
|
||||
// between -10 and 10
|
||||
Score int
|
||||
}
|
||||
|
||||
type Analyzer struct {
|
||||
// key is client ip + visitor ip
|
||||
records map[string]*MakeHoleRecords
|
||||
dataReserveDuration time.Duration
|
||||
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
func NewAnalyzer(dataReserveDuration time.Duration) *Analyzer {
|
||||
return &Analyzer{
|
||||
records: make(map[string]*MakeHoleRecords),
|
||||
dataReserveDuration: dataReserveDuration,
|
||||
}
|
||||
}
|
||||
|
||||
func (a *Analyzer) GetRecommandBehaviors(key string, c, v *NatFeature) (mode, index int, _ RecommandBehavior, _ RecommandBehavior) {
|
||||
a.mu.Lock()
|
||||
records, ok := a.records[key]
|
||||
if !ok {
|
||||
records = NewMakeHoleRecords(c, v)
|
||||
a.records[key] = records
|
||||
}
|
||||
a.mu.Unlock()
|
||||
|
||||
mode, index = records.Recommand()
|
||||
cBehavior, vBehavior := getBehaviorByModeAndIndex(mode, index)
|
||||
|
||||
switch mode {
|
||||
case DetectMode1:
|
||||
// HardNAT is always the sender
|
||||
if c.NatType == EasyNAT {
|
||||
cBehavior, vBehavior = vBehavior, cBehavior
|
||||
}
|
||||
case DetectMode2:
|
||||
// HardNAT is always the receiver
|
||||
if c.NatType == HardNAT {
|
||||
cBehavior, vBehavior = vBehavior, cBehavior
|
||||
}
|
||||
case DetectMode4:
|
||||
// Regular ports changes is always the sender
|
||||
if !c.RegularPortsChange {
|
||||
cBehavior, vBehavior = vBehavior, cBehavior
|
||||
}
|
||||
}
|
||||
return mode, index, cBehavior, vBehavior
|
||||
}
|
||||
|
||||
func (a *Analyzer) ReportSuccess(key string, mode, index int) {
|
||||
a.mu.Lock()
|
||||
records, ok := a.records[key]
|
||||
a.mu.Unlock()
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
records.ReportSuccess(mode, index)
|
||||
}
|
||||
|
||||
func (a *Analyzer) Clean() (int, int) {
|
||||
now := time.Now()
|
||||
total := 0
|
||||
count := 0
|
||||
|
||||
// cleanup 10w records may take 5ms
|
||||
a.mu.Lock()
|
||||
defer a.mu.Unlock()
|
||||
total = len(a.records)
|
||||
// clean up records that have not been used for a period of time.
|
||||
for key, records := range a.records {
|
||||
if now.Sub(records.LastUpdateTime) > a.dataReserveDuration {
|
||||
delete(a.records, key)
|
||||
count++
|
||||
}
|
||||
}
|
||||
return count, total
|
||||
}
|
||||
@ -0,0 +1,382 @@
|
||||
// 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 nathole
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/md5"
|
||||
"encoding/hex"
|
||||
"fmt"
|
||||
"net"
|
||||
"strconv"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/fatedier/golib/errors"
|
||||
"github.com/samber/lo"
|
||||
"golang.org/x/sync/errgroup"
|
||||
|
||||
"github.com/fatedier/frp/pkg/msg"
|
||||
"github.com/fatedier/frp/pkg/transport"
|
||||
"github.com/fatedier/frp/pkg/util/log"
|
||||
"github.com/fatedier/frp/pkg/util/util"
|
||||
)
|
||||
|
||||
// NatHoleTimeout seconds.
|
||||
var NatHoleTimeout int64 = 10
|
||||
|
||||
func NewTransactionID() string {
|
||||
id, _ := util.RandID()
|
||||
return fmt.Sprintf("%d%s", time.Now().Unix(), id)
|
||||
}
|
||||
|
||||
type ClientCfg struct {
|
||||
name string
|
||||
sk string
|
||||
sidCh chan string
|
||||
}
|
||||
|
||||
type Session struct {
|
||||
sid string
|
||||
analysisKey string
|
||||
recommandMode int
|
||||
recommandIndex int
|
||||
|
||||
visitorMsg *msg.NatHoleVisitor
|
||||
visitorTransporter transport.MessageTransporter
|
||||
vResp *msg.NatHoleResp
|
||||
vNatFeature *NatFeature
|
||||
vBehavior RecommandBehavior
|
||||
|
||||
clientMsg *msg.NatHoleClient
|
||||
clientTransporter transport.MessageTransporter
|
||||
cResp *msg.NatHoleResp
|
||||
cNatFeature *NatFeature
|
||||
cBehavior RecommandBehavior
|
||||
|
||||
notifyCh chan struct{}
|
||||
}
|
||||
|
||||
func (s *Session) genAnalysisKey() {
|
||||
hash := md5.New()
|
||||
vIPs := lo.Uniq(parseIPs(s.visitorMsg.MappedAddrs))
|
||||
if len(vIPs) > 0 {
|
||||
hash.Write([]byte(vIPs[0]))
|
||||
}
|
||||
hash.Write([]byte(s.vNatFeature.NatType))
|
||||
hash.Write([]byte(s.vNatFeature.Behavior))
|
||||
hash.Write([]byte(strconv.FormatBool(s.vNatFeature.RegularPortsChange)))
|
||||
|
||||
cIPs := lo.Uniq(parseIPs(s.clientMsg.MappedAddrs))
|
||||
if len(cIPs) > 0 {
|
||||
hash.Write([]byte(cIPs[0]))
|
||||
}
|
||||
hash.Write([]byte(s.cNatFeature.NatType))
|
||||
hash.Write([]byte(s.cNatFeature.Behavior))
|
||||
hash.Write([]byte(strconv.FormatBool(s.cNatFeature.RegularPortsChange)))
|
||||
s.analysisKey = hex.EncodeToString(hash.Sum(nil))
|
||||
}
|
||||
|
||||
type Controller struct {
|
||||
clientCfgs map[string]*ClientCfg
|
||||
sessions map[string]*Session
|
||||
analyzer *Analyzer
|
||||
|
||||
mu sync.RWMutex
|
||||
}
|
||||
|
||||
func NewController(analysisDataReserveDuration time.Duration) (*Controller, error) {
|
||||
return &Controller{
|
||||
clientCfgs: make(map[string]*ClientCfg),
|
||||
sessions: make(map[string]*Session),
|
||||
analyzer: NewAnalyzer(analysisDataReserveDuration),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (c *Controller) CleanWorker(ctx context.Context) {
|
||||
ticker := time.NewTicker(time.Hour)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-ticker.C:
|
||||
start := time.Now()
|
||||
count, total := c.analyzer.Clean()
|
||||
log.Trace("clean %d/%d nathole analysis data, cost %v", count, total, time.Since(start))
|
||||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Controller) ListenClient(name string, sk string) chan string {
|
||||
cfg := &ClientCfg{
|
||||
name: name,
|
||||
sk: sk,
|
||||
sidCh: make(chan string),
|
||||
}
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
c.clientCfgs[name] = cfg
|
||||
return cfg.sidCh
|
||||
}
|
||||
|
||||
func (c *Controller) CloseClient(name string) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
delete(c.clientCfgs, name)
|
||||
}
|
||||
|
||||
func (c *Controller) GenSid() string {
|
||||
t := time.Now().Unix()
|
||||
id, _ := util.RandID()
|
||||
return fmt.Sprintf("%d%s", t, id)
|
||||
}
|
||||
|
||||
func (c *Controller) HandleVisitor(m *msg.NatHoleVisitor, transporter transport.MessageTransporter) {
|
||||
if m.PreCheck {
|
||||
_, ok := c.clientCfgs[m.ProxyName]
|
||||
if !ok {
|
||||
_ = transporter.Send(c.GenNatHoleResponse(m.TransactionID, nil, fmt.Sprintf("xtcp server for [%s] doesn't exist", m.ProxyName)))
|
||||
} else {
|
||||
_ = transporter.Send(c.GenNatHoleResponse(m.TransactionID, nil, ""))
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
sid := c.GenSid()
|
||||
session := &Session{
|
||||
sid: sid,
|
||||
visitorMsg: m,
|
||||
visitorTransporter: transporter,
|
||||
notifyCh: make(chan struct{}, 1),
|
||||
}
|
||||
var (
|
||||
clientCfg *ClientCfg
|
||||
ok bool
|
||||
)
|
||||
err := func() error {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
|
||||
clientCfg, ok = c.clientCfgs[m.ProxyName]
|
||||
if !ok {
|
||||
return fmt.Errorf("xtcp server for [%s] doesn't exist", m.ProxyName)
|
||||
}
|
||||
if m.SignKey != util.GetAuthKey(clientCfg.sk, m.Timestamp) {
|
||||
return fmt.Errorf("xtcp connection of [%s] auth failed", m.ProxyName)
|
||||
}
|
||||
c.sessions[sid] = session
|
||||
return nil
|
||||
}()
|
||||
if err != nil {
|
||||
log.Warn("handle visitorMsg error: %v", err)
|
||||
_ = transporter.Send(c.GenNatHoleResponse(m.TransactionID, nil, err.Error()))
|
||||
return
|
||||
}
|
||||
log.Trace("handle visitor message, sid [%s]", sid)
|
||||
|
||||
defer func() {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
delete(c.sessions, sid)
|
||||
}()
|
||||
|
||||
if err := errors.PanicToError(func() {
|
||||
clientCfg.sidCh <- sid
|
||||
}); err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
// wait for NatHoleClient message
|
||||
select {
|
||||
case <-session.notifyCh:
|
||||
case <-time.After(time.Duration(NatHoleTimeout) * time.Second):
|
||||
log.Debug("wait for NatHoleClient message timeout, sid [%s]", sid)
|
||||
return
|
||||
}
|
||||
|
||||
// Make hole-punching decisions based on the NAT information of the client and visitor.
|
||||
vResp, cResp, err := c.analysis(session)
|
||||
if err != nil {
|
||||
log.Debug("sid [%s] analysis error: %v", err)
|
||||
vResp = c.GenNatHoleResponse(session.visitorMsg.TransactionID, nil, err.Error())
|
||||
cResp = c.GenNatHoleResponse(session.clientMsg.TransactionID, nil, err.Error())
|
||||
}
|
||||
session.cResp = cResp
|
||||
session.vResp = vResp
|
||||
|
||||
// send response to visitor and client
|
||||
var g errgroup.Group
|
||||
g.Go(func() error {
|
||||
// if it's sender, wait for a while to make sure the client has send the detect messages
|
||||
if vResp.DetectBehavior.Role == "sender" {
|
||||
time.Sleep(1 * time.Second)
|
||||
}
|
||||
_ = session.visitorTransporter.Send(vResp)
|
||||
return nil
|
||||
})
|
||||
g.Go(func() error {
|
||||
// if it's sender, wait for a while to make sure the client has send the detect messages
|
||||
if cResp.DetectBehavior.Role == "sender" {
|
||||
time.Sleep(1 * time.Second)
|
||||
}
|
||||
_ = session.clientTransporter.Send(cResp)
|
||||
return nil
|
||||
})
|
||||
_ = g.Wait()
|
||||
|
||||
time.Sleep(time.Duration(cResp.DetectBehavior.ReadTimeoutMs+30000) * time.Millisecond)
|
||||
}
|
||||
|
||||
func (c *Controller) HandleClient(m *msg.NatHoleClient, transporter transport.MessageTransporter) {
|
||||
c.mu.RLock()
|
||||
session, ok := c.sessions[m.Sid]
|
||||
c.mu.RUnlock()
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
log.Trace("handle client message, sid [%s]", session.sid)
|
||||
session.clientMsg = m
|
||||
session.clientTransporter = transporter
|
||||
select {
|
||||
case session.notifyCh <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Controller) HandleReport(m *msg.NatHoleReport) {
|
||||
c.mu.RLock()
|
||||
session, ok := c.sessions[m.Sid]
|
||||
c.mu.RUnlock()
|
||||
if !ok {
|
||||
log.Trace("sid [%s] report make hole success: %v, but session not found", m.Sid, m.Success)
|
||||
return
|
||||
}
|
||||
if m.Success {
|
||||
c.analyzer.ReportSuccess(session.analysisKey, session.recommandMode, session.recommandIndex)
|
||||
}
|
||||
log.Info("sid [%s] report make hole success: %v, mode %v, index %v",
|
||||
m.Sid, m.Success, session.recommandMode, session.recommandIndex)
|
||||
}
|
||||
|
||||
func (c *Controller) GenNatHoleResponse(transactionID string, session *Session, errInfo string) *msg.NatHoleResp {
|
||||
var sid string
|
||||
if session != nil {
|
||||
sid = session.sid
|
||||
}
|
||||
return &msg.NatHoleResp{
|
||||
TransactionID: transactionID,
|
||||
Sid: sid,
|
||||
Error: errInfo,
|
||||
}
|
||||
}
|
||||
|
||||
// analysis analyzes the NAT type and behavior of the visitor and client, then makes hole-punching decisions.
|
||||
// return the response to the visitor and client.
|
||||
func (c *Controller) analysis(session *Session) (*msg.NatHoleResp, *msg.NatHoleResp, error) {
|
||||
cm := session.clientMsg
|
||||
vm := session.visitorMsg
|
||||
|
||||
cNatFeature, err := ClassifyNATFeature(cm.MappedAddrs, parseIPs(cm.AssistedAddrs))
|
||||
if err != nil {
|
||||
return nil, nil, fmt.Errorf("classify client nat feature error: %v", err)
|
||||
}
|
||||
|
||||
vNatFeature, err := ClassifyNATFeature(vm.MappedAddrs, parseIPs(vm.AssistedAddrs))
|
||||
if err != nil {
|
||||
return nil, nil, fmt.Errorf("classify visitor nat feature error: %v", err)
|
||||
}
|
||||
session.cNatFeature = cNatFeature
|
||||
session.vNatFeature = vNatFeature
|
||||
session.genAnalysisKey()
|
||||
|
||||
mode, index, cBehavior, vBehavior := c.analyzer.GetRecommandBehaviors(session.analysisKey, cNatFeature, vNatFeature)
|
||||
session.recommandMode = mode
|
||||
session.recommandIndex = index
|
||||
session.cBehavior = cBehavior
|
||||
session.vBehavior = vBehavior
|
||||
|
||||
timeoutMs := lo.Max([]int{cBehavior.SendDelayMs, vBehavior.SendDelayMs}) + 5000
|
||||
if cBehavior.ListenRandomPorts > 0 || vBehavior.ListenRandomPorts > 0 {
|
||||
timeoutMs += 30000
|
||||
}
|
||||
|
||||
protocol := vm.Protocol
|
||||
vResp := &msg.NatHoleResp{
|
||||
TransactionID: vm.TransactionID,
|
||||
Sid: session.sid,
|
||||
Protocol: protocol,
|
||||
CandidateAddrs: lo.Uniq(cm.MappedAddrs),
|
||||
AssistedAddrs: lo.Uniq(cm.AssistedAddrs),
|
||||
DetectBehavior: msg.NatHoleDetectBehavior{
|
||||
Mode: mode,
|
||||
Role: vBehavior.Role,
|
||||
TTL: vBehavior.TTL,
|
||||
SendDelayMs: vBehavior.SendDelayMs,
|
||||
ReadTimeoutMs: timeoutMs - vBehavior.SendDelayMs,
|
||||
SendRandomPorts: vBehavior.PortsRandomNumber,
|
||||
ListenRandomPorts: vBehavior.ListenRandomPorts,
|
||||
CandidatePorts: getRangePorts(cm.MappedAddrs, cNatFeature.PortsDifference, vBehavior.PortsRangeNumber),
|
||||
},
|
||||
}
|
||||
cResp := &msg.NatHoleResp{
|
||||
TransactionID: cm.TransactionID,
|
||||
Sid: session.sid,
|
||||
Protocol: protocol,
|
||||
CandidateAddrs: lo.Uniq(vm.MappedAddrs),
|
||||
AssistedAddrs: lo.Uniq(vm.AssistedAddrs),
|
||||
DetectBehavior: msg.NatHoleDetectBehavior{
|
||||
Mode: mode,
|
||||
Role: cBehavior.Role,
|
||||
TTL: cBehavior.TTL,
|
||||
SendDelayMs: cBehavior.SendDelayMs,
|
||||
ReadTimeoutMs: timeoutMs - cBehavior.SendDelayMs,
|
||||
SendRandomPorts: cBehavior.PortsRandomNumber,
|
||||
ListenRandomPorts: cBehavior.ListenRandomPorts,
|
||||
CandidatePorts: getRangePorts(vm.MappedAddrs, vNatFeature.PortsDifference, cBehavior.PortsRangeNumber),
|
||||
},
|
||||
}
|
||||
|
||||
log.Debug("sid [%s] visitor nat: %+v, candidateAddrs: %v; client nat: %+v, candidateAddrs: %v, protocol: %s",
|
||||
session.sid, *vNatFeature, vm.MappedAddrs, *cNatFeature, cm.MappedAddrs, protocol)
|
||||
log.Debug("sid [%s] visitor detect behavior: %+v", session.sid, vResp.DetectBehavior)
|
||||
log.Debug("sid [%s] client detect behavior: %+v", session.sid, cResp.DetectBehavior)
|
||||
return vResp, cResp, nil
|
||||
}
|
||||
|
||||
func getRangePorts(addrs []string, difference, maxNumber int) []msg.PortsRange {
|
||||
if maxNumber <= 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
addr, err := lo.Last(addrs)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
var ports []msg.PortsRange
|
||||
_, portStr, err := net.SplitHostPort(addr)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
port, err := strconv.Atoi(portStr)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
ports = append(ports, msg.PortsRange{
|
||||
From: lo.Max([]int{port - difference - 5, port - maxNumber, 1}),
|
||||
To: lo.Min([]int{port + difference + 5, port + maxNumber, 65535}),
|
||||
})
|
||||
return ports
|
||||
}
|
||||
@ -0,0 +1,119 @@
|
||||
// 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 transport
|
||||
|
||||
import (
|
||||
"context"
|
||||
"reflect"
|
||||
"sync"
|
||||
|
||||
"github.com/fatedier/golib/errors"
|
||||
|
||||
"github.com/fatedier/frp/pkg/msg"
|
||||
)
|
||||
|
||||
type MessageTransporter interface {
|
||||
Send(msg.Message) error
|
||||
// Recv(ctx context.Context, laneKey string, msgType string) (Message, error)
|
||||
// Do will first send msg, then recv msg with the same laneKey and specified msgType.
|
||||
Do(ctx context.Context, req msg.Message, laneKey, recvMsgType string) (msg.Message, error)
|
||||
Dispatch(m msg.Message, laneKey string) bool
|
||||
DispatchWithType(m msg.Message, msgType, laneKey string) bool
|
||||
}
|
||||
|
||||
func NewMessageTransporter(sendCh chan msg.Message) MessageTransporter {
|
||||
return &transporterImpl{
|
||||
sendCh: sendCh,
|
||||
registry: make(map[string]map[string]chan msg.Message),
|
||||
}
|
||||
}
|
||||
|
||||
type transporterImpl struct {
|
||||
sendCh chan msg.Message
|
||||
|
||||
// First key is message type and second key is lane key.
|
||||
// Dispatch will dispatch message to releated channel by its message type
|
||||
// and lane key.
|
||||
registry map[string]map[string]chan msg.Message
|
||||
mu sync.RWMutex
|
||||
}
|
||||
|
||||
func (impl *transporterImpl) Send(m msg.Message) error {
|
||||
return errors.PanicToError(func() {
|
||||
impl.sendCh <- m
|
||||
})
|
||||
}
|
||||
|
||||
func (impl *transporterImpl) Do(ctx context.Context, req msg.Message, laneKey, recvMsgType string) (msg.Message, error) {
|
||||
ch := make(chan msg.Message, 1)
|
||||
defer close(ch)
|
||||
unregisterFn := impl.registerMsgChan(ch, laneKey, recvMsgType)
|
||||
defer unregisterFn()
|
||||
|
||||
if err := impl.Send(req); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
case resp := <-ch:
|
||||
return resp, nil
|
||||
}
|
||||
}
|
||||
|
||||
func (impl *transporterImpl) DispatchWithType(m msg.Message, msgType, laneKey string) bool {
|
||||
var ch chan msg.Message
|
||||
impl.mu.RLock()
|
||||
byLaneKey, ok := impl.registry[msgType]
|
||||
if ok {
|
||||
ch = byLaneKey[laneKey]
|
||||
}
|
||||
impl.mu.RUnlock()
|
||||
|
||||
if ch == nil {
|
||||
return false
|
||||
}
|
||||
|
||||
if err := errors.PanicToError(func() {
|
||||
ch <- m
|
||||
}); err != nil {
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func (impl *transporterImpl) Dispatch(m msg.Message, laneKey string) bool {
|
||||
msgType := reflect.TypeOf(m).Elem().Name()
|
||||
return impl.DispatchWithType(m, msgType, laneKey)
|
||||
}
|
||||
|
||||
func (impl *transporterImpl) registerMsgChan(recvCh chan msg.Message, laneKey string, msgType string) (unregister func()) {
|
||||
impl.mu.Lock()
|
||||
byLaneKey, ok := impl.registry[msgType]
|
||||
if !ok {
|
||||
byLaneKey = make(map[string]chan msg.Message)
|
||||
impl.registry[msgType] = byLaneKey
|
||||
}
|
||||
byLaneKey[laneKey] = recvCh
|
||||
impl.mu.Unlock()
|
||||
|
||||
unregister = func() {
|
||||
impl.mu.Lock()
|
||||
delete(byLaneKey, laneKey)
|
||||
impl.mu.Unlock()
|
||||
}
|
||||
return
|
||||
}
|
||||
@ -1,25 +0,0 @@
|
||||
package util
|
||||
|
||||
func InSlice[T comparable](v T, s []T) bool {
|
||||
for _, vv := range s {
|
||||
if v == vv {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func InSliceAny[T any](v T, s []T, equalFn func(a, b T) bool) bool {
|
||||
for _, vv := range s {
|
||||
if equalFn(v, vv) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func InSliceAnyFunc[T any](equalFn func(a, b T) bool) func(v T, s []T) bool {
|
||||
return func(v T, s []T) bool {
|
||||
return InSliceAny(v, s, equalFn)
|
||||
}
|
||||
}
|
||||
@ -1,49 +0,0 @@
|
||||
package util
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestInSlice(t *testing.T) {
|
||||
require := require.New(t)
|
||||
require.True(InSlice(1, []int{1, 2, 3}))
|
||||
require.False(InSlice(0, []int{1, 2, 3}))
|
||||
require.True(InSlice("foo", []string{"foo", "bar"}))
|
||||
require.False(InSlice("not exist", []string{"foo", "bar"}))
|
||||
}
|
||||
|
||||
type testStructA struct {
|
||||
Name string
|
||||
Age int
|
||||
}
|
||||
|
||||
func TestInSliceAny(t *testing.T) {
|
||||
require := require.New(t)
|
||||
|
||||
a := testStructA{Name: "foo", Age: 20}
|
||||
b := testStructA{Name: "foo", Age: 30}
|
||||
c := testStructA{Name: "bar", Age: 20}
|
||||
|
||||
equalFn := func(o, p testStructA) bool {
|
||||
return o.Name == p.Name
|
||||
}
|
||||
require.True(InSliceAny(a, []testStructA{b, c}, equalFn))
|
||||
require.False(InSliceAny(c, []testStructA{a, b}, equalFn))
|
||||
}
|
||||
|
||||
func TestInSliceAnyFunc(t *testing.T) {
|
||||
require := require.New(t)
|
||||
|
||||
a := testStructA{Name: "foo", Age: 20}
|
||||
b := testStructA{Name: "foo", Age: 30}
|
||||
c := testStructA{Name: "bar", Age: 20}
|
||||
|
||||
equalFn := func(o, p testStructA) bool {
|
||||
return o.Name == p.Name
|
||||
}
|
||||
testStructAInSlice := InSliceAnyFunc(equalFn)
|
||||
require.True(testStructAInSlice(a, []testStructA{b, c}))
|
||||
require.False(testStructAInSlice(c, []testStructA{a, b}))
|
||||
}
|
||||
Loading…
Reference in New Issue