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