simplify splitting of STREAM frames for retransmissions

packet_packer.go

@@ -238,19 +238,9 @@ func (p *packetPacker) PackRetransmission(packet *ackhandler.Packet) ([]*packedP
 			controlFrames = controlFrames[1:]
 		}
 
-		// temporarily increase the maxFrameSize by the (minimum) length of the DataLen field
-		// this leads to a properly sized packet in all cases, since we do all the packet length calculations with StreamFrames that have the DataLen set
-		// however, for the last STREAM frame in the packet, we can omit the DataLen, thus yielding a packet of exactly the correct size
-		// for gQUIC STREAM frames, DataLen is always 2 bytes
-		// for IETF draft style STREAM frames, the length is encoded to either 1 or 2 bytes
-		if p.version.UsesIETFFrameFormat() {
-			maxSize++
-		} else {
-			maxSize += 2
-		}
 		for len(streamFrames) > 0 && length+protocol.MinStreamFrameSize < maxSize {
-			// TODO: optimize by setting DataLenPresent = false on all but the last STREAM frame
 			frame := streamFrames[0]
+			frame.DataLenPresent = false
 			frameToAdd := frame
 
 			sf, err := frame.MaybeSplitOffFrame(maxSize-length, p.version)
@@ -262,6 +252,7 @@ func (p *packetPacker) PackRetransmission(packet *ackhandler.Packet) ([]*packedP
 			} else {
 				streamFrames = streamFrames[1:]
 			}
+			frame.DataLenPresent = true
 			length += frameToAdd.Length(p.version)
 			frames = append(frames, frameToAdd)
 		}

packet_packer_test.go

@@ -2,6 +2,8 @@ package quic
 
 import (
 	"bytes"
+	"fmt"
+	"math/rand"
 	"net"
 
 	"github.com/golang/mock/gomock"
@@ -89,6 +91,7 @@ var _ = Describe("Packet packer", func() {
 	}
 
 	BeforeEach(func() {
+		rand.Seed(GinkgoRandomSeed())
 		version := versionGQUICFrames
 		mockSender := NewMockStreamSender(mockCtrl)
 		mockSender.EXPECT().onHasStreamData(gomock.Any()).AnyTimes()
@@ -948,6 +951,45 @@ var _ = Describe("Packet packer", func() {
 			Expect(packets[0].raw).To(HaveLen(int(maxPacketSize)))
 		})
 
+		It("splits STREAM frames, if necessary", func() {
+			for i := 0; i < 100; i++ {
+				swf := &wire.StopWaitingFrame{}
+				mockAckFramer.EXPECT().GetStopWaitingFrame(true).Return(swf)
+				sf1 := &wire.StreamFrame{
+					StreamID: 42,
+					Offset:   1337,
+					Data:     bytes.Repeat([]byte{'a'}, 1+int(rand.Int31n(int32(maxPacketSize*4/5)))),
+				}
+				sf2 := &wire.StreamFrame{
+					StreamID: 2,
+					Offset:   42,
+					Data:     bytes.Repeat([]byte{'b'}, 1+int(rand.Int31n(int32(maxPacketSize*4/5)))),
+				}
+				expectedDataLen := sf1.DataLen() + sf2.DataLen()
+				fmt.Fprintf(GinkgoWriter, "STREAM frame 1: %d bytes, STREAM frame 2: %d\n", sf1.DataLen(), sf2.DataLen())
+				frames := []wire.Frame{sf1, sf2}
+				packets, err := packer.PackRetransmission(&ackhandler.Packet{
+					EncryptionLevel: protocol.EncryptionForwardSecure,
+					Frames:          frames,
+				})
+				Expect(err).ToNot(HaveOccurred())
+
+				if len(packets) > 1 {
+					Expect(packets[0].raw).To(HaveLen(int(maxPacketSize)))
+				}
+
+				var dataLen protocol.ByteCount
+				for _, p := range packets {
+					for _, f := range p.frames {
+						if sf, ok := f.(*wire.StreamFrame); ok {
+							dataLen += sf.DataLen()
+						}
+					}
+				}
+				Expect(dataLen).To(Equal(expectedDataLen))
+			}
+		})
+
 		It("packs two packets for retransmission if the original packet contained many STREAM frames", func() {
 			swf := &wire.StopWaitingFrame{}
 			mockAckFramer.EXPECT().GetStopWaitingFrame(true).Return(swf)