From 6157eeffdb2528a59d568a45a7209d8447c4b883 Mon Sep 17 00:00:00 2001
From: Lucas Clemente <lclemente@google.com>
Date: Tue, 20 Jun 2017 10:38:55 +0200
Subject: [PATCH] Factor crypto frame packing into a separate function

---
 packet_packer.go | 52 +++++++++++++++++++++++++++---------------------
 1 file changed, 29 insertions(+), 23 deletions(-)

diff --git a/packet_packer.go b/packet_packer.go
index f79727ae..783c067d 100644
--- a/packet_packer.go
+++ b/packet_packer.go
@@ -93,39 +93,26 @@ func (p *packetPacker) RetransmitNonForwardSecurePacket(packet *ackhandler.Packe
 // 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(leastUnacked protocol.PacketNumber) (*packedPacket, error) {
-	isCryptoStreamFrame := p.streamFramer.HasCryptoStreamFrame()
-
-	var sealer handshake.Sealer
-	var encLevel protocol.EncryptionLevel
-
-	// TODO(#656): Only do this for the crypto stream
-	if isCryptoStreamFrame {
-		encLevel, sealer = p.cryptoSetup.GetSealerForCryptoStream()
-	} else {
-		encLevel, sealer = p.cryptoSetup.GetSealer()
+	if p.streamFramer.HasCryptoStreamFrame() {
+		return p.packCryptoPacket(leastUnacked)
 	}
 
+	encLevel, sealer := p.cryptoSetup.GetSealer()
+
 	publicHeader := p.getPublicHeader(leastUnacked, encLevel)
 	publicHeaderLength, err := publicHeader.GetLength(p.perspective)
 	if err != nil {
 		return nil, err
 	}
-
 	if p.stopWaiting != nil {
 		p.stopWaiting.PacketNumber = publicHeader.PacketNumber
 		p.stopWaiting.PacketNumberLen = publicHeader.PacketNumberLen
 	}
 
-	var payloadFrames []frames.Frame
-	if isCryptoStreamFrame {
-		maxLen := protocol.MaxFrameAndPublicHeaderSize - protocol.NonForwardSecurePacketSizeReduction - publicHeaderLength
-		payloadFrames = []frames.Frame{p.streamFramer.PopCryptoStreamFrame(maxLen)}
-	} else {
-		maxSize := protocol.MaxFrameAndPublicHeaderSize - publicHeaderLength
-		payloadFrames, err = p.composeNextPacket(maxSize, p.canSendData(encLevel))
-		if err != nil {
-			return nil, err
-		}
+	maxSize := protocol.MaxFrameAndPublicHeaderSize - publicHeaderLength
+	payloadFrames, err := p.composeNextPacket(maxSize, p.canSendData(encLevel))
+	if err != nil {
+		return nil, err
 	}
 
 	// Check if we have enough frames to send
@@ -136,14 +123,12 @@ func (p *packetPacker) PackPacket(leastUnacked protocol.PacketNumber) (*packedPa
 	if len(payloadFrames) == 1 && p.stopWaiting != nil {
 		return nil, nil
 	}
-
 	p.stopWaiting = nil
 
 	raw, err := p.writeAndSealPacket(publicHeader, payloadFrames, sealer)
 	if err != nil {
 		return nil, err
 	}
-
 	return &packedPacket{
 		number:          publicHeader.PacketNumber,
 		raw:             raw,
@@ -152,6 +137,27 @@ func (p *packetPacker) PackPacket(leastUnacked protocol.PacketNumber) (*packedPa
 	}, nil
 }
 
+func (p *packetPacker) packCryptoPacket(leastUnacked protocol.PacketNumber) (*packedPacket, error) {
+	encLevel, sealer := p.cryptoSetup.GetSealerForCryptoStream()
+	publicHeader := p.getPublicHeader(leastUnacked, encLevel)
+	publicHeaderLength, err := publicHeader.GetLength(p.perspective)
+	if err != nil {
+		return nil, err
+	}
+	maxLen := protocol.MaxFrameAndPublicHeaderSize - protocol.NonForwardSecurePacketSizeReduction - publicHeaderLength
+	frames := []frames.Frame{p.streamFramer.PopCryptoStreamFrame(maxLen)}
+	raw, err := p.writeAndSealPacket(publicHeader, frames, sealer)
+	if err != nil {
+		return nil, err
+	}
+	return &packedPacket{
+		number:          publicHeader.PacketNumber,
+		raw:             raw,
+		frames:          frames,
+		encryptionLevel: encLevel,
+	}, nil
+}
+
 func (p *packetPacker) composeNextPacket(
 	maxFrameSize protocol.ByteCount,
 	canSendStreamFrames bool,