github.com/xraypb/xray-core@v1.6.6/common/mux/mux_test.go (about)

     1  package mux_test
     2  
     3  import (
     4  	"io"
     5  	"testing"
     6  
     7  	"github.com/google/go-cmp/cmp"
     8  	"github.com/xraypb/xray-core/common"
     9  	"github.com/xraypb/xray-core/common/buf"
    10  	. "github.com/xraypb/xray-core/common/mux"
    11  	"github.com/xraypb/xray-core/common/net"
    12  	"github.com/xraypb/xray-core/common/protocol"
    13  	"github.com/xraypb/xray-core/transport/pipe"
    14  )
    15  
    16  func readAll(reader buf.Reader) (buf.MultiBuffer, error) {
    17  	var mb buf.MultiBuffer
    18  	for {
    19  		b, err := reader.ReadMultiBuffer()
    20  		if err == io.EOF {
    21  			break
    22  		}
    23  		if err != nil {
    24  			return nil, err
    25  		}
    26  		mb = append(mb, b...)
    27  	}
    28  	return mb, nil
    29  }
    30  
    31  func TestReaderWriter(t *testing.T) {
    32  	pReader, pWriter := pipe.New(pipe.WithSizeLimit(1024))
    33  
    34  	dest := net.TCPDestination(net.DomainAddress("example.com"), 80)
    35  	writer := NewWriter(1, dest, pWriter, protocol.TransferTypeStream)
    36  
    37  	dest2 := net.TCPDestination(net.LocalHostIP, 443)
    38  	writer2 := NewWriter(2, dest2, pWriter, protocol.TransferTypeStream)
    39  
    40  	dest3 := net.TCPDestination(net.LocalHostIPv6, 18374)
    41  	writer3 := NewWriter(3, dest3, pWriter, protocol.TransferTypeStream)
    42  
    43  	writePayload := func(writer *Writer, payload ...byte) error {
    44  		b := buf.New()
    45  		b.Write(payload)
    46  		return writer.WriteMultiBuffer(buf.MultiBuffer{b})
    47  	}
    48  
    49  	common.Must(writePayload(writer, 'a', 'b', 'c', 'd'))
    50  	common.Must(writePayload(writer2))
    51  
    52  	common.Must(writePayload(writer, 'e', 'f', 'g', 'h'))
    53  	common.Must(writePayload(writer3, 'x'))
    54  
    55  	writer.Close()
    56  	writer3.Close()
    57  
    58  	common.Must(writePayload(writer2, 'y'))
    59  	writer2.Close()
    60  
    61  	bytesReader := &buf.BufferedReader{Reader: pReader}
    62  
    63  	{
    64  		var meta FrameMetadata
    65  		common.Must(meta.Unmarshal(bytesReader))
    66  		if r := cmp.Diff(meta, FrameMetadata{
    67  			SessionID:     1,
    68  			SessionStatus: SessionStatusNew,
    69  			Target:        dest,
    70  			Option:        OptionData,
    71  		}); r != "" {
    72  			t.Error("metadata: ", r)
    73  		}
    74  
    75  		data, err := readAll(NewStreamReader(bytesReader))
    76  		common.Must(err)
    77  		if s := data.String(); s != "abcd" {
    78  			t.Error("data: ", s)
    79  		}
    80  	}
    81  
    82  	{
    83  		var meta FrameMetadata
    84  		common.Must(meta.Unmarshal(bytesReader))
    85  		if r := cmp.Diff(meta, FrameMetadata{
    86  			SessionStatus: SessionStatusNew,
    87  			SessionID:     2,
    88  			Option:        0,
    89  			Target:        dest2,
    90  		}); r != "" {
    91  			t.Error("meta: ", r)
    92  		}
    93  	}
    94  
    95  	{
    96  		var meta FrameMetadata
    97  		common.Must(meta.Unmarshal(bytesReader))
    98  		if r := cmp.Diff(meta, FrameMetadata{
    99  			SessionID:     1,
   100  			SessionStatus: SessionStatusKeep,
   101  			Option:        1,
   102  		}); r != "" {
   103  			t.Error("meta: ", r)
   104  		}
   105  
   106  		data, err := readAll(NewStreamReader(bytesReader))
   107  		common.Must(err)
   108  		if s := data.String(); s != "efgh" {
   109  			t.Error("data: ", s)
   110  		}
   111  	}
   112  
   113  	{
   114  		var meta FrameMetadata
   115  		common.Must(meta.Unmarshal(bytesReader))
   116  		if r := cmp.Diff(meta, FrameMetadata{
   117  			SessionID:     3,
   118  			SessionStatus: SessionStatusNew,
   119  			Option:        1,
   120  			Target:        dest3,
   121  		}); r != "" {
   122  			t.Error("meta: ", r)
   123  		}
   124  
   125  		data, err := readAll(NewStreamReader(bytesReader))
   126  		common.Must(err)
   127  		if s := data.String(); s != "x" {
   128  			t.Error("data: ", s)
   129  		}
   130  	}
   131  
   132  	{
   133  		var meta FrameMetadata
   134  		common.Must(meta.Unmarshal(bytesReader))
   135  		if r := cmp.Diff(meta, FrameMetadata{
   136  			SessionID:     1,
   137  			SessionStatus: SessionStatusEnd,
   138  			Option:        0,
   139  		}); r != "" {
   140  			t.Error("meta: ", r)
   141  		}
   142  	}
   143  
   144  	{
   145  		var meta FrameMetadata
   146  		common.Must(meta.Unmarshal(bytesReader))
   147  		if r := cmp.Diff(meta, FrameMetadata{
   148  			SessionID:     3,
   149  			SessionStatus: SessionStatusEnd,
   150  			Option:        0,
   151  		}); r != "" {
   152  			t.Error("meta: ", r)
   153  		}
   154  	}
   155  
   156  	{
   157  		var meta FrameMetadata
   158  		common.Must(meta.Unmarshal(bytesReader))
   159  		if r := cmp.Diff(meta, FrameMetadata{
   160  			SessionID:     2,
   161  			SessionStatus: SessionStatusKeep,
   162  			Option:        1,
   163  		}); r != "" {
   164  			t.Error("meta: ", r)
   165  		}
   166  
   167  		data, err := readAll(NewStreamReader(bytesReader))
   168  		common.Must(err)
   169  		if s := data.String(); s != "y" {
   170  			t.Error("data: ", s)
   171  		}
   172  	}
   173  
   174  	{
   175  		var meta FrameMetadata
   176  		common.Must(meta.Unmarshal(bytesReader))
   177  		if r := cmp.Diff(meta, FrameMetadata{
   178  			SessionID:     2,
   179  			SessionStatus: SessionStatusEnd,
   180  			Option:        0,
   181  		}); r != "" {
   182  			t.Error("meta: ", r)
   183  		}
   184  	}
   185  
   186  	pWriter.Close()
   187  
   188  	{
   189  		var meta FrameMetadata
   190  		err := meta.Unmarshal(bytesReader)
   191  		if err == nil {
   192  			t.Error("nil error")
   193  		}
   194  	}
   195  }