123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137 |
- package udp
- import (
- "encoding/base64"
- "net"
- "sync"
- "time"
- "github.com/fatedier/frp/models/msg"
- "github.com/fatedier/golib/errors"
- "github.com/fatedier/golib/pool"
- )
- func NewUdpPacket(buf []byte, laddr, raddr *net.UDPAddr) *msg.UdpPacket {
- return &msg.UdpPacket{
- Content: base64.StdEncoding.EncodeToString(buf),
- LocalAddr: laddr,
- RemoteAddr: raddr,
- }
- }
- func GetContent(m *msg.UdpPacket) (buf []byte, err error) {
- buf, err = base64.StdEncoding.DecodeString(m.Content)
- return
- }
- func ForwardUserConn(udpConn *net.UDPConn, readCh <-chan *msg.UdpPacket, sendCh chan<- *msg.UdpPacket) {
-
- go func() {
- for udpMsg := range readCh {
- buf, err := GetContent(udpMsg)
- if err != nil {
- continue
- }
- udpConn.WriteToUDP(buf, udpMsg.RemoteAddr)
- }
- }()
-
- buf := pool.GetBuf(1500)
- defer pool.PutBuf(buf)
- for {
- n, remoteAddr, err := udpConn.ReadFromUDP(buf)
- if err != nil {
- udpConn.Close()
- return
- }
-
- udpMsg := NewUdpPacket(buf[:n], nil, remoteAddr)
- select {
- case sendCh <- udpMsg:
- default:
- }
- }
- return
- }
- func Forwarder(dstAddr *net.UDPAddr, readCh <-chan *msg.UdpPacket, sendCh chan<- msg.Message) {
- var (
- mu sync.RWMutex
- )
- udpConnMap := make(map[string]*net.UDPConn)
-
- writerFn := func(raddr *net.UDPAddr, udpConn *net.UDPConn) {
- addr := raddr.String()
- defer func() {
- mu.Lock()
- delete(udpConnMap, addr)
- mu.Unlock()
- udpConn.Close()
- }()
- buf := pool.GetBuf(1500)
- for {
- udpConn.SetReadDeadline(time.Now().Add(30 * time.Second))
- n, _, err := udpConn.ReadFromUDP(buf)
- if err != nil {
- return
- }
- udpMsg := NewUdpPacket(buf[:n], nil, raddr)
- if err = errors.PanicToError(func() {
- select {
- case sendCh <- udpMsg:
- default:
- }
- }); err != nil {
- return
- }
- }
- }
-
- go func() {
- for udpMsg := range readCh {
- buf, err := GetContent(udpMsg)
- if err != nil {
- continue
- }
- mu.Lock()
- udpConn, ok := udpConnMap[udpMsg.RemoteAddr.String()]
- if !ok {
- udpConn, err = net.DialUDP("udp", nil, dstAddr)
- if err != nil {
- continue
- }
- udpConnMap[udpMsg.RemoteAddr.String()] = udpConn
- }
- mu.Unlock()
- _, err = udpConn.Write(buf)
- if err != nil {
- udpConn.Close()
- }
- if !ok {
- go writerFn(udpMsg.RemoteAddr, udpConn)
- }
- }
- }()
- }
|