test case for socks udp

proxy/socks/protocol.go

@@ -344,6 +344,10 @@ type UDPReader struct {
 	reader io.Reader
 }
 
+func NewUDPReader(reader io.Reader) *UDPReader {
+	return &UDPReader{reader: reader}
+}
+
 func (r *UDPReader) Read() (*buf.Buffer, error) {
 	b := buf.NewSmall()
 	if err := b.AppendSupplier(buf.ReadFrom(r.reader)); err != nil {
@@ -363,6 +367,13 @@ type UDPWriter struct {
 	writer  io.Writer
 }
 
+func NewUDPWriter(request *protocol.RequestHeader, writer io.Writer) *UDPWriter {
+	return &UDPWriter{
+		request: request,
+		writer:  writer,
+	}
+}
+
 func (w *UDPWriter) Write(b *buf.Buffer) error {
 	eb := EncodeUDPPacket(w.request, b.Bytes())
 	b.Release()

proxy/socks/protocol_test.go

@@ -0,0 +1,34 @@
+package socks_test
+
+import (
+	"testing"
+
+	"v2ray.com/core/common/buf"
+	"v2ray.com/core/common/net"
+	"v2ray.com/core/common/protocol"
+	. "v2ray.com/core/proxy/socks"
+	"v2ray.com/core/testing/assert"
+)
+
+func TestUDPEncoding(t *testing.T) {
+	assert := assert.On(t)
+
+	b := buf.New()
+
+	request := &protocol.RequestHeader{
+		Address: net.IPAddress([]byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6}),
+		Port:    1024,
+	}
+	writer := NewUDPWriter(request, b)
+
+	content := []byte{'a'}
+	payload := buf.New()
+	payload.Append(content)
+	assert.Error(writer.Write(payload)).IsNil()
+
+	reader := NewUDPReader(b)
+
+	decodedPayload, err := reader.Read()
+	assert.Error(err).IsNil()
+	assert.Bytes(decodedPayload.Bytes()).Equals(content)
+}