Xray-core/transport/internet/udp/hub.go

134 lines
2.8 KiB
Go
Raw Normal View History

2020-11-25 11:01:53 +00:00
package udp
import (
"context"
2020-12-04 01:36:16 +00:00
"github.com/xtls/xray-core/common/buf"
"github.com/xtls/xray-core/common/errors"
2020-12-04 01:36:16 +00:00
"github.com/xtls/xray-core/common/net"
"github.com/xtls/xray-core/common/protocol/udp"
"github.com/xtls/xray-core/transport/internet"
2020-11-25 11:01:53 +00:00
)
type HubOption func(h *Hub)
func HubCapacity(capacity int) HubOption {
return func(h *Hub) {
h.capacity = capacity
}
}
func HubReceiveOriginalDestination(r bool) HubOption {
return func(h *Hub) {
h.recvOrigDest = r
}
}
type Hub struct {
conn *net.UDPConn
cache chan *udp.Packet
capacity int
recvOrigDest bool
}
func ListenUDP(ctx context.Context, address net.Address, port net.Port, streamSettings *internet.MemoryStreamConfig, options ...HubOption) (*Hub, error) {
hub := &Hub{
capacity: 256,
recvOrigDest: false,
}
for _, opt := range options {
opt(hub)
}
var sockopt *internet.SocketConfig
if streamSettings != nil {
sockopt = streamSettings.SocketSettings
}
if sockopt != nil && sockopt.ReceiveOriginalDestAddress {
hub.recvOrigDest = true
}
udpConn, err := internet.ListenSystemPacket(ctx, &net.UDPAddr{
IP: address.IP(),
Port: int(port),
}, sockopt)
if err != nil {
return nil, err
}
errors.LogInfo(ctx, "listening UDP on ", address, ":", port)
2020-11-25 11:01:53 +00:00
hub.conn = udpConn.(*net.UDPConn)
hub.cache = make(chan *udp.Packet, hub.capacity)
go hub.start()
return hub, nil
}
// Close implements net.Listener.
func (h *Hub) Close() error {
h.conn.Close()
return nil
}
func (h *Hub) WriteTo(payload []byte, dest net.Destination) (int, error) {
return h.conn.WriteToUDP(payload, &net.UDPAddr{
IP: dest.Address.IP(),
Port: int(dest.Port),
})
}
func (h *Hub) start() {
c := h.cache
defer close(c)
oobBytes := make([]byte, 256)
for {
buffer := buf.New()
var noob int
var addr *net.UDPAddr
rawBytes := buffer.Extend(buf.Size)
n, noob, _, addr, err := ReadUDPMsg(h.conn, rawBytes, oobBytes)
if err != nil {
errors.LogInfoInner(context.Background(), err, "failed to read UDP msg")
2020-11-25 11:01:53 +00:00
buffer.Release()
break
}
buffer.Resize(0, int32(n))
if buffer.IsEmpty() {
buffer.Release()
continue
}
payload := &udp.Packet{
Payload: buffer,
Source: net.UDPDestination(net.IPAddress(addr.IP), net.Port(addr.Port)),
}
if h.recvOrigDest && noob > 0 {
payload.Target = RetrieveOriginalDest(oobBytes[:noob])
if payload.Target.IsValid() {
errors.LogDebug(context.Background(), "UDP original destination: ", payload.Target)
2020-11-25 11:01:53 +00:00
} else {
errors.LogInfo(context.Background(), "failed to read UDP original destination")
2020-11-25 11:01:53 +00:00
}
}
select {
case c <- payload:
default:
buffer.Release()
payload.Payload = nil
}
}
}
// Addr implements net.Listener.
func (h *Hub) Addr() net.Addr {
return h.conn.LocalAddr()
}
func (h *Hub) Receive() <-chan *udp.Packet {
return h.cache
}