v2fly/common/mux/mux_test.go

package mux_test

import (
	"io"
	"testing"

	"v2ray.com/core/common"
	"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 = append(mb, 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.MultiBuffer{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)
	common.Must(err)
	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))
	common.Must(err)
	assert(len(data), Equals, 1)
	assert(data[0].String(), Equals, "abcd")

	err = meta.Unmarshal(bytesReader)
	common.Must(err)
	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)
	common.Must(err)
	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))
	common.Must(err)
	assert(len(data), Equals, 1)
	assert(data[0].String(), Equals, "efgh")

	err = meta.Unmarshal(bytesReader)
	common.Must(err)
	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))
	common.Must(err)
	assert(len(data), Equals, 1)
	assert(data[0].String(), Equals, "x")

	err = meta.Unmarshal(bytesReader)
	common.Must(err)
	assert(byte(meta.SessionStatus), Equals, byte(SessionStatusEnd))
	assert(meta.SessionID, Equals, uint16(1))
	assert(byte(meta.Option), Equals, byte(0))

	err = meta.Unmarshal(bytesReader)
	common.Must(err)
	assert(byte(meta.SessionStatus), Equals, byte(SessionStatusEnd))
	assert(meta.SessionID, Equals, uint16(3))
	assert(byte(meta.Option), Equals, byte(0))

	err = meta.Unmarshal(bytesReader)
	common.Must(err)
	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))
	common.Must(err)
	assert(len(data), Equals, 1)
	assert(data[0].String(), Equals, "y")

	err = meta.Unmarshal(bytesReader)
	common.Must(err)
	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)
}