v1.0.0

transport/internet/udp/hub.go

@@ -0,0 +1,132 @@
+package udp
+
+import (
+	"context"
+
+	"github.com/xtls/xray-core/v1/common/buf"
+	"github.com/xtls/xray-core/v1/common/net"
+	"github.com/xtls/xray-core/v1/common/protocol/udp"
+	"github.com/xtls/xray-core/v1/transport/internet"
+)
+
+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
+	}
+	newError("listening UDP on ", address, ":", port).WriteToLog()
+	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 {
+			newError("failed to read UDP msg").Base(err).WriteToLog()
+			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() {
+				newError("UDP original destination: ", payload.Target).AtDebug().WriteToLog()
+			} else {
+				newError("failed to read UDP original destination").WriteToLog()
+			}
+		}
+
+		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
+}