github.com/metacubex/mihomo@v1.18.5/listener/shadowsocks/udp.go (about)

     1  package shadowsocks
     2  
     3  import (
     4  	"net"
     5  
     6  	"github.com/metacubex/mihomo/adapter/inbound"
     7  	N "github.com/metacubex/mihomo/common/net"
     8  	"github.com/metacubex/mihomo/common/sockopt"
     9  	C "github.com/metacubex/mihomo/constant"
    10  	"github.com/metacubex/mihomo/log"
    11  	"github.com/metacubex/mihomo/transport/shadowsocks/core"
    12  	"github.com/metacubex/mihomo/transport/socks5"
    13  )
    14  
    15  type UDPListener struct {
    16  	packetConn net.PacketConn
    17  	closed     bool
    18  }
    19  
    20  func NewUDP(addr string, pickCipher core.Cipher, tunnel C.Tunnel) (*UDPListener, error) {
    21  	l, err := net.ListenPacket("udp", addr)
    22  	if err != nil {
    23  		return nil, err
    24  	}
    25  
    26  	err = sockopt.UDPReuseaddr(l.(*net.UDPConn))
    27  	if err != nil {
    28  		log.Warnln("Failed to Reuse UDP Address: %s", err)
    29  	}
    30  
    31  	sl := &UDPListener{l, false}
    32  	conn := pickCipher.PacketConn(N.NewEnhancePacketConn(l))
    33  	go func() {
    34  		for {
    35  			data, put, remoteAddr, err := conn.WaitReadFrom()
    36  			if err != nil {
    37  				if put != nil {
    38  					put()
    39  				}
    40  				if sl.closed {
    41  					break
    42  				}
    43  				continue
    44  			}
    45  			handleSocksUDP(conn, tunnel, data, put, remoteAddr)
    46  		}
    47  	}()
    48  
    49  	return sl, nil
    50  }
    51  
    52  func (l *UDPListener) Close() error {
    53  	l.closed = true
    54  	return l.packetConn.Close()
    55  }
    56  
    57  func (l *UDPListener) LocalAddr() net.Addr {
    58  	return l.packetConn.LocalAddr()
    59  }
    60  
    61  func handleSocksUDP(pc net.PacketConn, tunnel C.Tunnel, buf []byte, put func(), addr net.Addr, additions ...inbound.Addition) {
    62  	tgtAddr := socks5.SplitAddr(buf)
    63  	if tgtAddr == nil {
    64  		// Unresolved UDP packet, return buffer to the pool
    65  		if put != nil {
    66  			put()
    67  		}
    68  		return
    69  	}
    70  	target := tgtAddr
    71  	payload := buf[len(tgtAddr):]
    72  
    73  	packet := &packet{
    74  		pc:      pc,
    75  		rAddr:   addr,
    76  		payload: payload,
    77  		put:     put,
    78  	}
    79  	tunnel.HandleUDPPacket(inbound.NewPacket(target, packet, C.SHADOWSOCKS, additions...))
    80  }