github.com/apernet/quic-go@v0.43.1-0.20240515053213-5e9e635fd9f0/integrationtests/self/uni_stream_test.go (about)

     1  package self_test
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"io"
     7  	"net"
     8  	"sync"
     9  
    10  	"github.com/apernet/quic-go"
    11  	"github.com/apernet/quic-go/internal/protocol"
    12  
    13  	. "github.com/onsi/ginkgo/v2"
    14  	. "github.com/onsi/gomega"
    15  )
    16  
    17  var _ = Describe("Unidirectional Streams", func() {
    18  	const numStreams = 500
    19  
    20  	var (
    21  		server     *quic.Listener
    22  		serverAddr string
    23  	)
    24  
    25  	BeforeEach(func() {
    26  		var err error
    27  		server, err = quic.ListenAddr("localhost:0", getTLSConfig(), getQuicConfig(nil))
    28  		Expect(err).ToNot(HaveOccurred())
    29  		serverAddr = fmt.Sprintf("localhost:%d", server.Addr().(*net.UDPAddr).Port)
    30  	})
    31  
    32  	AfterEach(func() {
    33  		server.Close()
    34  	})
    35  
    36  	dataForStream := func(id protocol.StreamID) []byte {
    37  		return GeneratePRData(10 * int(id))
    38  	}
    39  
    40  	runSendingPeer := func(conn quic.Connection) {
    41  		for i := 0; i < numStreams; i++ {
    42  			str, err := conn.OpenUniStreamSync(context.Background())
    43  			Expect(err).ToNot(HaveOccurred())
    44  			go func() {
    45  				defer GinkgoRecover()
    46  				_, err := str.Write(dataForStream(str.StreamID()))
    47  				Expect(err).ToNot(HaveOccurred())
    48  				Expect(str.Close()).To(Succeed())
    49  			}()
    50  		}
    51  	}
    52  
    53  	runReceivingPeer := func(conn quic.Connection) {
    54  		var wg sync.WaitGroup
    55  		wg.Add(numStreams)
    56  		for i := 0; i < numStreams; i++ {
    57  			str, err := conn.AcceptUniStream(context.Background())
    58  			Expect(err).ToNot(HaveOccurred())
    59  			go func() {
    60  				defer GinkgoRecover()
    61  				defer wg.Done()
    62  				data, err := io.ReadAll(str)
    63  				Expect(err).ToNot(HaveOccurred())
    64  				Expect(data).To(Equal(dataForStream(str.StreamID())))
    65  			}()
    66  		}
    67  		wg.Wait()
    68  	}
    69  
    70  	It(fmt.Sprintf("client opening %d streams to a server", numStreams), func() {
    71  		go func() {
    72  			defer GinkgoRecover()
    73  			conn, err := server.Accept(context.Background())
    74  			Expect(err).ToNot(HaveOccurred())
    75  			runReceivingPeer(conn)
    76  			conn.CloseWithError(0, "")
    77  		}()
    78  
    79  		client, err := quic.DialAddr(
    80  			context.Background(),
    81  			serverAddr,
    82  			getTLSClientConfig(),
    83  			getQuicConfig(nil),
    84  		)
    85  		Expect(err).ToNot(HaveOccurred())
    86  		runSendingPeer(client)
    87  		<-client.Context().Done()
    88  	})
    89  
    90  	It(fmt.Sprintf("server opening %d streams to a client", numStreams), func() {
    91  		done := make(chan struct{})
    92  		go func() {
    93  			defer GinkgoRecover()
    94  			defer close(done)
    95  			conn, err := server.Accept(context.Background())
    96  			Expect(err).ToNot(HaveOccurred())
    97  			runSendingPeer(conn)
    98  			<-conn.Context().Done()
    99  		}()
   100  
   101  		client, err := quic.DialAddr(
   102  			context.Background(),
   103  			serverAddr,
   104  			getTLSClientConfig(),
   105  			getQuicConfig(nil),
   106  		)
   107  		Expect(err).ToNot(HaveOccurred())
   108  		runReceivingPeer(client)
   109  		client.CloseWithError(0, "")
   110  	})
   111  
   112  	It(fmt.Sprintf("client and server opening %d streams each and sending data to the peer", numStreams), func() {
   113  		done1 := make(chan struct{})
   114  		go func() {
   115  			defer GinkgoRecover()
   116  			conn, err := server.Accept(context.Background())
   117  			Expect(err).ToNot(HaveOccurred())
   118  			done := make(chan struct{})
   119  			go func() {
   120  				defer GinkgoRecover()
   121  				runReceivingPeer(conn)
   122  				close(done)
   123  			}()
   124  			runSendingPeer(conn)
   125  			<-done
   126  			close(done1)
   127  		}()
   128  
   129  		client, err := quic.DialAddr(
   130  			context.Background(),
   131  			serverAddr,
   132  			getTLSClientConfig(),
   133  			getQuicConfig(nil),
   134  		)
   135  		Expect(err).ToNot(HaveOccurred())
   136  		done2 := make(chan struct{})
   137  		go func() {
   138  			defer GinkgoRecover()
   139  			runSendingPeer(client)
   140  			close(done2)
   141  		}()
   142  		runReceivingPeer(client)
   143  		<-done1
   144  		<-done2
   145  		client.CloseWithError(0, "")
   146  	})
   147  })