github.com/xraypb/Xray-core@v1.8.1/common/mux/client_test.go (about)

     1  package mux_test
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  	"time"
     7  
     8  	"github.com/golang/mock/gomock"
     9  	"github.com/xraypb/Xray-core/common"
    10  	"github.com/xraypb/Xray-core/common/errors"
    11  	"github.com/xraypb/Xray-core/common/mux"
    12  	"github.com/xraypb/Xray-core/common/net"
    13  	"github.com/xraypb/Xray-core/common/session"
    14  	"github.com/xraypb/Xray-core/testing/mocks"
    15  	"github.com/xraypb/Xray-core/transport"
    16  	"github.com/xraypb/Xray-core/transport/pipe"
    17  )
    18  
    19  func TestIncrementalPickerFailure(t *testing.T) {
    20  	mockCtl := gomock.NewController(t)
    21  	defer mockCtl.Finish()
    22  
    23  	mockWorkerFactory := mocks.NewMuxClientWorkerFactory(mockCtl)
    24  	mockWorkerFactory.EXPECT().Create().Return(nil, errors.New("test"))
    25  
    26  	picker := mux.IncrementalWorkerPicker{
    27  		Factory: mockWorkerFactory,
    28  	}
    29  
    30  	_, err := picker.PickAvailable()
    31  	if err == nil {
    32  		t.Error("expected error, but nil")
    33  	}
    34  }
    35  
    36  func TestClientWorkerEOF(t *testing.T) {
    37  	reader, writer := pipe.New(pipe.WithoutSizeLimit())
    38  	common.Must(writer.Close())
    39  
    40  	worker, err := mux.NewClientWorker(transport.Link{Reader: reader, Writer: writer}, mux.ClientStrategy{})
    41  	common.Must(err)
    42  
    43  	time.Sleep(time.Millisecond * 500)
    44  
    45  	f := worker.Dispatch(context.Background(), nil)
    46  	if f {
    47  		t.Error("expected failed dispatching, but actually not")
    48  	}
    49  }
    50  
    51  func TestClientWorkerClose(t *testing.T) {
    52  	mockCtl := gomock.NewController(t)
    53  	defer mockCtl.Finish()
    54  
    55  	r1, w1 := pipe.New(pipe.WithoutSizeLimit())
    56  	worker1, err := mux.NewClientWorker(transport.Link{
    57  		Reader: r1,
    58  		Writer: w1,
    59  	}, mux.ClientStrategy{
    60  		MaxConcurrency: 4,
    61  		MaxConnection:  4,
    62  	})
    63  	common.Must(err)
    64  
    65  	r2, w2 := pipe.New(pipe.WithoutSizeLimit())
    66  	worker2, err := mux.NewClientWorker(transport.Link{
    67  		Reader: r2,
    68  		Writer: w2,
    69  	}, mux.ClientStrategy{
    70  		MaxConcurrency: 4,
    71  		MaxConnection:  4,
    72  	})
    73  	common.Must(err)
    74  
    75  	factory := mocks.NewMuxClientWorkerFactory(mockCtl)
    76  	gomock.InOrder(
    77  		factory.EXPECT().Create().Return(worker1, nil),
    78  		factory.EXPECT().Create().Return(worker2, nil),
    79  	)
    80  
    81  	picker := &mux.IncrementalWorkerPicker{
    82  		Factory: factory,
    83  	}
    84  	manager := &mux.ClientManager{
    85  		Picker: picker,
    86  	}
    87  
    88  	tr1, tw1 := pipe.New(pipe.WithoutSizeLimit())
    89  	ctx1 := session.ContextWithOutbound(context.Background(), &session.Outbound{
    90  		Target: net.TCPDestination(net.DomainAddress("www.example.com"), 80),
    91  	})
    92  	common.Must(manager.Dispatch(ctx1, &transport.Link{
    93  		Reader: tr1,
    94  		Writer: tw1,
    95  	}))
    96  	defer tw1.Close()
    97  
    98  	common.Must(w1.Close())
    99  
   100  	time.Sleep(time.Millisecond * 500)
   101  	if !worker1.Closed() {
   102  		t.Error("worker1 is not finished")
   103  	}
   104  
   105  	tr2, tw2 := pipe.New(pipe.WithoutSizeLimit())
   106  	ctx2 := session.ContextWithOutbound(context.Background(), &session.Outbound{
   107  		Target: net.TCPDestination(net.DomainAddress("www.example.com"), 80),
   108  	})
   109  	common.Must(manager.Dispatch(ctx2, &transport.Link{
   110  		Reader: tr2,
   111  		Writer: tw2,
   112  	}))
   113  	defer tw2.Close()
   114  
   115  	common.Must(w2.Close())
   116  }