v2ray-core/transport/internet/udp/udp.go

package udp

import (
	"net"
	"sync"

	"github.com/v2ray/v2ray-core/common/alloc"
	v2net "github.com/v2ray/v2ray-core/common/net"
)

type UDPPayloadHandler func(*alloc.Buffer, v2net.Destination)

type UDPHub struct {
	sync.RWMutex
	conn      *net.UDPConn
	callback  UDPPayloadHandler
	accepting bool
}

func ListenUDP(address v2net.Address, port v2net.Port, callback UDPPayloadHandler) (*UDPHub, error) {
	udpConn, err := net.ListenUDP("udp", &net.UDPAddr{
		IP:   address.IP(),
		Port: int(port),
	})
	if err != nil {
		return nil, err
	}
	hub := &UDPHub{
		conn:     udpConn,
		callback: callback,
	}
	go hub.start()
	return hub, nil
}

func (this *UDPHub) Close() {
	this.Lock()
	defer this.Unlock()

	this.accepting = false
	this.conn.Close()
}

func (this *UDPHub) WriteTo(payload []byte, dest v2net.Destination) (int, error) {
	return this.conn.WriteToUDP(payload, &net.UDPAddr{
		IP:   dest.Address().IP(),
		Port: int(dest.Port()),
	})
}

func (this *UDPHub) start() {
	this.Lock()
	this.accepting = true
	this.Unlock()

	for this.Running() {
		buffer := alloc.NewBuffer()
		nBytes, addr, err := this.conn.ReadFromUDP(buffer.Value)
		if err != nil {
			buffer.Release()
			continue
		}
		buffer.Slice(0, nBytes)
		dest := v2net.UDPDestination(v2net.IPAddress(addr.IP), v2net.Port(addr.Port))
		go this.callback(buffer, dest)
	}
}

func (this *UDPHub) Running() bool {
	this.RLock()
	defer this.RUnlock()

	return this.accepting
}