From d70b997d84a0bd9cb073e22b1de5c31ba3bee18b Mon Sep 17 00:00:00 2001
From: Darien Raymond <admin@v2ray.com>
Date: Sun, 6 Nov 2016 14:32:04 +0100
Subject: [PATCH] test case for shadowsocks UDP

---
 proxy/shadowsocks/protocol.go      |  8 +++---
 proxy/shadowsocks/protocol_test.go | 41 ++++++++++++++++++++++++++++++
 2 files changed, 45 insertions(+), 4 deletions(-)

diff --git a/proxy/shadowsocks/protocol.go b/proxy/shadowsocks/protocol.go
index 8ad71fc5e..74f756f06 100644
--- a/proxy/shadowsocks/protocol.go
+++ b/proxy/shadowsocks/protocol.go
@@ -250,10 +250,10 @@ func EncodeUDPPacket(request *protocol.RequestHeader, payload *alloc.Buffer) (*a
 	case v2net.AddressFamilyIPv4:
 		buffer.AppendBytes(AddrTypeIPv4)
 		buffer.Append([]byte(request.Address.IP()))
-	case AddrTypeIPv6:
+	case v2net.AddressFamilyIPv6:
 		buffer.AppendBytes(AddrTypeIPv6)
 		buffer.Append([]byte(request.Address.IP()))
-	case AddrTypeDomain:
+	case v2net.AddressFamilyDomain:
 		buffer.AppendBytes(AddrTypeDomain, byte(len(request.Address.Domain())))
 		buffer.Append([]byte(request.Address.Domain()))
 	default:
@@ -265,7 +265,7 @@ func EncodeUDPPacket(request *protocol.RequestHeader, payload *alloc.Buffer) (*a
 
 	if request.Option.Has(RequestOptionOneTimeAuth) {
 		authenticator := NewAuthenticator(HeaderKeyGenerator(account.Key, iv))
-		buffer.Value[ivLen] |= 0x0F
+		buffer.Value[ivLen] |= 0x10
 
 		buffer.Value = authenticator.Authenticate(buffer.Value, buffer.Value[ivLen:])
 	}
@@ -334,7 +334,7 @@ func DecodeUDPPacket(user *protocol.User, payload *alloc.Buffer) (*protocol.Requ
 		request.Address = v2net.DomainAddress(string(payload.Value[1 : 1+domainLength]))
 		payload.SliceFrom(1 + domainLength)
 	default:
-		return nil, nil, errors.New("Shadowsocks|UDP: Unknown address type: " + err.Error())
+		return nil, nil, errors.New("Shadowsocks|UDP: Unknown address type")
 	}
 
 	request.Port = v2net.PortFromBytes(payload.Value[:2])
diff --git a/proxy/shadowsocks/protocol_test.go b/proxy/shadowsocks/protocol_test.go
index 3ab3a0bcf..0508540fe 100644
--- a/proxy/shadowsocks/protocol_test.go
+++ b/proxy/shadowsocks/protocol_test.go
@@ -75,3 +75,44 @@ func TestTCPRequest(t *testing.T) {
 	assert.Error(err).IsNil()
 	assert.Bytes(decodedData.Value).Equals([]byte("test string"))
 }
+
+func TestUDPReaderWriter(t *testing.T) {
+	assert := assert.On(t)
+
+	user := &protocol.User{
+		Account: loader.NewTypedSettings(&Account{
+			Password:   "test-password",
+			CipherType: CipherType_CHACHA20_IEFT,
+		}),
+	}
+	cache := alloc.NewBuffer().Clear()
+	writer := &UDPWriter{
+		Writer: cache,
+		Request: &protocol.RequestHeader{
+			Version: Version,
+			Address: v2net.DomainAddress("v2ray.com"),
+			Port:    123,
+			User:    user,
+			Option:  RequestOptionOneTimeAuth,
+		},
+	}
+
+	reader := &UDPReader{
+		Reader: cache,
+		User:   user,
+	}
+
+	err := writer.Write(alloc.NewBuffer().Clear().AppendString("test payload"))
+	assert.Error(err).IsNil()
+
+	payload, err := reader.Read()
+	assert.Error(err).IsNil()
+	assert.String(payload.String()).Equals("test payload")
+
+	err = writer.Write(alloc.NewBuffer().Clear().AppendString("test payload 2"))
+	assert.Error(err).IsNil()
+
+	payload, err = reader.Read()
+	assert.Error(err).IsNil()
+	assert.String(payload.String()).Equals("test payload 2")
+}