fix: MaybePackProbePacket also use QUIC spec (#34)

Patch MaybePackProbePacket to also generate the initial packet based on the QUIC spec if set. This fixes the incorrect behavior observed on automatic retry on timeout (sending probe packet), where uquic was inccorectly sending the default frames (PADDING, CRYPTO) instead of specified frames by QUIC spec.

Signed-off-by: Gaukas Wang <i@gaukas.wang>

internal/ackhandler/u_sent_packet_handler.go

@@ -28,3 +28,72 @@ func SetInitialPacketNumberLength(h SentPacketHandler, pnLen protocol.PacketNumb
 		sph.initialPacketNumberLength = pnLen
 	}
 }
+
+// func (h *uSentPacketHandler) OnLossDetectionTimeout() error {
+// 	defer h.setLossDetectionTimer()
+// 	earliestLossTime, encLevel := h.getLossTimeAndSpace()
+// 	if !earliestLossTime.IsZero() {
+// 		if h.logger.Debug() {
+// 			h.logger.Debugf("Loss detection alarm fired in loss timer mode. Loss time: %s", earliestLossTime)
+// 		}
+// 		if h.tracer != nil && h.tracer.LossTimerExpired != nil {
+// 			h.tracer.LossTimerExpired(logging.TimerTypeACK, encLevel)
+// 		}
+// 		// Early retransmit or time loss detection
+// 		return h.detectLostPackets(time.Now(), encLevel)
+// 	}
+
+// 	// PTO
+// 	// When all outstanding are acknowledged, the alarm is canceled in
+// 	// setLossDetectionTimer. This doesn't reset the timer in the session though.
+// 	// When OnAlarm is called, we therefore need to make sure that there are
+// 	// actually packets outstanding.
+// 	if h.bytesInFlight == 0 && !h.peerCompletedAddressValidation {
+// 		h.ptoCount++
+// 		h.numProbesToSend++
+// 		if h.initialPackets != nil {
+// 			h.ptoMode = SendPTOInitial
+// 		} else if h.handshakePackets != nil {
+// 			h.ptoMode = SendPTOHandshake
+// 		} else {
+// 			return errors.New("sentPacketHandler BUG: PTO fired, but bytes_in_flight is 0 and Initial and Handshake already dropped")
+// 		}
+// 		return nil
+// 	}
+
+// 	_, encLevel, ok := h.getPTOTimeAndSpace()
+// 	if !ok {
+// 		return nil
+// 	}
+// 	if ps := h.getPacketNumberSpace(encLevel); !ps.history.HasOutstandingPackets() && !h.peerCompletedAddressValidation {
+// 		return nil
+// 	}
+// 	h.ptoCount++
+// 	if h.logger.Debug() {
+// 		h.logger.Debugf("Loss detection alarm for %s fired in PTO mode. PTO count: %d", encLevel, h.ptoCount)
+// 	}
+// 	if h.tracer != nil {
+// 		if h.tracer.LossTimerExpired != nil {
+// 			h.tracer.LossTimerExpired(logging.TimerTypePTO, encLevel)
+// 		}
+// 		if h.tracer.UpdatedPTOCount != nil {
+// 			h.tracer.UpdatedPTOCount(h.ptoCount)
+// 		}
+// 	}
+// 	h.numProbesToSend += 2
+// 	//nolint:exhaustive // We never arm a PTO timer for 0-RTT packets.
+// 	switch encLevel {
+// 	case protocol.EncryptionInitial:
+// 		// h.ptoMode = SendPTOInitial // or quic-go will send fallback initial packets with different FRAME architecture
+// 	case protocol.EncryptionHandshake:
+// 		h.ptoMode = SendPTOHandshake
+// 	case protocol.Encryption1RTT:
+// 		// skip a packet number in order to elicit an immediate ACK
+// 		pn := h.PopPacketNumber(protocol.Encryption1RTT)
+// 		h.getPacketNumberSpace(protocol.Encryption1RTT).history.SkippedPacket(pn)
+// 		h.ptoMode = SendPTOAppData
+// 	default:
+// 		return fmt.Errorf("PTO timer in unexpected encryption level: %s", encLevel)
+// 	}
+// 	return nil
+// }

u_packet_packer.go

@@ -123,9 +123,9 @@ func (p *uPacketPacker) PackCoalescedPacket(onlyAck bool, maxPacketSize protocol
 	}
 	if initialPayload.length > 0 {
 		if onlyAck || len(initialPayload.frames) == 0 {
-			// TODO: uQUIC should send Initial Packet if requested.
+			// TODO: uQUIC should send Initial Packet ACK if requested.
 			// However, it should be otherwise configurable whether to request
-			// to send Initial Packet or not. See quic-go#4007
+			// to send Initial Packet ACK or not. See quic-go#4007
 			padding := p.initialPaddingLen(initialPayload.frames, size, maxPacketSize)
 			cont, err := p.appendLongHeaderPacket(buffer, initialHdr, initialPayload, padding, protocol.EncryptionInitial, initialSealer, v)
 			if err != nil {
@@ -256,3 +256,79 @@ func (p *uPacketPacker) MarshalInitialPacketPayload(pl payload, v protocol.Versi
 	}
 	return p.uSpec.InitialPacketSpec.FrameBuilder.Build(cryptoData)
 }
+
+func (p *uPacketPacker) MaybePackProbePacket(encLevel protocol.EncryptionLevel, maxPacketSize protocol.ByteCount, v protocol.Version) (*coalescedPacket, error) {
+	if encLevel == protocol.Encryption1RTT {
+		s, err := p.cryptoSetup.Get1RTTSealer()
+		if err != nil {
+			return nil, err
+		}
+		kp := s.KeyPhase()
+		connID := p.getDestConnID()
+		pn, pnLen := p.pnManager.PeekPacketNumber(protocol.Encryption1RTT)
+		hdrLen := wire.ShortHeaderLen(connID, pnLen)
+		pl := p.maybeGetAppDataPacket(maxPacketSize-protocol.ByteCount(s.Overhead())-hdrLen, false, true, v)
+		if pl.length == 0 {
+			return nil, nil
+		}
+		buffer := getPacketBuffer()
+		packet := &coalescedPacket{buffer: buffer}
+		shp, err := p.appendShortHeaderPacket(buffer, connID, pn, pnLen, kp, pl, 0, maxPacketSize, s, false, v)
+		if err != nil {
+			return nil, err
+		}
+		packet.shortHdrPacket = &shp
+		return packet, nil
+	}
+
+	var hdr *wire.ExtendedHeader
+	var pl payload
+	var sealer handshake.LongHeaderSealer
+	//nolint:exhaustive // Probe packets are never sent for 0-RTT.
+	switch encLevel {
+	case protocol.EncryptionInitial:
+		var err error
+		sealer, err = p.cryptoSetup.GetInitialSealer()
+		if err != nil {
+			return nil, err
+		}
+		hdr, pl = p.maybeGetCryptoPacket(maxPacketSize-protocol.ByteCount(sealer.Overhead()), protocol.EncryptionInitial, false, true, v)
+	case protocol.EncryptionHandshake:
+		var err error
+		sealer, err = p.cryptoSetup.GetHandshakeSealer()
+		if err != nil {
+			return nil, err
+		}
+		hdr, pl = p.maybeGetCryptoPacket(maxPacketSize-protocol.ByteCount(sealer.Overhead()), protocol.EncryptionHandshake, false, true, v)
+	default:
+		panic("unknown encryption level")
+	}
+
+	if pl.length == 0 {
+		return nil, nil
+	}
+	buffer := getPacketBuffer()
+	packet := &coalescedPacket{buffer: buffer}
+	size := p.longHeaderPacketLength(hdr, pl, v) + protocol.ByteCount(sealer.Overhead())
+	var padding protocol.ByteCount
+	if encLevel == protocol.EncryptionInitial {
+		if p.uSpec == nil { // default behavior
+			padding = p.initialPaddingLen(pl.frames, size, maxPacketSize)
+		} else { // otherwise we resend the spec-based initial packet
+			initPkt, err := p.appendInitialPacket(buffer, hdr, pl, protocol.EncryptionInitial, sealer, v)
+			if err != nil {
+				return nil, err
+			}
+
+			packet.longHdrPackets = []*longHeaderPacket{initPkt}
+			return packet, nil
+		}
+	}
+
+	longHdrPacket, err := p.appendLongHeaderPacket(buffer, hdr, pl, padding, encLevel, sealer, v)
+	if err != nil {
+		return nil, err
+	}
+	packet.longHdrPackets = []*longHeaderPacket{longHdrPacket}
+	return packet, nil
+}