github.com/metacubex/quic-go@v0.44.1-0.20240520163451-20b689a59136/crypto_stream_manager_test.go (about)

     1  package quic
     2  
     3  import (
     4  	"github.com/metacubex/quic-go/internal/protocol"
     5  	"github.com/metacubex/quic-go/internal/wire"
     6  
     7  	. "github.com/onsi/ginkgo/v2"
     8  	. "github.com/onsi/gomega"
     9  )
    10  
    11  var _ = Describe("Crypto Stream Manager", func() {
    12  	var (
    13  		csm *cryptoStreamManager
    14  		cs  *MockCryptoDataHandler
    15  
    16  		initialStream   *MockCryptoStream
    17  		handshakeStream *MockCryptoStream
    18  		oneRTTStream    *MockCryptoStream
    19  	)
    20  
    21  	BeforeEach(func() {
    22  		initialStream = NewMockCryptoStream(mockCtrl)
    23  		handshakeStream = NewMockCryptoStream(mockCtrl)
    24  		oneRTTStream = NewMockCryptoStream(mockCtrl)
    25  		cs = NewMockCryptoDataHandler(mockCtrl)
    26  		csm = newCryptoStreamManager(cs, initialStream, handshakeStream, oneRTTStream)
    27  	})
    28  
    29  	It("passes messages to the initial stream", func() {
    30  		cf := &wire.CryptoFrame{Data: []byte("foobar")}
    31  		initialStream.EXPECT().HandleCryptoFrame(cf)
    32  		initialStream.EXPECT().GetCryptoData().Return([]byte("foobar"))
    33  		initialStream.EXPECT().GetCryptoData()
    34  		cs.EXPECT().HandleMessage([]byte("foobar"), protocol.EncryptionInitial)
    35  		Expect(csm.HandleCryptoFrame(cf, protocol.EncryptionInitial)).To(Succeed())
    36  	})
    37  
    38  	It("passes messages to the handshake stream", func() {
    39  		cf := &wire.CryptoFrame{Data: []byte("foobar")}
    40  		handshakeStream.EXPECT().HandleCryptoFrame(cf)
    41  		handshakeStream.EXPECT().GetCryptoData().Return([]byte("foobar"))
    42  		handshakeStream.EXPECT().GetCryptoData()
    43  		cs.EXPECT().HandleMessage([]byte("foobar"), protocol.EncryptionHandshake)
    44  		Expect(csm.HandleCryptoFrame(cf, protocol.EncryptionHandshake)).To(Succeed())
    45  	})
    46  
    47  	It("passes messages to the 1-RTT stream", func() {
    48  		cf := &wire.CryptoFrame{Data: []byte("foobar")}
    49  		oneRTTStream.EXPECT().HandleCryptoFrame(cf)
    50  		oneRTTStream.EXPECT().GetCryptoData().Return([]byte("foobar"))
    51  		oneRTTStream.EXPECT().GetCryptoData()
    52  		cs.EXPECT().HandleMessage([]byte("foobar"), protocol.Encryption1RTT)
    53  		Expect(csm.HandleCryptoFrame(cf, protocol.Encryption1RTT)).To(Succeed())
    54  	})
    55  
    56  	It("doesn't call the message handler, if there's no message", func() {
    57  		cf := &wire.CryptoFrame{Data: []byte("foobar")}
    58  		handshakeStream.EXPECT().HandleCryptoFrame(cf)
    59  		handshakeStream.EXPECT().GetCryptoData() // don't return any data to handle
    60  		// don't EXPECT any calls to HandleMessage()
    61  		Expect(csm.HandleCryptoFrame(cf, protocol.EncryptionHandshake)).To(Succeed())
    62  	})
    63  
    64  	It("processes all messages", func() {
    65  		cf := &wire.CryptoFrame{Data: []byte("foobar")}
    66  		handshakeStream.EXPECT().HandleCryptoFrame(cf)
    67  		handshakeStream.EXPECT().GetCryptoData().Return([]byte("foo"))
    68  		handshakeStream.EXPECT().GetCryptoData().Return([]byte("bar"))
    69  		handshakeStream.EXPECT().GetCryptoData()
    70  		cs.EXPECT().HandleMessage([]byte("foo"), protocol.EncryptionHandshake)
    71  		cs.EXPECT().HandleMessage([]byte("bar"), protocol.EncryptionHandshake)
    72  		Expect(csm.HandleCryptoFrame(cf, protocol.EncryptionHandshake)).To(Succeed())
    73  	})
    74  
    75  	It("errors for unknown encryption levels", func() {
    76  		err := csm.HandleCryptoFrame(&wire.CryptoFrame{}, 42)
    77  		Expect(err).To(HaveOccurred())
    78  		Expect(err.Error()).To(ContainSubstring("received CRYPTO frame with unexpected encryption level"))
    79  	})
    80  
    81  	It("drops Initial", func() {
    82  		initialStream.EXPECT().Finish()
    83  		Expect(csm.Drop(protocol.EncryptionInitial)).To(Succeed())
    84  	})
    85  
    86  	It("drops Handshake", func() {
    87  		handshakeStream.EXPECT().Finish()
    88  		Expect(csm.Drop(protocol.EncryptionHandshake)).To(Succeed())
    89  	})
    90  })