From c4de242751e7e5cdcaa8159c5fd133f0fcfd686e Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Sun, 6 Oct 2019 08:07:19 +0200
Subject: [PATCH] add a MaxDataLen function for DATAGRAM frames

---
 internal/wire/datagram_frame.go      | 18 +++++++++
 internal/wire/datagram_frame_test.go | 55 ++++++++++++++++++++++++++++
 2 files changed, 73 insertions(+)

diff --git a/internal/wire/datagram_frame.go b/internal/wire/datagram_frame.go
index 653167f2..bd5b67c0 100644
--- a/internal/wire/datagram_frame.go
+++ b/internal/wire/datagram_frame.go
@@ -57,6 +57,24 @@ func (f *DatagramFrame) Write(b *bytes.Buffer, _ protocol.VersionNumber) error {
 	return nil
 }
 
+// MaxDataLen returns the maximum data length
+func (f *DatagramFrame) MaxDataLen(maxSize protocol.ByteCount, version protocol.VersionNumber) protocol.ByteCount {
+	headerLen := protocol.ByteCount(1)
+	if f.DataLenPresent {
+		// pretend that the data size will be 1 bytes
+		// if it turns out that varint encoding the length will consume 2 bytes, we need to adjust the data length afterwards
+		headerLen++
+	}
+	if headerLen > maxSize {
+		return 0
+	}
+	maxDataLen := maxSize - headerLen
+	if f.DataLenPresent && utils.VarIntLen(uint64(maxDataLen)) != 1 {
+		maxDataLen--
+	}
+	return maxDataLen
+}
+
 // Length of a written frame
 func (f *DatagramFrame) Length(_ protocol.VersionNumber) protocol.ByteCount {
 	length := 1 + protocol.ByteCount(len(f.Data))
diff --git a/internal/wire/datagram_frame_test.go b/internal/wire/datagram_frame_test.go
index 52ab1bfc..f4d9e3f2 100644
--- a/internal/wire/datagram_frame_test.go
+++ b/internal/wire/datagram_frame_test.go
@@ -95,4 +95,59 @@ var _ = Describe("STREAM frame", func() {
 			Expect(f.Length(versionIETFFrames)).To(Equal(protocol.ByteCount(1 + 6)))
 		})
 	})
+
+	Context("max data length", func() {
+		const maxSize = 3000
+
+		It("returns a data length such that the resulting frame has the right size, if data length is not present", func() {
+			data := make([]byte, maxSize)
+			f := &DatagramFrame{}
+			b := &bytes.Buffer{}
+			for i := 1; i < 3000; i++ {
+				b.Reset()
+				f.Data = nil
+				maxDataLen := f.MaxDataLen(protocol.ByteCount(i), versionIETFFrames)
+				if maxDataLen == 0 { // 0 means that no valid STREAM frame can be written
+					// check that writing a minimal size STREAM frame (i.e. with 1 byte data) is actually larger than the desired size
+					f.Data = []byte{0}
+					Expect(f.Write(b, versionIETFFrames)).To(Succeed())
+					Expect(b.Len()).To(BeNumerically(">", i))
+					continue
+				}
+				f.Data = data[:int(maxDataLen)]
+				Expect(f.Write(b, versionIETFFrames)).To(Succeed())
+				Expect(b.Len()).To(Equal(i))
+			}
+		})
+
+		It("always returns a data length such that the resulting frame has the right size, if data length is present", func() {
+			data := make([]byte, maxSize)
+			f := &DatagramFrame{DataLenPresent: true}
+			b := &bytes.Buffer{}
+			var frameOneByteTooSmallCounter int
+			for i := 1; i < 3000; i++ {
+				b.Reset()
+				f.Data = nil
+				maxDataLen := f.MaxDataLen(protocol.ByteCount(i), versionIETFFrames)
+				if maxDataLen == 0 { // 0 means that no valid STREAM frame can be written
+					// check that writing a minimal size STREAM frame (i.e. with 1 byte data) is actually larger than the desired size
+					f.Data = []byte{0}
+					Expect(f.Write(b, versionIETFFrames)).To(Succeed())
+					Expect(b.Len()).To(BeNumerically(">", i))
+					continue
+				}
+				f.Data = data[:int(maxDataLen)]
+				Expect(f.Write(b, versionIETFFrames)).To(Succeed())
+				// There's *one* pathological case, where a data length of x can be encoded into 1 byte
+				// but a data lengths of x+1 needs 2 bytes
+				// In that case, it's impossible to create a STREAM frame of the desired size
+				if b.Len() == i-1 {
+					frameOneByteTooSmallCounter++
+					continue
+				}
+				Expect(b.Len()).To(Equal(i))
+			}
+			Expect(frameOneByteTooSmallCounter).To(Equal(1))
+		})
+	})
 })