move mux to common

common/mux/mux_test.go

@@ -0,0 +1,140 @@
+package mux_test
+
+import (
+	"io"
+	"testing"
+
+	"v2ray.com/core/common/buf"
+	. "v2ray.com/core/common/mux"
+	"v2ray.com/core/common/net"
+	"v2ray.com/core/common/protocol"
+	"v2ray.com/core/transport/pipe"
+	. "v2ray.com/ext/assert"
+)
+
+func readAll(reader buf.Reader) (buf.MultiBuffer, error) {
+	var mb buf.MultiBuffer
+	for {
+		b, err := reader.ReadMultiBuffer()
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			return nil, err
+		}
+		mb.AppendMulti(b)
+	}
+	return mb, nil
+}
+
+func TestReaderWriter(t *testing.T) {
+	assert := With(t)
+
+	pReader, pWriter := pipe.New(pipe.WithSizeLimit(1024))
+
+	dest := net.TCPDestination(net.DomainAddress("v2ray.com"), 80)
+	writer := NewWriter(1, dest, pWriter, protocol.TransferTypeStream)
+
+	dest2 := net.TCPDestination(net.LocalHostIP, 443)
+	writer2 := NewWriter(2, dest2, pWriter, protocol.TransferTypeStream)
+
+	dest3 := net.TCPDestination(net.LocalHostIPv6, 18374)
+	writer3 := NewWriter(3, dest3, pWriter, protocol.TransferTypeStream)
+
+	writePayload := func(writer *Writer, payload ...byte) error {
+		b := buf.New()
+		b.Write(payload)
+		return writer.WriteMultiBuffer(buf.NewMultiBufferValue(b))
+	}
+
+	assert(writePayload(writer, 'a', 'b', 'c', 'd'), IsNil)
+	assert(writePayload(writer2), IsNil)
+
+	assert(writePayload(writer, 'e', 'f', 'g', 'h'), IsNil)
+	assert(writePayload(writer3, 'x'), IsNil)
+
+	writer.Close()
+	writer3.Close()
+
+	assert(writePayload(writer2, 'y'), IsNil)
+	writer2.Close()
+
+	bytesReader := &buf.BufferedReader{Reader: pReader}
+
+	var meta FrameMetadata
+	err := meta.Unmarshal(bytesReader)
+	assert(err, IsNil)
+	assert(meta.SessionID, Equals, uint16(1))
+	assert(byte(meta.SessionStatus), Equals, byte(SessionStatusNew))
+	assert(meta.Target, Equals, dest)
+	assert(byte(meta.Option), Equals, byte(OptionData))
+
+	data, err := readAll(NewStreamReader(bytesReader))
+	assert(err, IsNil)
+	assert(len(data), Equals, 1)
+	assert(data[0].String(), Equals, "abcd")
+
+	err = meta.Unmarshal(bytesReader)
+	assert(err, IsNil)
+	assert(byte(meta.SessionStatus), Equals, byte(SessionStatusNew))
+	assert(meta.SessionID, Equals, uint16(2))
+	assert(byte(meta.Option), Equals, byte(0))
+	assert(meta.Target, Equals, dest2)
+
+	err = meta.Unmarshal(bytesReader)
+	assert(err, IsNil)
+	assert(byte(meta.SessionStatus), Equals, byte(SessionStatusKeep))
+	assert(meta.SessionID, Equals, uint16(1))
+	assert(byte(meta.Option), Equals, byte(1))
+
+	data, err = readAll(NewStreamReader(bytesReader))
+	assert(err, IsNil)
+	assert(len(data), Equals, 1)
+	assert(data[0].String(), Equals, "efgh")
+
+	err = meta.Unmarshal(bytesReader)
+	assert(err, IsNil)
+	assert(byte(meta.SessionStatus), Equals, byte(SessionStatusNew))
+	assert(meta.SessionID, Equals, uint16(3))
+	assert(byte(meta.Option), Equals, byte(1))
+	assert(meta.Target, Equals, dest3)
+
+	data, err = readAll(NewStreamReader(bytesReader))
+	assert(err, IsNil)
+	assert(len(data), Equals, 1)
+	assert(data[0].String(), Equals, "x")
+
+	err = meta.Unmarshal(bytesReader)
+	assert(err, IsNil)
+	assert(byte(meta.SessionStatus), Equals, byte(SessionStatusEnd))
+	assert(meta.SessionID, Equals, uint16(1))
+	assert(byte(meta.Option), Equals, byte(0))
+
+	err = meta.Unmarshal(bytesReader)
+	assert(err, IsNil)
+	assert(byte(meta.SessionStatus), Equals, byte(SessionStatusEnd))
+	assert(meta.SessionID, Equals, uint16(3))
+	assert(byte(meta.Option), Equals, byte(0))
+
+	err = meta.Unmarshal(bytesReader)
+	assert(err, IsNil)
+	assert(byte(meta.SessionStatus), Equals, byte(SessionStatusKeep))
+	assert(meta.SessionID, Equals, uint16(2))
+	assert(byte(meta.Option), Equals, byte(1))
+
+	data, err = readAll(NewStreamReader(bytesReader))
+	assert(err, IsNil)
+	assert(len(data), Equals, 1)
+	assert(data[0].String(), Equals, "y")
+
+	err = meta.Unmarshal(bytesReader)
+	assert(err, IsNil)
+	assert(byte(meta.SessionStatus), Equals, byte(SessionStatusEnd))
+	assert(meta.SessionID, Equals, uint16(2))
+	assert(byte(meta.Option), Equals, byte(0))
+
+	pWriter.Close()
+
+	err = meta.Unmarshal(bytesReader)
+	assert(err, IsNotNil)
+}