1
0
mirror of https://github.com/v2fly/v2ray-core.git synced 2024-12-22 18:17:52 -05:00
v2fly/transport/internet/udp/hub.go
2016-11-01 11:46:34 +01:00

114 lines
2.4 KiB
Go

package udp
import (
"net"
"sync"
"v2ray.com/core/common/alloc"
"v2ray.com/core/common/log"
v2net "v2ray.com/core/common/net"
"v2ray.com/core/proxy"
"v2ray.com/core/transport/internet/internal"
)
type UDPPayloadHandler func(*alloc.Buffer, *proxy.SessionInfo)
type UDPHub struct {
sync.RWMutex
conn *net.UDPConn
option ListenOption
accepting bool
pool *alloc.BufferPool
}
type ListenOption struct {
Callback UDPPayloadHandler
ReceiveOriginalDest bool
}
func ListenUDP(address v2net.Address, port v2net.Port, option ListenOption) (*UDPHub, error) {
udpConn, err := net.ListenUDP("udp", &net.UDPAddr{
IP: address.IP(),
Port: int(port),
})
if err != nil {
return nil, err
}
if option.ReceiveOriginalDest {
fd, err := internal.GetSysFd(udpConn)
if err != nil {
log.Warning("UDP|Listener: Failed to get fd: ", err)
return nil, err
}
err = SetOriginalDestOptions(fd)
if err != nil {
log.Warning("UDP|Listener: Failed to set socket options: ", err)
return nil, err
}
}
hub := &UDPHub{
conn: udpConn,
option: option,
pool: alloc.NewBufferPool(2048, 64),
}
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()
oobBytes := make([]byte, 256)
for this.Running() {
buffer := this.pool.Allocate()
nBytes, noob, _, addr, err := ReadUDPMsg(this.conn, buffer.Value, oobBytes)
if err != nil {
log.Info("UDP|Hub: Failed to read UDP msg: ", err)
buffer.Release()
continue
}
buffer.Slice(0, nBytes)
session := new(proxy.SessionInfo)
session.Source = v2net.UDPDestination(v2net.IPAddress(addr.IP), v2net.Port(addr.Port))
if this.option.ReceiveOriginalDest && noob > 0 {
session.Destination = RetrieveOriginalDest(oobBytes[:noob])
}
go this.option.Callback(buffer, session)
}
}
func (this *UDPHub) Running() bool {
this.RLock()
defer this.RUnlock()
return this.accepting
}
// Connection return the net.Conn underneath this hub.
// Private: Visible for testing only
func (this *UDPHub) Connection() net.Conn {
return this.conn
}
func (this *UDPHub) Addr() net.Addr {
return this.conn.LocalAddr()
}