From c7e3b34e2d789758183ec1cd078e6ad42b9a6b21 Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Wed, 10 Aug 2016 18:52:27 +0700
Subject: [PATCH] implement a PacketNumberGenerator

ref #186
---
 packet_number_generator.go      | 60 ++++++++++++++++++++++++++++
 packet_number_generator_test.go | 71 +++++++++++++++++++++++++++++++++
 2 files changed, 131 insertions(+)
 create mode 100644 packet_number_generator.go
 create mode 100644 packet_number_generator_test.go

diff --git a/packet_number_generator.go b/packet_number_generator.go
new file mode 100644
index 00000000..183d7f0b
--- /dev/null
+++ b/packet_number_generator.go
@@ -0,0 +1,60 @@
+package quic
+
+import (
+	"crypto/rand"
+	"math"
+
+	"github.com/lucas-clemente/quic-go/protocol"
+)
+
+type packetNumberGenerator struct {
+	last          protocol.PacketNumber
+	nextToSkip    protocol.PacketNumber
+	averagePeriod protocol.PacketNumber
+}
+
+func newPacketNumberGenerator(averagePeriod protocol.PacketNumber) *packetNumberGenerator {
+	return &packetNumberGenerator{
+		averagePeriod: averagePeriod,
+	}
+}
+
+func (p *packetNumberGenerator) GetNextPacketNumber() protocol.PacketNumber {
+	p.last++
+
+	if p.last == p.nextToSkip {
+		return p.GetNextPacketNumber()
+	}
+
+	if p.last > p.nextToSkip {
+		p.generateNewSkip()
+	}
+
+	return p.last
+}
+
+func (p *packetNumberGenerator) generateNewSkip() error {
+	num, err := p.getRandomNumber()
+	if err != nil {
+		return err
+	}
+
+	skip := protocol.PacketNumber(num) * (p.averagePeriod - 1) / (math.MaxUint16 / 2)
+	// make sure that there are never two consecutive packet numbers that are skipped
+	p.nextToSkip = p.last + 2 + skip
+
+	return nil
+}
+
+// getRandomNumber() generates a cryptographically secure random number between 0 and MaxUint16 (= 65535)
+// The expectation value is 65535/2
+func (p *packetNumberGenerator) getRandomNumber() (uint16, error) {
+	b := make([]byte, 2)
+	_, err := rand.Read(b)
+	if err != nil {
+		return 0, err
+	}
+
+	num := uint16(b[0])<<8 + uint16(b[1])
+	return num, nil
+}
diff --git a/packet_number_generator_test.go b/packet_number_generator_test.go
new file mode 100644
index 00000000..ab10d297
--- /dev/null
+++ b/packet_number_generator_test.go
@@ -0,0 +1,71 @@
+package quic
+
+import (
+	"math"
+
+	"github.com/lucas-clemente/quic-go/protocol"
+	. "github.com/onsi/ginkgo"
+	. "github.com/onsi/gomega"
+)
+
+var _ = Describe("Packet Number Generator", func() {
+	var png packetNumberGenerator
+
+	BeforeEach(func() {
+		png = packetNumberGenerator{}
+	})
+
+	It("gets 1 as the first packet number", func() {
+		num := png.GetNextPacketNumber()
+		Expect(num).To(Equal(protocol.PacketNumber(1)))
+	})
+
+	It("skips a packet number", func() {
+		png.nextToSkip = 2
+		num := png.GetNextPacketNumber()
+		Expect(num).To(Equal(protocol.PacketNumber(1)))
+		num = png.GetNextPacketNumber()
+		Expect(num).To(Equal(protocol.PacketNumber(3)))
+	})
+
+	It("generates a new packet number to skip", func() {
+		png.last = 100
+		png.averagePeriod = 100
+
+		rep := 1000
+		var sum protocol.PacketNumber
+
+		for i := 0; i < rep; i++ {
+			png.generateNewSkip()
+			Expect(png.nextToSkip).ToNot(Equal(protocol.PacketNumber(101)))
+			sum += png.nextToSkip
+		}
+
+		average := sum / protocol.PacketNumber(rep)
+		Expect(average).To(BeNumerically("==", protocol.PacketNumber(200), 10))
+	})
+
+	It("uses random numbers", func() {
+		var smallest uint16 = math.MaxUint16
+		var largest uint16
+		var sum uint64
+
+		rep := 10000
+
+		for i := 0; i < rep; i++ {
+			num, err := png.getRandomNumber()
+			Expect(err).ToNot(HaveOccurred())
+			sum += uint64(num)
+			if num > largest {
+				largest = num
+			}
+			if num < smallest {
+				smallest = num
+			}
+		}
+
+		Expect(smallest).To(BeNumerically("<", 300))
+		Expect(largest).To(BeNumerically(">", math.MaxUint16-300))
+		Expect(sum / uint64(rep)).To(BeNumerically("==", uint64(math.MaxUint16/2), 1000))
+	})
+})