github.com/database64128/shadowsocks-go@v1.10.2-0.20240315062903-143a773533f1/direct/stream_test.go (about)

     1  package direct
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"crypto/rand"
     7  	"net/netip"
     8  	"sync"
     9  	"testing"
    10  
    11  	"github.com/database64128/shadowsocks-go/conn"
    12  	"github.com/database64128/shadowsocks-go/pipe"
    13  	"github.com/database64128/shadowsocks-go/zerocopy"
    14  )
    15  
    16  func TestDirectStreamReadWriter(t *testing.T) {
    17  	pl, pr := pipe.NewDuplexPipe()
    18  
    19  	l := DirectStreamReadWriter{
    20  		rw: pl,
    21  	}
    22  	r := DirectStreamReadWriter{
    23  		rw: pr,
    24  	}
    25  
    26  	zerocopy.ReadWriterTestFunc(t, &l, &r)
    27  }
    28  
    29  func testShadowsocksNoneStreamReadWriter(t *testing.T, ctx context.Context, clientInitialPayload []byte) {
    30  	pl, pr := pipe.NewDuplexPipe()
    31  	plo := zerocopy.SimpleDirectReadWriteCloserOpener{DirectReadWriteCloser: pl}
    32  
    33  	clientTargetAddr := conn.AddrFromIPPort(netip.AddrPortFrom(netip.IPv6Unspecified(), 53))
    34  	serverInitialPayload := make([]byte, len(clientInitialPayload))
    35  
    36  	var (
    37  		c                *DirectStreamReadWriter
    38  		s                *DirectStreamReadWriter
    39  		serverTargetAddr conn.Addr
    40  		nr               int
    41  		cerr, serr       error
    42  	)
    43  
    44  	var wg sync.WaitGroup
    45  	wg.Add(2)
    46  
    47  	go func() {
    48  		c, _, cerr = NewShadowsocksNoneStreamClientReadWriter(ctx, &plo, clientTargetAddr, clientInitialPayload)
    49  		wg.Done()
    50  	}()
    51  
    52  	go func() {
    53  		s, serverTargetAddr, serr = NewShadowsocksNoneStreamServerReadWriter(pr)
    54  		if len(serverInitialPayload) > 0 && serr == nil {
    55  			nr, serr = s.ReadZeroCopy(serverInitialPayload, 0, len(serverInitialPayload))
    56  		}
    57  		wg.Done()
    58  	}()
    59  
    60  	wg.Wait()
    61  	if cerr != nil {
    62  		t.Fatal(cerr)
    63  	}
    64  	if serr != nil {
    65  		t.Fatal(serr)
    66  	}
    67  
    68  	if nr != len(serverInitialPayload) {
    69  		t.Fatalf("Expected server initial payload bytes %d, got %d", len(serverInitialPayload), nr)
    70  	}
    71  	if !clientTargetAddr.Equals(serverTargetAddr) {
    72  		t.Errorf("Target address mismatch: c: %s, s: %s", clientTargetAddr, serverTargetAddr)
    73  	}
    74  	if !bytes.Equal(clientInitialPayload, serverInitialPayload) {
    75  		t.Error("Initial payload mismatch")
    76  	}
    77  
    78  	zerocopy.ReadWriterTestFunc(t, c, s)
    79  }
    80  
    81  func TestShadowsocksNoneStreamReadWriter(t *testing.T) {
    82  	ctx := context.Background()
    83  	initialPayload := make([]byte, 1024)
    84  	_, err := rand.Read(initialPayload)
    85  	if err != nil {
    86  		t.Fatal(err)
    87  	}
    88  
    89  	testShadowsocksNoneStreamReadWriter(t, ctx, nil)
    90  	testShadowsocksNoneStreamReadWriter(t, ctx, initialPayload)
    91  }
    92  
    93  func TestSocks5StreamReadWriter(t *testing.T) {
    94  	pl, pr := pipe.NewDuplexPipe()
    95  
    96  	clientTargetAddr := conn.AddrFromIPPort(netip.AddrPortFrom(netip.IPv6Unspecified(), 53))
    97  
    98  	var (
    99  		c                *DirectStreamReadWriter
   100  		s                *DirectStreamReadWriter
   101  		serverTargetAddr conn.Addr
   102  		cerr, serr       error
   103  	)
   104  
   105  	var wg sync.WaitGroup
   106  	wg.Add(2)
   107  
   108  	go func() {
   109  		c, cerr = NewSocks5StreamClientReadWriter(pl, clientTargetAddr)
   110  		wg.Done()
   111  	}()
   112  
   113  	go func() {
   114  		s, serverTargetAddr, serr = NewSocks5StreamServerReadWriter(pr, true, false)
   115  		wg.Done()
   116  	}()
   117  
   118  	wg.Wait()
   119  	if cerr != nil {
   120  		t.Fatal(cerr)
   121  	}
   122  	if serr != nil {
   123  		t.Fatal(serr)
   124  	}
   125  
   126  	if !clientTargetAddr.Equals(serverTargetAddr) {
   127  		t.Errorf("Target address mismatch: c: %s, s: %s", clientTargetAddr, serverTargetAddr)
   128  	}
   129  
   130  	zerocopy.ReadWriterTestFunc(t, c, s)
   131  }