From d5f931ae8bd320a1d1d0888a66540817c4383eba Mon Sep 17 00:00:00 2001
From: Darien Raymond <admin@v2ray.com>
Date: Tue, 2 May 2017 23:36:37 +0200
Subject: [PATCH] fix udp handling in dokodemo and mux

---
 app/proxyman/mux/mux.go         |  9 +++++----
 proxy/dokodemo/dokodemo.go      |  9 +++++++--
 testing/scenarios/vmess_test.go | 21 ++++++++++++++++-----
 3 files changed, 28 insertions(+), 11 deletions(-)

diff --git a/app/proxyman/mux/mux.go b/app/proxyman/mux/mux.go
index 650496f0..a0778c1e 100644
--- a/app/proxyman/mux/mux.go
+++ b/app/proxyman/mux/mux.go
@@ -133,11 +133,12 @@ func (m *Client) monitor() {
 
 func fetchInput(ctx context.Context, s *Session, output buf.Writer) {
 	dest, _ := proxy.TargetFromContext(ctx)
-	writer := &Writer{
-		dest:   dest,
-		id:     s.ID,
-		writer: output,
+	transferType := protocol.TransferTypeStream
+	if dest.Network == net.Network_UDP {
+		transferType = protocol.TransferTypePacket
 	}
+	s.transferType = transferType
+	writer := NewWriter(s.ID, dest, output, transferType)
 	defer writer.Close()
 	defer s.CloseUplink()
 
diff --git a/proxy/dokodemo/dokodemo.go b/proxy/dokodemo/dokodemo.go
index c6500579..5b031aa2 100644
--- a/proxy/dokodemo/dokodemo.go
+++ b/proxy/dokodemo/dokodemo.go
@@ -84,9 +84,14 @@ func (d *DokodemoDoor) Process(ctx context.Context, network net.Network, conn in
 	})
 
 	responseDone := signal.ExecuteAsync(func() error {
-		v2writer := buf.NewWriter(conn)
+		var writer buf.Writer
+		if network == net.Network_TCP {
+			writer = buf.NewWriter(conn)
+		} else {
+			writer = buf.NewSequentialWriter(conn)
+		}
 
-		if err := buf.Copy(inboundRay.InboundOutput(), v2writer, buf.UpdateActivity(timer)); err != nil {
+		if err := buf.Copy(inboundRay.InboundOutput(), writer, buf.UpdateActivity(timer)); err != nil {
 			return newError("failed to transport response").Base(err)
 		}
 		return nil
diff --git a/testing/scenarios/vmess_test.go b/testing/scenarios/vmess_test.go
index 9a89b42f..3fa82a7a 100644
--- a/testing/scenarios/vmess_test.go
+++ b/testing/scenarios/vmess_test.go
@@ -1198,15 +1198,26 @@ func TestVMessGCMMuxUDP(t *testing.T) {
 				})
 				assert.Error(err).IsNil()
 
+				conn.SetDeadline(time.Now().Add(time.Second * 10))
+
 				payload := make([]byte, 1024)
 				rand.Read(payload)
 
-				nBytes, err := conn.Write(payload)
-				assert.Error(err).IsNil()
-				assert.Int(nBytes).Equals(len(payload))
+				for j := 0; j < 10; j++ {
+					nBytes, _, err := conn.WriteMsgUDP(payload, nil, nil)
+					assert.Error(err).IsNil()
+					assert.Int(nBytes).Equals(len(payload))
+				}
+
+				response := make([]byte, 1024)
+				oob := make([]byte, 16)
+				for j := 0; j < 10; j++ {
+					nBytes, _, _, _, err := conn.ReadMsgUDP(response, oob)
+					assert.Error(err).IsNil()
+					assert.Int(nBytes).Equals(1024)
+					assert.Bytes(response).Equals(xor(payload))
+				}
 
-				response := readFrom(conn, time.Second*5, 1024)
-				assert.Bytes(response).Equals(xor(payload))
 				assert.Error(conn.Close()).IsNil()
 				wg.Done()
 			}()