reduce allocations when handling crypto data (quic-go#4615)

crypto_stream.go

@@ -21,8 +21,7 @@ type cryptoStream interface {
 }
 
 type cryptoStreamImpl struct {
-	queue  *frameSorter
-	msgBuf []byte
+	queue frameSorter
 
 	highestOffset protocol.ByteCount
 	finished      bool
@@ -32,7 +31,7 @@ type cryptoStreamImpl struct {
 }
 
 func newCryptoStream() cryptoStream {
-	return &cryptoStreamImpl{queue: newFrameSorter()}
+	return &cryptoStreamImpl{queue: *newFrameSorter()}
 }
 
 func (s *cryptoStreamImpl) HandleCryptoFrame(f *wire.CryptoFrame) error {
@@ -56,23 +55,13 @@ func (s *cryptoStreamImpl) HandleCryptoFrame(f *wire.CryptoFrame) error {
 		return nil
 	}
 	s.highestOffset = max(s.highestOffset, highestOffset)
-	if err := s.queue.Push(f.Data, f.Offset, nil); err != nil {
-		return err
-	}
-	for {
-		_, data, _ := s.queue.Pop()
-		if data == nil {
-			return nil
-		}
-		s.msgBuf = append(s.msgBuf, data...)
-	}
+	return s.queue.Push(f.Data, f.Offset, nil)
 }
 
 // GetCryptoData retrieves data that was received in CRYPTO frames
 func (s *cryptoStreamImpl) GetCryptoData() []byte {
-	b := s.msgBuf
-	s.msgBuf = nil
-	return b
+	_, data, _ := s.queue.Pop()
+	return data
 }
 
 func (s *cryptoStreamImpl) Finish() error {

crypto_stream_test.go

@@ -41,12 +41,30 @@ var _ = Describe("Crypto Stream", func() {
 		It("handles out-of-order CRYPTO frames", func() {
 			Expect(str.HandleCryptoFrame(&wire.CryptoFrame{Offset: 3, Data: []byte("bar")})).To(Succeed())
 			Expect(str.HandleCryptoFrame(&wire.CryptoFrame{Data: []byte("foo")})).To(Succeed())
-			Expect(str.GetCryptoData()).To(Equal([]byte("foobar")))
+			var data []byte
+			for {
+				b := str.GetCryptoData()
+				if b == nil {
+					break
+				}
+				data = append(data, b...)
+			}
+			Expect(data).To(Equal([]byte("foobar")))
 			Expect(str.GetCryptoData()).To(BeNil())
 		})
 
 		Context("finishing", func() {
-			It("errors if there's still data to read after finishing", func() {
+			It("errors if there's still data to read at the current offset after finishing", func() {
+				Expect(str.HandleCryptoFrame(&wire.CryptoFrame{
+					Data: []byte("foo"),
+				})).To(Succeed())
+				Expect(str.Finish()).To(MatchError(&qerr.TransportError{
+					ErrorCode:    qerr.ProtocolViolation,
+					ErrorMessage: "encryption level changed, but crypto stream has more data to read",
+				}))
+			})
+
+			It("errors if there's still data to read at a higher offset after finishing", func() {
 				Expect(str.HandleCryptoFrame(&wire.CryptoFrame{
 					Data:   []byte("foobar"),
 					Offset: 10,
@@ -67,6 +85,8 @@ var _ = Describe("Crypto Stream", func() {
 				}
 				Expect(str.HandleCryptoFrame(f2)).To(Succeed())
 				Expect(str.HandleCryptoFrame(f1)).To(Succeed())
+				Expect(str.GetCryptoData()).To(HaveLen(3))
+				Expect(str.GetCryptoData()).To(HaveLen(3))
 				Expect(str.Finish()).To(Succeed())
 				Expect(str.HandleCryptoFrame(f2)).To(Succeed())
 			})