reduce allocations when handling crypto data (#4615)

This commit is contained in:
Marten Seemann
2024-08-02 21:01:53 -07:00
committed by GitHub
parent b8ea5c7981
commit bb4eb8597c
2 changed files with 27 additions and 18 deletions

View File

@@ -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 {

View File

@@ -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())
})