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

     1  package socks
     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/socks5"
    12  )
    13  
    14  type UDPListener struct {
    15  	packetConn net.PacketConn
    16  	addr       string
    17  	closed     bool
    18  }
    19  
    20  // RawAddress implements C.Listener
    21  func (l *UDPListener) RawAddress() string {
    22  	return l.addr
    23  }
    24  
    25  // Address implements C.Listener
    26  func (l *UDPListener) Address() string {
    27  	return l.packetConn.LocalAddr().String()
    28  }
    29  
    30  // Close implements C.Listener
    31  func (l *UDPListener) Close() error {
    32  	l.closed = true
    33  	return l.packetConn.Close()
    34  }
    35  
    36  func NewUDP(addr string, tunnel C.Tunnel, additions ...inbound.Addition) (*UDPListener, error) {
    37  	if len(additions) == 0 {
    38  		additions = []inbound.Addition{
    39  			inbound.WithInName("DEFAULT-SOCKS"),
    40  			inbound.WithSpecialRules(""),
    41  		}
    42  	}
    43  	l, err := net.ListenPacket("udp", addr)
    44  	if err != nil {
    45  		return nil, err
    46  	}
    47  
    48  	if err := sockopt.UDPReuseaddr(l.(*net.UDPConn)); err != nil {
    49  		log.Warnln("Failed to Reuse UDP Address: %s", err)
    50  	}
    51  
    52  	sl := &UDPListener{
    53  		packetConn: l,
    54  		addr:       addr,
    55  	}
    56  	conn := N.NewEnhancePacketConn(l)
    57  	go func() {
    58  		for {
    59  			data, put, remoteAddr, err := conn.WaitReadFrom()
    60  			if err != nil {
    61  				if put != nil {
    62  					put()
    63  				}
    64  				if sl.closed {
    65  					break
    66  				}
    67  				continue
    68  			}
    69  			handleSocksUDP(l, tunnel, data, put, remoteAddr, additions...)
    70  		}
    71  	}()
    72  
    73  	return sl, nil
    74  }
    75  
    76  func handleSocksUDP(pc net.PacketConn, tunnel C.Tunnel, buf []byte, put func(), addr net.Addr, additions ...inbound.Addition) {
    77  	target, payload, err := socks5.DecodeUDPPacket(buf)
    78  	if err != nil {
    79  		// Unresolved UDP packet, return buffer to the pool
    80  		if put != nil {
    81  			put()
    82  		}
    83  		return
    84  	}
    85  	packet := &packet{
    86  		pc:      pc,
    87  		rAddr:   addr,
    88  		payload: payload,
    89  		put:     put,
    90  	}
    91  	tunnel.HandleUDPPacket(inbound.NewPacket(target, packet, C.SOCKS5, additions...))
    92  }