github.com/sagernet/sing@v0.4.0-beta.19.0.20240518125136-f67a0988a636/common/bufio/vectorised_test.go (about)

     1  package bufio
     2  
     3  import (
     4  	"crypto/rand"
     5  	"io"
     6  	"testing"
     7  
     8  	"github.com/stretchr/testify/require"
     9  )
    10  
    11  func TestWriteVectorised(t *testing.T) {
    12  	t.Parallel()
    13  	inputConn, outputConn := TCPPipe(t)
    14  	vectorisedWriter, created := CreateVectorisedWriter(inputConn)
    15  	require.True(t, created)
    16  	require.NotNil(t, vectorisedWriter)
    17  	var bufA [1024]byte
    18  	var bufB [1024]byte
    19  	var bufC [2048]byte
    20  	_, err := io.ReadFull(rand.Reader, bufA[:])
    21  	require.NoError(t, err)
    22  	_, err = io.ReadFull(rand.Reader, bufB[:])
    23  	require.NoError(t, err)
    24  	copy(bufC[:], bufA[:])
    25  	copy(bufC[1024:], bufB[:])
    26  	finish := Timeout(t)
    27  	_, err = WriteVectorised(vectorisedWriter, [][]byte{bufA[:], bufB[:]})
    28  	require.NoError(t, err)
    29  	output := make([]byte, 2048)
    30  	_, err = io.ReadFull(outputConn, output)
    31  	finish()
    32  	require.NoError(t, err)
    33  	require.Equal(t, bufC[:], output)
    34  }
    35  
    36  func TestWriteVectorisedPacket(t *testing.T) {
    37  	t.Parallel()
    38  	inputConn, outputConn, outputAddr := UDPPipe(t)
    39  	vectorisedWriter, created := CreateVectorisedPacketWriter(inputConn)
    40  	require.True(t, created)
    41  	require.NotNil(t, vectorisedWriter)
    42  	var bufA [1024]byte
    43  	var bufB [1024]byte
    44  	var bufC [2048]byte
    45  	_, err := io.ReadFull(rand.Reader, bufA[:])
    46  	require.NoError(t, err)
    47  	_, err = io.ReadFull(rand.Reader, bufB[:])
    48  	require.NoError(t, err)
    49  	copy(bufC[:], bufA[:])
    50  	copy(bufC[1024:], bufB[:])
    51  	finish := Timeout(t)
    52  	_, err = WriteVectorisedPacket(vectorisedWriter, [][]byte{bufA[:], bufB[:]}, outputAddr)
    53  	require.NoError(t, err)
    54  	output := make([]byte, 2048)
    55  	n, _, err := outputConn.ReadFrom(output)
    56  	finish()
    57  	require.NoError(t, err)
    58  	require.Equal(t, 2048, n)
    59  	require.Equal(t, bufC[:], output)
    60  }