From 3c0726e132ccf22b3434ee3e6ac4e4a7955bbe11 Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Wed, 17 Feb 2021 11:34:19 +0800
Subject: [PATCH] move the random number generator to the utils package

---
 .../ackhandler/packet_number_generator.go     | 27 +---------------
 internal/utils/rand.go                        | 29 +++++++++++++++++
 internal/utils/rand_test.go                   | 32 +++++++++++++++++++
 3 files changed, 62 insertions(+), 26 deletions(-)
 create mode 100644 internal/utils/rand.go
 create mode 100644 internal/utils/rand_test.go

diff --git a/internal/ackhandler/packet_number_generator.go b/internal/ackhandler/packet_number_generator.go
index adb39fcd..7d58650c 100644
--- a/internal/ackhandler/packet_number_generator.go
+++ b/internal/ackhandler/packet_number_generator.go
@@ -1,9 +1,6 @@
 package ackhandler
 
 import (
-	"crypto/rand"
-	"encoding/binary"
-
 	"github.com/lucas-clemente/quic-go/internal/protocol"
 	"github.com/lucas-clemente/quic-go/internal/utils"
 )
@@ -33,28 +30,6 @@ func (p *sequentialPacketNumberGenerator) Pop() protocol.PacketNumber {
 	return next
 }
 
-type rng struct {
-	buf [4]byte
-}
-
-func (r *rng) Int31() int32 {
-	rand.Read(r.buf[:])
-	return int32(binary.BigEndian.Uint32(r.buf[:]) & ^uint32(1<<31))
-}
-
-// copied from the standard library math/rand implementation of Int63n
-func (r *rng) Int31n(n int32) int32 {
-	if n&(n-1) == 0 { // n is power of two, can mask
-		return r.Int31() & (n - 1)
-	}
-	max := int32((1 << 31) - 1 - (1<<31)%uint32(n))
-	v := r.Int31()
-	for v > max {
-		v = r.Int31()
-	}
-	return v % n
-}
-
 // The skippingPacketNumberGenerator generates the packet number for the next packet
 // it randomly skips a packet number every averagePeriod packets (on average).
 // It is guaranteed to never skip two consecutive packet numbers.
@@ -65,7 +40,7 @@ type skippingPacketNumberGenerator struct {
 	next       protocol.PacketNumber
 	nextToSkip protocol.PacketNumber
 
-	rng rng
+	rng utils.Rand
 }
 
 var _ packetNumberGenerator = &skippingPacketNumberGenerator{}
diff --git a/internal/utils/rand.go b/internal/utils/rand.go
new file mode 100644
index 00000000..30069144
--- /dev/null
+++ b/internal/utils/rand.go
@@ -0,0 +1,29 @@
+package utils
+
+import (
+	"crypto/rand"
+	"encoding/binary"
+)
+
+// Rand is a wrapper around crypto/rand that adds some convenience functions known from math/rand.
+type Rand struct {
+	buf [4]byte
+}
+
+func (r *Rand) Int31() int32 {
+	rand.Read(r.buf[:])
+	return int32(binary.BigEndian.Uint32(r.buf[:]) & ^uint32(1<<31))
+}
+
+// copied from the standard library math/rand implementation of Int63n
+func (r *Rand) Int31n(n int32) int32 {
+	if n&(n-1) == 0 { // n is power of two, can mask
+		return r.Int31() & (n - 1)
+	}
+	max := int32((1 << 31) - 1 - (1<<31)%uint32(n))
+	v := r.Int31()
+	for v > max {
+		v = r.Int31()
+	}
+	return v % n
+}
diff --git a/internal/utils/rand_test.go b/internal/utils/rand_test.go
new file mode 100644
index 00000000..4e865c37
--- /dev/null
+++ b/internal/utils/rand_test.go
@@ -0,0 +1,32 @@
+package utils
+
+import (
+	. "github.com/onsi/ginkgo"
+	. "github.com/onsi/gomega"
+)
+
+var _ = Describe("Rand", func() {
+	It("generates random numbers", func() {
+		const (
+			num = 1000
+			max = 123456
+		)
+
+		var values [num]int32
+		var r Rand
+		for i := 0; i < num; i++ {
+			v := r.Int31n(max)
+			Expect(v).To(And(
+				BeNumerically(">=", 0),
+				BeNumerically("<", max),
+			))
+			values[i] = v
+		}
+
+		var sum uint64
+		for _, n := range values {
+			sum += uint64(n)
+		}
+		Expect(float64(sum) / num).To(BeNumerically("~", max/2, max/25))
+	})
+})