From 2e0eae1a1d8fcc19fbbc9f7ef9fc0d3a4367a6ff Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Sun, 13 Nov 2016 11:30:44 +0700
Subject: [PATCH] calculate required padding size in CHLOs

---
 handshake/crypto_setup_client.go      | 15 ++++++++++++++-
 handshake/crypto_setup_client_test.go |  7 +++++++
 2 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/handshake/crypto_setup_client.go b/handshake/crypto_setup_client.go
index 1240d2e7..8015357c 100644
--- a/handshake/crypto_setup_client.go
+++ b/handshake/crypto_setup_client.go
@@ -183,6 +183,8 @@ func (h *cryptoSetupClient) sendCHLO() error {
 	b := &bytes.Buffer{}
 
 	tags := h.getTags()
+	h.addPadding(tags)
+
 	WriteHandshakeMessage(b, TagCHLO, tags)
 
 	_, err := h.cryptoStream.Write(b.Bytes())
@@ -199,7 +201,6 @@ func (h *cryptoSetupClient) getTags() map[Tag][]byte {
 	tags := make(map[Tag][]byte)
 	tags[TagSNI] = []byte("quic.clemente.io") // TODO: use real SNI here
 	tags[TagPDMD] = []byte("X509")
-	tags[TagPAD] = bytes.Repeat([]byte("0"), protocol.ClientHelloMinimumSize)
 
 	versionTag := make([]byte, 4, 4)
 	binary.LittleEndian.PutUint32(versionTag, protocol.VersionNumberToTag(h.version))
@@ -226,6 +227,18 @@ func (h *cryptoSetupClient) getTags() map[Tag][]byte {
 	return tags
 }
 
+// add a TagPAD to a tagMap, such that the total size will be bigger than the ClientHelloMinimumSize
+func (h *cryptoSetupClient) addPadding(tags map[Tag][]byte) {
+	var size int
+	for _, tag := range tags {
+		size += 8 + len(tag) // 4 bytes for the tag + 4 bytes for the offset + the length of the data
+	}
+	paddingSize := protocol.ClientHelloMinimumSize - size
+	if paddingSize > 0 {
+		tags[TagPAD] = bytes.Repeat([]byte{0}, paddingSize)
+	}
+}
+
 func (h *cryptoSetupClient) maybeUpgradeCrypto() error {
 	leafCert := h.certManager.GetLeafCert()
 
diff --git a/handshake/crypto_setup_client_test.go b/handshake/crypto_setup_client_test.go
index 20a2bb71..d722a070 100644
--- a/handshake/crypto_setup_client_test.go
+++ b/handshake/crypto_setup_client_test.go
@@ -141,6 +141,13 @@ var _ = Describe("Crypto setup", func() {
 			Expect(cs.cryptoStream.(*mockStream).dataWritten.Len()).To(BeNumerically(">", protocol.ClientHelloMinimumSize))
 		})
 
+		It("doesn't overflow the packet with padding", func() {
+			tagMap := make(map[Tag][]byte)
+			tagMap[TagSCID] = bytes.Repeat([]byte{0}, protocol.ClientHelloMinimumSize*6/10)
+			cs.addPadding(tagMap)
+			Expect(len(tagMap[TagPAD])).To(BeNumerically("<", protocol.ClientHelloMinimumSize/2))
+		})
+
 		It("saves the last sent CHLO", func() {
 			// send first CHLO
 			err := cs.sendCHLO()