forked from v2fly/v2ray-core
-
Notifications
You must be signed in to change notification settings - Fork 0
/
udp.go
74 lines (62 loc) · 1.4 KB
/
udp.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
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
}