TUN inbound: Preserve UDP packet destinations with traffic stats (#6747)

https://github.com/XTLS/Xray-core/pull/6747#issuecomment-5647964551

---------

Co-authored-by: 风扇滑翔翼 <Fangliding.fshxy@outlook.com>

app/proxyman/outbound/outbound.go

@@ -2,43 +2,33 @@ package outbound
 
 import (
 	"context"
-	"slices"
 	"sort"
 	"strings"
-	"sync/atomic"
+	"sync"
 
 	"github.com/xtls/xray-core/app/proxyman"
 	"github.com/xtls/xray-core/common"
 	"github.com/xtls/xray-core/common/errors"
-	"github.com/xtls/xray-core/common/utils"
 	"github.com/xtls/xray-core/core"
 	"github.com/xtls/xray-core/features/outbound"
-	"github.com/xtls/xray-core/features/routing"
 )
 
 // Manager is to manage all outbound handlers.
 type Manager struct {
-	defaultHandler   atomic.Pointer[outbound.Handler]
-	taggedHandler    *utils.TypedSyncMap[string, outbound.Handler]
-	untaggedHandlers atomic.Pointer[[]outbound.Handler]
-	running          atomic.Bool
-	tagsCache        *utils.TypedSyncMap[string, []string]
-	balancerPicker   routing.BalancerPicker
+	access           sync.RWMutex
+	defaultHandler   outbound.Handler
+	taggedHandler    map[string]outbound.Handler
+	untaggedHandlers []outbound.Handler
+	running          bool
+	tagsCache        *sync.Map
 }
 
 // New creates a new Manager.
 func New(ctx context.Context, config *proxyman.OutboundConfig) (*Manager, error) {
 	m := &Manager{
-		taggedHandler: utils.NewTypedSyncMap[string, outbound.Handler](),
+		taggedHandler: make(map[string]outbound.Handler),
+		tagsCache:     &sync.Map{},
 	}
-	m.tagsCache = utils.NewTypedSyncMap[string, []string]()
-	empty := make([]outbound.Handler, 0)
-	m.untaggedHandlers.Store(&empty)
-	_ = core.OptionalFeatures(ctx, func(router routing.Router) {
-		if picker, ok := router.(routing.BalancerPicker); ok {
-			m.balancerPicker = picker
-		}
-	})
 	return m, nil
 }
 
@@ -49,25 +39,20 @@ func (m *Manager) Type() interface{} {
 
 // Start implements core.Feature
 func (m *Manager) Start() error {
-	m.running.Store(true)
+	m.access.Lock()
+	defer m.access.Unlock()
 
-	var startErr error
-	m.taggedHandler.Range(func(_ string, h outbound.Handler) bool {
+	m.running = true
+
+	for _, h := range m.taggedHandler {
 		if err := h.Start(); err != nil {
-			startErr = err
-			return false
+			return err
 		}
-		return true
-	})
-	if startErr != nil {
-		return startErr
 	}
 
-	if untagged := m.untaggedHandlers.Load(); untagged != nil {
-		for _, h := range *untagged {
-			if err := h.Start(); err != nil {
-				return err
-			}
+	for _, h := range m.untaggedHandlers {
+		if err := h.Start(); err != nil {
+			return err
 		}
 	}
 
@@ -76,18 +61,18 @@ func (m *Manager) Start() error {
 
 // Close implements core.Feature
 func (m *Manager) Close() error {
-	m.running.Store(false)
+	m.access.Lock()
+	defer m.access.Unlock()
+
+	m.running = false
 
 	var errs []error
-	m.taggedHandler.Range(func(_ string, h outbound.Handler) bool {
+	for _, h := range m.taggedHandler {
 		errs = append(errs, h.Close())
-		return true
-	})
+	}
 
-	if untagged := m.untaggedHandlers.Load(); untagged != nil {
-		for _, h := range *untagged {
-			errs = append(errs, h.Close())
-		}
+	for _, h := range m.untaggedHandlers {
+		errs = append(errs, h.Close())
 	}
 
 	return errors.Combine(errs...)
@@ -95,69 +80,47 @@ func (m *Manager) Close() error {
 
 // GetDefaultHandler implements outbound.Manager.
 func (m *Manager) GetDefaultHandler() outbound.Handler {
-	if h := m.defaultHandler.Load(); h != nil {
-		return *h
+	m.access.RLock()
+	defer m.access.RUnlock()
+
+	if m.defaultHandler == nil {
+		return nil
 	}
-	return nil
+	return m.defaultHandler
 }
 
 // GetHandler implements outbound.Manager.
 func (m *Manager) GetHandler(tag string) outbound.Handler {
-	if handler, found := m.taggedHandler.Load(tag); found {
+	m.access.RLock()
+	defer m.access.RUnlock()
+	if handler, found := m.taggedHandler[tag]; found {
 		return handler
 	}
-	if strings.HasPrefix(tag, "balancer:") && m.balancerPicker != nil {
-		targetTag := m.tryGetOutboundTagWithBalancer(tag, nil)
-		if targetTag == "" {
-			return nil
-		}
-		if handler, found := m.taggedHandler.Load(targetTag); found {
-			return handler
-		}
-	}
 	return nil
 }
 
-func (m *Manager) tryGetOutboundTagWithBalancer(tag string, parents []string) string {
-	balancerTag := tag[len("balancer:"):]
-	targetTag, err := m.balancerPicker.GetBalancerOutboundTag(balancerTag)
-	if err != nil {
-		errors.LogWarning(context.Background(), "failed to pick outbound from balancer [", balancerTag, "]: ", err)
-		return ""
-	}
-	if strings.HasPrefix(targetTag, "balancer:") {
-		if slices.Contains(parents, balancerTag) {
-			errors.LogWarning(context.Background(), "detected balancer loop for [", balancerTag, "]")
-			return ""
-		}
-		return m.tryGetOutboundTagWithBalancer(targetTag, append(parents, balancerTag))
-	}
-	return targetTag
-}
-
 // AddHandler implements outbound.Manager.
 func (m *Manager) AddHandler(ctx context.Context, handler outbound.Handler) error {
-	m.defaultHandler.CompareAndSwap(nil, &handler)
+	m.access.Lock()
+	defer m.access.Unlock()
+
+	m.tagsCache = &sync.Map{}
+
+	if m.defaultHandler == nil {
+		m.defaultHandler = handler
+	}
 
 	tag := handler.Tag()
 	if len(tag) > 0 {
-		if _, found := m.taggedHandler.LoadOrStore(tag, handler); found {
+		if _, found := m.taggedHandler[tag]; found {
 			return errors.New("existing tag found: " + tag)
 		}
+		m.taggedHandler[tag] = handler
 	} else {
-		for {
-			oldUntagged := m.untaggedHandlers.Load()
-			newUntagged := make([]outbound.Handler, 0, len(*oldUntagged)+1)
-			newUntagged = append(newUntagged, *oldUntagged...)
-			newUntagged = append(newUntagged, handler)
-			if m.untaggedHandlers.CompareAndSwap(oldUntagged, &newUntagged) {
-				m.tagsCache.Clear()
-				break
-			}
-		}
+		m.untaggedHandlers = append(m.untaggedHandlers, handler)
 	}
 
-	if m.running.Load() {
+	if m.running {
 		return handler.Start()
 	}
 
@@ -169,27 +132,30 @@ func (m *Manager) RemoveHandler(ctx context.Context, tag string) error {
 	if tag == "" {
 		return common.ErrNoClue
 	}
+	m.access.Lock()
+	defer m.access.Unlock()
 
-	m.taggedHandler.Delete(tag)
-	if cur := m.defaultHandler.Load(); cur != nil && (*cur).Tag() == tag {
-		m.defaultHandler.CompareAndSwap(cur, nil)
+	m.tagsCache = &sync.Map{}
+
+	delete(m.taggedHandler, tag)
+	if m.defaultHandler != nil && m.defaultHandler.Tag() == tag {
+		m.defaultHandler = nil
 	}
-	m.tagsCache.Clear()
 
 	return nil
 }
 
 // ListHandlers implements outbound.Manager.
 func (m *Manager) ListHandlers(ctx context.Context) []outbound.Handler {
-	var response []outbound.Handler
-	if untagged := m.untaggedHandlers.Load(); untagged != nil {
-		response = slices.Clone(*untagged)
-	}
+	m.access.RLock()
+	defer m.access.RUnlock()
 
-	m.taggedHandler.Range(func(_ string, v outbound.Handler) bool {
+	response := make([]outbound.Handler, len(m.untaggedHandlers))
+	copy(response, m.untaggedHandlers)
+
+	for _, v := range m.taggedHandler {
 		response = append(response, v)
-		return true
-	})
+	}
 
 	return response
 }
@@ -197,21 +163,23 @@ func (m *Manager) ListHandlers(ctx context.Context) []outbound.Handler {
 // Select implements outbound.HandlerSelector.
 func (m *Manager) Select(selectors []string) []string {
 	key := strings.Join(selectors, ",")
-	if result, ok := m.tagsCache.Load(key); ok {
-		return result
+	if cache, ok := m.tagsCache.Load(key); ok {
+		return cache.([]string)
 	}
 
+	m.access.RLock()
+	defer m.access.RUnlock()
+
 	tags := make([]string, 0, len(selectors))
 
-	m.taggedHandler.Range(func(tag string, _ outbound.Handler) bool {
+	for tag := range m.taggedHandler {
 		for _, selector := range selectors {
 			if strings.HasPrefix(tag, selector) {
 				tags = append(tags, tag)
 				break
 			}
 		}
-		return true
-	})
+	}
 
 	sort.Strings(tags)
 	m.tagsCache.Store(key, tags)

app/router/balancing.go

@@ -165,11 +165,3 @@ func (r *Router) GetOverrideTarget(tag string) (string, error) {
 	}
 	return "", errors.New("cannot find tag")
 }
-
-// PickOutbound implements routing.BalancerPicker
-func (r *Router) GetBalancerOutboundTag(balancerTag string) (string, error) {
-	if b, ok := (*r.balancers.Load())[balancerTag]; ok {
-		return b.PickOutbound()
-	}
-	return "", errors.New("cannot find tag")
-}

common/buf/readv_windows.go

@@ -5,7 +5,8 @@ import (
 )
 
 type windowsReader struct {
-	bufs []syscall.WSABuf
+	bufs  []syscall.WSABuf
+	ready bool
 }
 
 func (r *windowsReader) Init(bs []*Buffer) {
@@ -15,6 +16,7 @@ func (r *windowsReader) Init(bs []*Buffer) {
 	for _, b := range bs {
 		r.bufs = append(r.bufs, syscall.WSABuf{Len: uint32(Size), Buf: &b.v[0]})
 	}
+	r.ready = false
 }
 
 func (r *windowsReader) Clear() {
@@ -25,6 +27,14 @@ func (r *windowsReader) Clear() {
 }
 
 func (r *windowsReader) Read(fd uintptr) int32 {
+	// On the first invocation, we return -1 to indicate "not ready"
+	// to make rawConn.Read wait for readability using the runtime's own mechanism
+	// because syscall.WSARecv() is a blocking call when used with nil OVERLAPPED
+	if !r.ready {
+		r.ready = true
+		return -1
+	}
+
 	var nBytes uint32
 	var flags uint32
 	err := syscall.WSARecv(syscall.Handle(fd), &r.bufs[0], uint32(len(r.bufs)), &nBytes, &flags, nil, nil)

features/routing/balancer.go

@@ -8,8 +8,3 @@ type BalancerOverrider interface {
 type BalancerPrincipleTarget interface {
 	GetPrincipleTarget(tag string) ([]string, error)
 }
-
-// BalancerPicker return the outbound tag selected by the given balancer
-type BalancerPicker interface {
-	GetBalancerOutboundTag(balancerTag string) (string, error)
-}

proxy/tun/handler.go

@@ -37,6 +37,25 @@ type Handler struct {
 	downlinkCounter stats.Counter
 }
 
+type tunUDPStatsWriter struct {
+	writer  buf.Writer
+	counter stats.Counter
+}
+
+func (w *tunUDPStatsWriter) WriteMultiBuffer(mb buf.MultiBuffer) error {
+	for len(mb) > 0 {
+		remaining, packet := buf.SplitFirst(mb)
+		packetSize := packet.Len()
+		if err := w.writer.WriteMultiBuffer(buf.MultiBuffer{packet}); err != nil {
+			buf.ReleaseMulti(remaining)
+			return err
+		}
+		w.counter.Add(int64(packetSize))
+		mb = remaining
+	}
+	return nil
+}
+
 // ConnectionHandler interface with the only method that stack is going to push new connections to
 type ConnectionHandler interface {
 	HandleConnection(conn net.Conn, destination net.Destination)
@@ -171,7 +190,8 @@ func (t *Handler) HandleConnection(conn net.Conn, destination net.Destination) {
 		return
 	}
 	source := net.DestinationFromAddr(remote)
-	if t.uplinkCounter != nil || t.downlinkCounter != nil {
+	isUDP := destination.Network == net.Network_UDP
+	if !isUDP && (t.uplinkCounter != nil || t.downlinkCounter != nil) {
 		conn = &stat.CounterConnection{
 			Connection:   conn,
 			ReadCounter:  t.uplinkCounter,
@@ -203,9 +223,18 @@ func (t *Handler) HandleConnection(conn net.Conn, destination net.Destination) {
 	})
 	errors.LogInfo(ctx, "processing from ", source, " to ", destination)
 
+	reader := &buf.TimeoutWrapperReader{Reader: buf.NewReader(conn)}
+	writer := buf.NewWriter(conn)
+	if isUDP {
+		reader.Counter = t.uplinkCounter
+		if t.downlinkCounter != nil {
+			writer = &tunUDPStatsWriter{writer: writer, counter: t.downlinkCounter}
+		}
+	}
+
 	link := &transport.Link{
-		Reader: &buf.TimeoutWrapperReader{Reader: buf.NewReader(conn)},
-		Writer: buf.NewWriter(conn),
+		Reader: reader,
+		Writer: writer,
 	}
 	if err := t.dispatcher.DispatchLink(ctx, destination, link); err != nil {
 		errors.LogError(ctx, errors.New("connection closed").Base(err))