From b4c1a560265569682f8c26687331d467dc9811e5 Mon Sep 17 00:00:00 2001
From: RPRX <63339210+RPRX@users.noreply.github.com>
Date: Sun, 16 Apr 2023 21:15:27 +0000
Subject: [PATCH] XUDP practice: MUST check the flag first & Add more comments

---
 common/mux/frame.go  |  8 +++++---
 common/mux/server.go |  2 +-
 common/xudp/xudp.go  | 15 ++++++++-------
 3 files changed, 14 insertions(+), 11 deletions(-)

diff --git a/common/mux/frame.go b/common/mux/frame.go
index ab57d771..d53ac202 100644
--- a/common/mux/frame.go
+++ b/common/mux/frame.go
@@ -82,8 +82,8 @@ func (f FrameMetadata) WriteTo(b *buf.Buffer) error {
 		if err := addrParser.WriteAddressPort(b, f.Target.Address, f.Target.Port); err != nil {
 			return err
 		}
-		if b.UDP != nil {
-			b.Write(f.GlobalID[:])
+		if b.UDP != nil { // make sure it's user's proxy request
+			b.Write(f.GlobalID[:]) // no need to check whether it's empty
 		}
 	} else if b.UDP != nil {
 		b.WriteByte(byte(TargetNetworkUDP))
@@ -126,7 +126,8 @@ func (f *FrameMetadata) UnmarshalFromBuffer(b *buf.Buffer) error {
 	f.Option = bitmask.Byte(b.Byte(3))
 	f.Target.Network = net.Network_Unknown
 
-	if f.SessionStatus == SessionStatusNew || (f.SessionStatus == SessionStatusKeep && b.Len() != 4) {
+	if f.SessionStatus == SessionStatusNew || (f.SessionStatus == SessionStatusKeep && b.Len() > 4 &&
+		TargetNetwork(b.Byte(4)) == TargetNetworkUDP) { // MUST check the flag first
 		if b.Len() < 8 {
 			return newError("insufficient buffer: ", b.Len())
 		}
@@ -148,6 +149,7 @@ func (f *FrameMetadata) UnmarshalFromBuffer(b *buf.Buffer) error {
 		}
 	}
 
+	// Application data is essential, to test whether the pipe is closed.
 	if f.SessionStatus == SessionStatusNew && f.Option.Has(OptionData) &&
 		f.Target.Network == net.Network_UDP && b.Len() >= 8 {
 		copy(f.GlobalID[:], b.Bytes())
diff --git a/common/mux/server.go b/common/mux/server.go
index 47a6d3dc..2d33189f 100644
--- a/common/mux/server.go
+++ b/common/mux/server.go
@@ -138,7 +138,7 @@ func (w *ServerWorker) handleStatusNew(ctx context.Context, meta *FrameMetadata,
 		}
 	}
 
-	if meta.GlobalID != [8]byte{} {
+	if meta.GlobalID != [8]byte{} { // MUST ignore empty Global ID
 		mb, err := NewPacketReader(reader, &meta.Target).ReadMultiBuffer()
 		if err != nil {
 			return err
diff --git a/common/xudp/xudp.go b/common/xudp/xudp.go
index a3df4e8b..513247c3 100644
--- a/common/xudp/xudp.go
+++ b/common/xudp/xudp.go
@@ -86,21 +86,21 @@ func (w *PacketWriter) WriteMultiBuffer(mb buf.MultiBuffer) error {
 		}
 
 		eb := buf.New()
-		eb.Write([]byte{0, 0, 0, 0})
+		eb.Write([]byte{0, 0, 0, 0}) // Meta data length; Mux Session ID
 		if w.Dest.Network == net.Network_UDP {
 			eb.WriteByte(1) // New
 			eb.WriteByte(1) // Opt
 			eb.WriteByte(2) // UDP
 			AddrParser.WriteAddressPort(eb, w.Dest.Address, w.Dest.Port)
 			if b.UDP != nil { // make sure it's user's proxy request
-				eb.Write(w.GlobalID[:])
+				eb.Write(w.GlobalID[:]) // no need to check whether it's empty
 			}
 			w.Dest.Network = net.Network_Unknown
 		} else {
 			eb.WriteByte(2) // Keep
-			eb.WriteByte(1)
+			eb.WriteByte(1) // Opt
 			if b.UDP != nil {
-				eb.WriteByte(2)
+				eb.WriteByte(2) // UDP
 				AddrParser.WriteAddressPort(eb, b.UDP.Address, b.UDP.Port)
 			}
 		}
@@ -148,9 +148,10 @@ func (r *PacketReader) ReadMultiBuffer() (buf.MultiBuffer, error) {
 		discard := false
 		switch b.Byte(2) {
 		case 2:
-			if l != 4 {
+			if l > 4 && b.Byte(4) == 2 { // MUST check the flag first
 				b.Advance(5)
-				addr, port, err := AddrParser.ReadAddressPort(nil, b) // read addr will read all content and clear b
+				// b.Clear() will be called automatically if all data had been read.
+				addr, port, err := AddrParser.ReadAddressPort(nil, b)
 				if err != nil {
 					b.Release()
 					return nil, err
@@ -167,6 +168,7 @@ func (r *PacketReader) ReadMultiBuffer() (buf.MultiBuffer, error) {
 			b.Release()
 			return nil, io.EOF
 		}
+		b.Clear() // in case there is padding (empty bytes) attached
 		if b.Byte(3) == 1 {
 			if _, err := io.ReadFull(r.Reader, r.cache); err != nil {
 				b.Release()
@@ -174,7 +176,6 @@ func (r *PacketReader) ReadMultiBuffer() (buf.MultiBuffer, error) {
 			}
 			length := int32(r.cache[0])<<8 | int32(r.cache[1])
 			if length > 0 {
-				b.Clear()
 				if _, err := b.ReadFullFrom(r.Reader, length); err != nil {
 					b.Release()
 					return nil, err