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>

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)

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))