fix application bug

vendor/github.com/lucas-clemente/quic-go/packet_packer.go

@@ -21,14 +21,34 @@ type packer interface {
 	PackConnectionClose(*wire.ConnectionCloseFrame) (*packedPacket, error)
 
 	HandleTransportParameters(*handshake.TransportParameters)
+	SetToken([]byte)
 	ChangeDestConnectionID(protocol.ConnectionID)
 }
 
 type packedPacket struct {
-	header          *wire.Header
-	raw             []byte
-	frames          []wire.Frame
-	encryptionLevel protocol.EncryptionLevel
+	header *wire.ExtendedHeader
+	raw    []byte
+	frames []wire.Frame
+
+	buffer *packetBuffer
+}
+
+func (p *packedPacket) EncryptionLevel() protocol.EncryptionLevel {
+	if !p.header.IsLongHeader {
+		return protocol.Encryption1RTT
+	}
+	switch p.header.Type {
+	case protocol.PacketTypeInitial:
+		return protocol.EncryptionInitial
+	case protocol.PacketTypeHandshake:
+		return protocol.EncryptionHandshake
+	default:
+		return protocol.EncryptionUnspecified
+	}
+}
+
+func (p *packedPacket) IsRetransmittable() bool {
+	return ackhandler.HasRetransmittableFrames(p.frames)
 }
 
 func (p *packedPacket) ToAckHandlerPacket() *ackhandler.Packet {
@@ -37,7 +57,7 @@ func (p *packedPacket) ToAckHandlerPacket() *ackhandler.Packet {
 		PacketType:      p.header.Type,
 		Frames:          p.frames,
 		Length:          protocol.ByteCount(len(p.raw)),
-		EncryptionLevel: p.encryptionLevel,
+		EncryptionLevel: p.EncryptionLevel(),
 		SendTime:        time.Now(),
 	}
 }
@@ -59,9 +79,13 @@ func getMaxPacketSize(addr net.Addr) protocol.ByteCount {
 	return maxSize
 }
 
+type packetNumberManager interface {
+	PeekPacketNumber(protocol.EncryptionLevel) (protocol.PacketNumber, protocol.PacketNumberLen)
+	PopPacketNumber(protocol.EncryptionLevel) protocol.PacketNumber
+}
+
 type sealingManager interface {
 	GetSealer() (protocol.EncryptionLevel, handshake.Sealer)
-	GetSealerForCryptoStream() (protocol.EncryptionLevel, handshake.Sealer)
 	GetSealerWithEncryptionLevel(protocol.EncryptionLevel) (handshake.Sealer, error)
 }
 
@@ -71,8 +95,7 @@ type frameSource interface {
 }
 
 type ackFrameSource interface {
-	GetAckFrame() *wire.AckFrame
-	GetStopWaitingFrame(bool) *wire.StopWaitingFrame
+	GetAckFrame(protocol.EncryptionLevel) *wire.AckFrame
 }
 
 type packetPacker struct {
@@ -83,16 +106,16 @@ type packetPacker struct {
 	version     protocol.VersionNumber
 	cryptoSetup sealingManager
 
+	initialStream   cryptoStream
+	handshakeStream cryptoStream
+
 	token []byte
 
-	packetNumberGenerator *packetNumberGenerator
-	getPacketNumberLen    func(protocol.PacketNumber) protocol.PacketNumberLen
-	cryptoStream          cryptoStream
-	framer                frameSource
-	acks                  ackFrameSource
+	pnManager packetNumberManager
+	framer    frameSource
+	acks      ackFrameSource
 
 	maxPacketSize             protocol.ByteCount
-	hasSentPacket             bool // has the packetPacker already sent a packet
 	numNonRetransmittableAcks int
 }
 
@@ -101,11 +124,10 @@ var _ packer = &packetPacker{}
 func newPacketPacker(
 	destConnID protocol.ConnectionID,
 	srcConnID protocol.ConnectionID,
-	initialPacketNumber protocol.PacketNumber,
-	getPacketNumberLen func(protocol.PacketNumber) protocol.PacketNumberLen,
+	initialStream cryptoStream,
+	handshakeStream cryptoStream,
+	packetNumberManager packetNumberManager,
 	remoteAddr net.Addr, // only used for determining the max packet size
-	token []byte,
-	cryptoStream cryptoStream,
 	cryptoSetup sealingManager,
 	framer frameSource,
 	acks ackFrameSource,
@@ -113,18 +135,17 @@ func newPacketPacker(
 	version protocol.VersionNumber,
 ) *packetPacker {
 	return &packetPacker{
-		cryptoStream:          cryptoStream,
-		cryptoSetup:           cryptoSetup,
-		token:                 token,
-		destConnID:            destConnID,
-		srcConnID:             srcConnID,
-		perspective:           perspective,
-		version:               version,
-		framer:                framer,
-		acks:                  acks,
-		getPacketNumberLen:    getPacketNumberLen,
-		packetNumberGenerator: newPacketNumberGenerator(initialPacketNumber, protocol.SkipPacketAveragePeriodLength),
-		maxPacketSize:         getMaxPacketSize(remoteAddr),
+		cryptoSetup:     cryptoSetup,
+		destConnID:      destConnID,
+		srcConnID:       srcConnID,
+		initialStream:   initialStream,
+		handshakeStream: handshakeStream,
+		perspective:     perspective,
+		version:         version,
+		framer:          framer,
+		acks:            acks,
+		pnManager:       packetNumberManager,
+		maxPacketSize:   getMaxPacketSize(remoteAddr),
 	}
 }
 
@@ -133,44 +154,31 @@ func (p *packetPacker) PackConnectionClose(ccf *wire.ConnectionCloseFrame) (*pac
 	frames := []wire.Frame{ccf}
 	encLevel, sealer := p.cryptoSetup.GetSealer()
 	header := p.getHeader(encLevel)
-	raw, err := p.writeAndSealPacket(header, frames, sealer)
-	return &packedPacket{
-		header:          header,
-		raw:             raw,
-		frames:          frames,
-		encryptionLevel: encLevel,
-	}, err
+	return p.writeAndSealPacket(header, frames, encLevel, sealer)
 }
 
 func (p *packetPacker) MaybePackAckPacket() (*packedPacket, error) {
-	ack := p.acks.GetAckFrame()
+	ack := p.acks.GetAckFrame(protocol.Encryption1RTT)
 	if ack == nil {
 		return nil, nil
 	}
+	// TODO(#1534): only pack ACKs with the right encryption level
 	encLevel, sealer := p.cryptoSetup.GetSealer()
 	header := p.getHeader(encLevel)
 	frames := []wire.Frame{ack}
-	raw, err := p.writeAndSealPacket(header, frames, sealer)
-	return &packedPacket{
-		header:          header,
-		raw:             raw,
-		frames:          frames,
-		encryptionLevel: encLevel,
-	}, err
+	return p.writeAndSealPacket(header, frames, encLevel, sealer)
 }
 
 // PackRetransmission packs a retransmission
 // For packets sent after completion of the handshake, it might happen that 2 packets have to be sent.
 // This can happen e.g. when a longer packet number is used in the header.
 func (p *packetPacker) PackRetransmission(packet *ackhandler.Packet) ([]*packedPacket, error) {
-	if packet.EncryptionLevel != protocol.EncryptionForwardSecure {
-		p, err := p.packHandshakeRetransmission(packet)
-		return []*packedPacket{p}, err
-	}
-
 	var controlFrames []wire.Frame
 	var streamFrames []*wire.StreamFrame
 	for _, f := range packet.Frames {
+		// CRYPTO frames are treated as control frames here.
+		// Since we're making sure that the header can never be larger for a retransmission,
+		// we never have to split CRYPTO frames.
 		if sf, ok := f.(*wire.StreamFrame); ok {
 			sf.DataLenPresent = true
 			streamFrames = append(streamFrames, sf)
@@ -180,17 +188,18 @@ func (p *packetPacker) PackRetransmission(packet *ackhandler.Packet) ([]*packedP
 	}
 
 	var packets []*packedPacket
-	encLevel, sealer := p.cryptoSetup.GetSealer()
+	encLevel := packet.EncryptionLevel
+	sealer, err := p.cryptoSetup.GetSealerWithEncryptionLevel(encLevel)
+	if err != nil {
+		return nil, err
+	}
 	for len(controlFrames) > 0 || len(streamFrames) > 0 {
 		var frames []wire.Frame
 		var length protocol.ByteCount
 
 		header := p.getHeader(encLevel)
-		headerLength, err := header.GetLength(p.version)
-		if err != nil {
-			return nil, err
-		}
-		maxSize := p.maxPacketSize - protocol.ByteCount(sealer.Overhead()) - headerLength
+		headerLen := header.GetLength(p.version)
+		maxSize := p.maxPacketSize - protocol.ByteCount(sealer.Overhead()) - headerLen
 
 		for len(controlFrames) > 0 {
 			frame := controlFrames[0]
@@ -224,41 +233,15 @@ func (p *packetPacker) PackRetransmission(packet *ackhandler.Packet) ([]*packedP
 		if sf, ok := frames[len(frames)-1].(*wire.StreamFrame); ok {
 			sf.DataLenPresent = false
 		}
-		raw, err := p.writeAndSealPacket(header, frames, sealer)
+		p, err := p.writeAndSealPacket(header, frames, encLevel, sealer)
 		if err != nil {
 			return nil, err
 		}
-		packets = append(packets, &packedPacket{
-			header:          header,
-			raw:             raw,
-			frames:          frames,
-			encryptionLevel: encLevel,
-		})
+		packets = append(packets, p)
 	}
 	return packets, nil
 }
 
-// packHandshakeRetransmission retransmits a handshake packet, that was sent with less than forward-secure encryption
-func (p *packetPacker) packHandshakeRetransmission(packet *ackhandler.Packet) (*packedPacket, error) {
-	sealer, err := p.cryptoSetup.GetSealerWithEncryptionLevel(packet.EncryptionLevel)
-	if err != nil {
-		return nil, err
-	}
-	// make sure that the retransmission for an Initial packet is sent as an Initial packet
-	if packet.PacketType == protocol.PacketTypeInitial {
-		p.hasSentPacket = false
-	}
-	header := p.getHeader(packet.EncryptionLevel)
-	header.Type = packet.PacketType
-	raw, err := p.writeAndSealPacket(header, packet.Frames, sealer)
-	return &packedPacket{
-		header:          header,
-		raw:             raw,
-		frames:          packet.Frames,
-		encryptionLevel: packet.EncryptionLevel,
-	}, err
-}
-
 // PackPacket packs a new packet
 // the other controlFrames are sent in the next packet, but might be queued and sent in the next packet if the packet would overflow MaxPacketSize otherwise
 func (p *packetPacker) PackPacket() (*packedPacket, error) {
@@ -269,21 +252,16 @@ func (p *packetPacker) PackPacket() (*packedPacket, error) {
 	if packet != nil {
 		return packet, nil
 	}
-	// if this is the first packet to be send, make sure it contains stream data
-	if !p.hasSentPacket && packet == nil {
-		return nil, nil
-	}
 
 	encLevel, sealer := p.cryptoSetup.GetSealer()
-
 	header := p.getHeader(encLevel)
-	headerLength, err := header.GetLength(p.version)
+	headerLen := header.GetLength(p.version)
 	if err != nil {
 		return nil, err
 	}
 
-	maxSize := p.maxPacketSize - protocol.ByteCount(sealer.Overhead()) - headerLength
-	frames, err := p.composeNextPacket(maxSize, p.canSendData(encLevel))
+	maxSize := p.maxPacketSize - protocol.ByteCount(sealer.Overhead()) - headerLen
+	frames, err := p.composeNextPacket(maxSize)
 	if err != nil {
 		return nil, err
 	}
@@ -304,53 +282,56 @@ func (p *packetPacker) PackPacket() (*packedPacket, error) {
 		p.numNonRetransmittableAcks = 0
 	}
 
-	raw, err := p.writeAndSealPacket(header, frames, sealer)
-	if err != nil {
-		return nil, err
-	}
-	return &packedPacket{
-		header:          header,
-		raw:             raw,
-		frames:          frames,
-		encryptionLevel: encLevel,
-	}, nil
+	return p.writeAndSealPacket(header, frames, encLevel, sealer)
 }
 
 func (p *packetPacker) maybePackCryptoPacket() (*packedPacket, error) {
-	if !p.cryptoStream.hasData() {
+	var s cryptoStream
+	var encLevel protocol.EncryptionLevel
+
+	hasData := p.initialStream.HasData()
+	ack := p.acks.GetAckFrame(protocol.EncryptionInitial)
+	if hasData || ack != nil {
+		s = p.initialStream
+		encLevel = protocol.EncryptionInitial
+	} else {
+		hasData = p.handshakeStream.HasData()
+		ack = p.acks.GetAckFrame(protocol.EncryptionHandshake)
+		if hasData || ack != nil {
+			s = p.handshakeStream
+			encLevel = protocol.EncryptionHandshake
+		}
+	}
+	if s == nil {
 		return nil, nil
 	}
-	encLevel, sealer := p.cryptoSetup.GetSealerForCryptoStream()
-	header := p.getHeader(encLevel)
-	headerLength, err := header.GetLength(p.version)
+	sealer, err := p.cryptoSetup.GetSealerWithEncryptionLevel(encLevel)
 	if err != nil {
+		// The sealer
 		return nil, err
 	}
-	maxLen := p.maxPacketSize - protocol.ByteCount(sealer.Overhead()) - protocol.NonForwardSecurePacketSizeReduction - headerLength
-	sf, _ := p.cryptoStream.popStreamFrame(maxLen)
-	sf.DataLenPresent = false
-	frames := []wire.Frame{sf}
-	raw, err := p.writeAndSealPacket(header, frames, sealer)
-	if err != nil {
-		return nil, err
+
+	hdr := p.getHeader(encLevel)
+	hdrLen := hdr.GetLength(p.version)
+	var length protocol.ByteCount
+	frames := make([]wire.Frame, 0, 2)
+	if ack != nil {
+		frames = append(frames, ack)
+		length += ack.Length(p.version)
 	}
-	return &packedPacket{
-		header:          header,
-		raw:             raw,
-		frames:          frames,
-		encryptionLevel: encLevel,
-	}, nil
+	if hasData {
+		cf := s.PopCryptoFrame(p.maxPacketSize - hdrLen - protocol.ByteCount(sealer.Overhead()) - length)
+		frames = append(frames, cf)
+	}
+	return p.writeAndSealPacket(hdr, frames, encLevel, sealer)
 }
 
-func (p *packetPacker) composeNextPacket(
-	maxFrameSize protocol.ByteCount,
-	canSendStreamFrames bool,
-) ([]wire.Frame, error) {
+func (p *packetPacker) composeNextPacket(maxFrameSize protocol.ByteCount) ([]wire.Frame, error) {
 	var length protocol.ByteCount
 	var frames []wire.Frame
 
 	// ACKs need to go first, so that the sentPacketHandler will recognize them
-	if ack := p.acks.GetAckFrame(); ack != nil {
+	if ack := p.acks.GetAckFrame(protocol.Encryption1RTT); ack != nil {
 		frames = append(frames, ack)
 		length += ack.Length(p.version)
 	}
@@ -359,10 +340,6 @@ func (p *packetPacker) composeNextPacket(
 	frames, lengthAdded = p.framer.AppendControlFrames(frames, maxFrameSize-length)
 	length += lengthAdded
 
-	if !canSendStreamFrames {
-		return frames, nil
-	}
-
 	// 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 STREAM frames 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
@@ -379,27 +356,28 @@ func (p *packetPacker) composeNextPacket(
 	return frames, nil
 }
 
-func (p *packetPacker) getHeader(encLevel protocol.EncryptionLevel) *wire.Header {
-	pnum := p.packetNumberGenerator.Peek()
-	packetNumberLen := p.getPacketNumberLen(pnum)
+func (p *packetPacker) getHeader(encLevel protocol.EncryptionLevel) *wire.ExtendedHeader {
+	pn, pnLen := p.pnManager.PeekPacketNumber(encLevel)
+	header := &wire.ExtendedHeader{}
+	header.PacketNumber = pn
+	header.PacketNumberLen = pnLen
+	header.Version = p.version
+	header.DestConnectionID = p.destConnID
 
-	header := &wire.Header{
-		PacketNumber:     pnum,
-		PacketNumberLen:  packetNumberLen,
-		Version:          p.version,
-		DestConnectionID: p.destConnID,
-	}
-
-	if encLevel != protocol.EncryptionForwardSecure {
+	if encLevel != protocol.Encryption1RTT {
 		header.IsLongHeader = true
+		// Always send Initial and Handshake packets with the maximum packet number length.
+		// This simplifies retransmissions: Since the header can't get any larger,
+		// we don't need to split CRYPTO frames.
+		header.PacketNumberLen = protocol.PacketNumberLen4
 		header.SrcConnectionID = p.srcConnID
-		// Set the payload len to maximum size.
+		// Set the length to the maximum packet size.
 		// Since it is encoded as a varint, this guarantees us that the header will end up at most as big as GetLength() returns.
-		header.PayloadLen = p.maxPacketSize
-		if !p.hasSentPacket && p.perspective == protocol.PerspectiveClient {
+		header.Length = p.maxPacketSize
+		switch encLevel {
+		case protocol.EncryptionInitial:
 			header.Type = protocol.PacketTypeInitial
-			header.Token = p.token
-		} else {
+		case protocol.EncryptionHandshake:
 			header.Type = protocol.PacketTypeHandshake
 		}
 	}
@@ -408,46 +386,63 @@ func (p *packetPacker) getHeader(encLevel protocol.EncryptionLevel) *wire.Header
 }
 
 func (p *packetPacker) writeAndSealPacket(
-	header *wire.Header,
+	header *wire.ExtendedHeader,
 	frames []wire.Frame,
+	encLevel protocol.EncryptionLevel,
 	sealer handshake.Sealer,
-) ([]byte, error) {
-	raw := *getPacketBuffer()
-	buffer := bytes.NewBuffer(raw[:0])
+) (*packedPacket, error) {
+	packetBuffer := getPacketBuffer()
+	buffer := bytes.NewBuffer(packetBuffer.Slice[:0])
+
+	addPaddingForInitial := p.perspective == protocol.PerspectiveClient && header.Type == protocol.PacketTypeInitial
 
-	// the payload length is only needed for Long Headers
 	if header.IsLongHeader {
-		if header.Type == protocol.PacketTypeInitial {
-			headerLen, _ := header.GetLength(p.version)
-			header.PayloadLen = protocol.ByteCount(protocol.MinInitialPacketSize) - headerLen
+		if p.perspective == protocol.PerspectiveClient && header.Type == protocol.PacketTypeInitial {
+			header.Token = p.token
+		}
+		if addPaddingForInitial {
+			headerLen := header.GetLength(p.version)
+			header.Length = protocol.ByteCount(header.PacketNumberLen) + protocol.MinInitialPacketSize - headerLen
 		} else {
-			payloadLen := protocol.ByteCount(sealer.Overhead())
+			// long header packets always use 4 byte packet number, so we never need to pad short payloads
+			length := protocol.ByteCount(sealer.Overhead()) + protocol.ByteCount(header.PacketNumberLen)
 			for _, frame := range frames {
-				payloadLen += frame.Length(p.version)
+				length += frame.Length(p.version)
 			}
-			header.PayloadLen = payloadLen
+			header.Length = length
 		}
 	}
 
-	if err := header.Write(buffer, p.perspective, p.version); err != nil {
+	if err := header.Write(buffer, p.version); err != nil {
 		return nil, err
 	}
-	payloadStartIndex := buffer.Len()
+	payloadOffset := buffer.Len()
 
-	// the Initial packet needs to be padded, so the last STREAM frame must have the data length present
-	if header.Type == protocol.PacketTypeInitial {
-		lastFrame := frames[len(frames)-1]
-		if sf, ok := lastFrame.(*wire.StreamFrame); ok {
-			sf.DataLenPresent = true
-		}
-	}
-	for _, frame := range frames {
+	// write all frames but the last one
+	for _, frame := range frames[:len(frames)-1] {
 		if err := frame.Write(buffer, p.version); err != nil {
 			return nil, err
 		}
 	}
-	// if this is an Initial packet, we need to pad it to fulfill the minimum size requirement
-	if header.Type == protocol.PacketTypeInitial {
+	lastFrame := frames[len(frames)-1]
+	if addPaddingForInitial {
+		// when appending padding, we need to make sure that the last STREAM frames has the data length set
+		if sf, ok := lastFrame.(*wire.StreamFrame); ok {
+			sf.DataLenPresent = true
+		}
+	} else {
+		payloadLen := buffer.Len() - payloadOffset + int(lastFrame.Length(p.version))
+		if paddingLen := 4 - int(header.PacketNumberLen) - payloadLen; paddingLen > 0 {
+			// Pad the packet such that packet number length + payload length is 4 bytes.
+			// This is needed to enable the peer to get a 16 byte sample for header protection.
+			buffer.Write(bytes.Repeat([]byte{0}, paddingLen))
+		}
+	}
+	if err := lastFrame.Write(buffer, p.version); err != nil {
+		return nil, err
+	}
+
+	if addPaddingForInitial {
 		paddingLen := protocol.MinInitialPacketSize - sealer.Overhead() - buffer.Len()
 		if paddingLen > 0 {
 			buffer.Write(bytes.Repeat([]byte{0}, paddingLen))
@@ -458,29 +453,37 @@ func (p *packetPacker) writeAndSealPacket(
 		return nil, fmt.Errorf("PacketPacker BUG: packet too large (%d bytes, allowed %d bytes)", size, p.maxPacketSize)
 	}
 
-	raw = raw[0:buffer.Len()]
-	_ = sealer.Seal(raw[payloadStartIndex:payloadStartIndex], raw[payloadStartIndex:], header.PacketNumber, raw[:payloadStartIndex])
+	raw := buffer.Bytes()
+	_ = sealer.Seal(raw[payloadOffset:payloadOffset], raw[payloadOffset:], header.PacketNumber, raw[:payloadOffset])
 	raw = raw[0 : buffer.Len()+sealer.Overhead()]
 
-	num := p.packetNumberGenerator.Pop()
+	pnOffset := payloadOffset - int(header.PacketNumberLen)
+	sealer.EncryptHeader(
+		raw[pnOffset+4:pnOffset+4+16],
+		&raw[0],
+		raw[pnOffset:payloadOffset],
+	)
+
+	num := p.pnManager.PopPacketNumber(encLevel)
 	if num != header.PacketNumber {
 		return nil, errors.New("packetPacker BUG: Peeked and Popped packet numbers do not match")
 	}
-	p.hasSentPacket = true
-	return raw, nil
-}
-
-func (p *packetPacker) canSendData(encLevel protocol.EncryptionLevel) bool {
-	if p.perspective == protocol.PerspectiveClient {
-		return encLevel >= protocol.EncryptionSecure
-	}
-	return encLevel == protocol.EncryptionForwardSecure
+	return &packedPacket{
+		header: header,
+		raw:    raw,
+		frames: frames,
+		buffer: packetBuffer,
+	}, nil
 }
 
 func (p *packetPacker) ChangeDestConnectionID(connID protocol.ConnectionID) {
 	p.destConnID = connID
 }
 
+func (p *packetPacker) SetToken(token []byte) {
+	p.token = token
+}
+
 func (p *packetPacker) HandleTransportParameters(params *handshake.TransportParameters) {
 	if params.MaxPacketSize != 0 {
 		p.maxPacketSize = utils.MinByteCount(p.maxPacketSize, params.MaxPacketSize)