2016-06-14 16:54:08 -04:00
|
|
|
package udp
|
2016-01-28 11:43:47 -05:00
|
|
|
|
|
|
|
import (
|
|
|
|
"net"
|
2016-07-13 15:39:18 -04:00
|
|
|
"sync"
|
2016-01-28 11:43:47 -05:00
|
|
|
|
2017-02-10 10:42:24 -05:00
|
|
|
"v2ray.com/core/app/log"
|
2016-12-09 05:35:27 -05:00
|
|
|
"v2ray.com/core/common/buf"
|
2016-11-18 15:30:03 -05:00
|
|
|
"v2ray.com/core/common/dice"
|
2016-08-20 14:55:45 -04:00
|
|
|
v2net "v2ray.com/core/common/net"
|
2016-11-18 15:30:03 -05:00
|
|
|
"v2ray.com/core/common/signal"
|
2016-08-20 14:55:45 -04:00
|
|
|
"v2ray.com/core/transport/internet/internal"
|
2016-01-28 11:43:47 -05:00
|
|
|
)
|
|
|
|
|
2016-12-21 09:48:39 -05:00
|
|
|
// Payload represents a single UDP payload.
|
|
|
|
type Payload struct {
|
2017-01-26 14:46:44 -05:00
|
|
|
payload *buf.Buffer
|
|
|
|
source v2net.Destination
|
|
|
|
originalDest v2net.Destination
|
2016-11-18 15:30:03 -05:00
|
|
|
}
|
|
|
|
|
2016-12-21 09:48:39 -05:00
|
|
|
// PayloadHandler is function to handle Payload.
|
2017-01-26 14:46:44 -05:00
|
|
|
type PayloadHandler func(payload *buf.Buffer, source v2net.Destination, originalDest v2net.Destination)
|
2016-01-28 11:43:47 -05:00
|
|
|
|
2016-12-21 09:48:39 -05:00
|
|
|
// PayloadQueue is a queue of Payload.
|
|
|
|
type PayloadQueue struct {
|
|
|
|
queue []chan Payload
|
|
|
|
callback PayloadHandler
|
2016-11-18 15:30:03 -05:00
|
|
|
}
|
|
|
|
|
2016-12-21 09:48:39 -05:00
|
|
|
// NewPayloadQueue returns a new PayloadQueue.
|
|
|
|
func NewPayloadQueue(option ListenOption) *PayloadQueue {
|
|
|
|
queue := &PayloadQueue{
|
2016-11-18 15:30:03 -05:00
|
|
|
callback: option.Callback,
|
2016-12-21 09:48:39 -05:00
|
|
|
queue: make([]chan Payload, option.Concurrency),
|
2016-11-18 15:30:03 -05:00
|
|
|
}
|
|
|
|
for i := range queue.queue {
|
2016-12-21 09:48:39 -05:00
|
|
|
queue.queue[i] = make(chan Payload, 64)
|
2016-11-18 15:30:03 -05:00
|
|
|
go queue.Dequeue(queue.queue[i])
|
|
|
|
}
|
|
|
|
return queue
|
|
|
|
}
|
|
|
|
|
2016-12-21 09:48:39 -05:00
|
|
|
func (v *PayloadQueue) Enqueue(payload Payload) {
|
2016-11-27 15:39:09 -05:00
|
|
|
size := len(v.queue)
|
2016-11-18 15:34:42 -05:00
|
|
|
idx := 0
|
|
|
|
if size > 1 {
|
|
|
|
idx = dice.Roll(size)
|
|
|
|
}
|
2016-11-18 15:30:03 -05:00
|
|
|
for i := 0; i < size; i++ {
|
|
|
|
select {
|
2016-11-27 15:39:09 -05:00
|
|
|
case v.queue[idx%size] <- payload:
|
2016-11-18 15:30:03 -05:00
|
|
|
return
|
|
|
|
default:
|
2016-11-18 15:34:42 -05:00
|
|
|
idx++
|
2016-11-18 15:30:03 -05:00
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2016-12-21 09:48:39 -05:00
|
|
|
func (v *PayloadQueue) Dequeue(queue <-chan Payload) {
|
2016-11-18 18:34:30 -05:00
|
|
|
for payload := range queue {
|
2017-01-26 14:46:44 -05:00
|
|
|
v.callback(payload.payload, payload.source, payload.originalDest)
|
2016-11-18 15:30:03 -05:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2016-12-21 09:48:39 -05:00
|
|
|
func (v *PayloadQueue) Close() {
|
2016-11-27 15:39:09 -05:00
|
|
|
for _, queue := range v.queue {
|
2016-11-18 15:30:03 -05:00
|
|
|
close(queue)
|
|
|
|
}
|
2016-01-28 11:43:47 -05:00
|
|
|
}
|
|
|
|
|
2016-08-15 11:44:46 -04:00
|
|
|
type ListenOption struct {
|
2016-12-21 09:48:39 -05:00
|
|
|
Callback PayloadHandler
|
2016-08-15 11:44:46 -04:00
|
|
|
ReceiveOriginalDest bool
|
2016-11-18 15:30:03 -05:00
|
|
|
Concurrency int
|
|
|
|
}
|
|
|
|
|
2016-12-21 09:48:39 -05:00
|
|
|
type Hub struct {
|
2016-11-18 15:30:03 -05:00
|
|
|
sync.RWMutex
|
|
|
|
conn *net.UDPConn
|
|
|
|
cancel *signal.CancelSignal
|
2016-12-21 09:48:39 -05:00
|
|
|
queue *PayloadQueue
|
2016-11-18 15:30:03 -05:00
|
|
|
option ListenOption
|
2016-08-15 11:44:46 -04:00
|
|
|
}
|
|
|
|
|
2016-12-21 09:48:39 -05:00
|
|
|
func ListenUDP(address v2net.Address, port v2net.Port, option ListenOption) (*Hub, error) {
|
2016-11-18 15:30:03 -05:00
|
|
|
if option.Concurrency < 1 {
|
|
|
|
option.Concurrency = 1
|
|
|
|
}
|
2016-01-28 11:43:47 -05:00
|
|
|
udpConn, err := net.ListenUDP("udp", &net.UDPAddr{
|
2016-05-29 10:37:52 -04:00
|
|
|
IP: address.IP(),
|
2016-01-28 11:43:47 -05:00
|
|
|
Port: int(port),
|
|
|
|
})
|
|
|
|
if err != nil {
|
|
|
|
return nil, err
|
|
|
|
}
|
2017-01-26 14:46:44 -05:00
|
|
|
log.Info("UDP|Hub: Listening on ", address, ":", port)
|
2016-08-15 11:44:46 -04:00
|
|
|
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
|
|
|
|
}
|
|
|
|
}
|
2016-12-21 09:48:39 -05:00
|
|
|
hub := &Hub{
|
2016-08-15 11:44:46 -04:00
|
|
|
conn: udpConn,
|
2016-12-21 09:48:39 -05:00
|
|
|
queue: NewPayloadQueue(option),
|
2016-11-18 15:30:03 -05:00
|
|
|
option: option,
|
|
|
|
cancel: signal.NewCloseSignal(),
|
2016-01-28 11:43:47 -05:00
|
|
|
}
|
|
|
|
go hub.start()
|
|
|
|
return hub, nil
|
|
|
|
}
|
|
|
|
|
2016-12-21 09:48:39 -05:00
|
|
|
func (v *Hub) Close() {
|
2016-11-27 15:39:09 -05:00
|
|
|
v.Lock()
|
|
|
|
defer v.Unlock()
|
2016-07-13 15:39:18 -04:00
|
|
|
|
2016-11-27 15:39:09 -05:00
|
|
|
v.cancel.Cancel()
|
|
|
|
v.conn.Close()
|
|
|
|
v.cancel.WaitForDone()
|
|
|
|
v.queue.Close()
|
2016-01-28 11:43:47 -05:00
|
|
|
}
|
|
|
|
|
2016-12-21 09:48:39 -05:00
|
|
|
func (v *Hub) WriteTo(payload []byte, dest v2net.Destination) (int, error) {
|
2016-11-27 15:39:09 -05:00
|
|
|
return v.conn.WriteToUDP(payload, &net.UDPAddr{
|
2016-09-20 05:53:05 -04:00
|
|
|
IP: dest.Address.IP(),
|
|
|
|
Port: int(dest.Port),
|
2016-01-28 11:43:47 -05:00
|
|
|
})
|
|
|
|
}
|
|
|
|
|
2016-12-21 09:48:39 -05:00
|
|
|
func (v *Hub) start() {
|
2016-11-27 15:39:09 -05:00
|
|
|
v.cancel.WaitThread()
|
|
|
|
defer v.cancel.FinishThread()
|
2016-07-13 15:39:18 -04:00
|
|
|
|
2016-08-15 11:44:46 -04:00
|
|
|
oobBytes := make([]byte, 256)
|
2016-11-27 15:39:09 -05:00
|
|
|
for v.Running() {
|
2016-12-09 06:08:25 -05:00
|
|
|
buffer := buf.NewSmall()
|
2016-12-06 05:03:42 -05:00
|
|
|
var noob int
|
|
|
|
var addr *net.UDPAddr
|
2016-12-09 06:08:25 -05:00
|
|
|
err := buffer.AppendSupplier(func(b []byte) (int, error) {
|
2016-12-06 05:03:42 -05:00
|
|
|
n, nb, _, a, e := ReadUDPMsg(v.conn, b, oobBytes)
|
|
|
|
noob = nb
|
|
|
|
addr = a
|
2016-12-09 06:08:25 -05:00
|
|
|
return n, e
|
2016-12-06 05:03:42 -05:00
|
|
|
})
|
|
|
|
|
2016-01-28 11:43:47 -05:00
|
|
|
if err != nil {
|
2016-08-19 06:58:26 -04:00
|
|
|
log.Info("UDP|Hub: Failed to read UDP msg: ", err)
|
2016-01-28 11:43:47 -05:00
|
|
|
buffer.Release()
|
|
|
|
continue
|
|
|
|
}
|
2016-08-15 11:44:46 -04:00
|
|
|
|
2017-01-26 14:46:44 -05:00
|
|
|
payload := Payload{
|
|
|
|
payload: buffer,
|
|
|
|
}
|
|
|
|
payload.source = v2net.UDPDestination(v2net.IPAddress(addr.IP), v2net.Port(addr.Port))
|
2016-11-27 15:39:09 -05:00
|
|
|
if v.option.ReceiveOriginalDest && noob > 0 {
|
2017-01-26 14:46:44 -05:00
|
|
|
payload.originalDest = RetrieveOriginalDest(oobBytes[:noob])
|
2016-08-15 11:44:46 -04:00
|
|
|
}
|
2017-01-26 14:46:44 -05:00
|
|
|
v.queue.Enqueue(payload)
|
2016-01-28 11:43:47 -05:00
|
|
|
}
|
|
|
|
}
|
2016-07-13 15:39:18 -04:00
|
|
|
|
2016-12-21 09:48:39 -05:00
|
|
|
func (v *Hub) Running() bool {
|
2016-11-27 15:39:09 -05:00
|
|
|
return !v.cancel.Cancelled()
|
2016-07-13 15:39:18 -04:00
|
|
|
}
|
2016-08-15 15:33:21 -04:00
|
|
|
|
2016-12-21 09:48:39 -05:00
|
|
|
// Connection return the net.Conn underneath this hub.
|
2016-08-15 15:33:21 -04:00
|
|
|
// Private: Visible for testing only
|
2016-12-21 09:48:39 -05:00
|
|
|
func (v *Hub) Connection() net.Conn {
|
2016-11-27 15:39:09 -05:00
|
|
|
return v.conn
|
2016-08-15 15:33:21 -04:00
|
|
|
}
|
2016-08-15 16:20:16 -04:00
|
|
|
|
2016-12-21 09:48:39 -05:00
|
|
|
func (v *Hub) Addr() net.Addr {
|
2016-11-27 15:39:09 -05:00
|
|
|
return v.conn.LocalAddr()
|
2016-08-15 16:20:16 -04:00
|
|
|
}
|