From a063500d737f871b628fc86771e03034b9ab965f Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Tue, 8 Sep 2020 11:11:16 +0700
Subject: [PATCH] use a package-level variable to set the key update frequency

---
 integrationtests/self/key_update_test.go  |  5 ++--
 internal/handshake/updatable_aead.go      | 31 +++--------------------
 internal/handshake/updatable_aead_test.go | 24 ------------------
 internal/protocol/params.go               |  2 +-
 4 files changed, 8 insertions(+), 54 deletions(-)

diff --git a/integrationtests/self/key_update_test.go b/integrationtests/self/key_update_test.go
index d90e9e328..2eec9ed02 100644
--- a/integrationtests/self/key_update_test.go
+++ b/integrationtests/self/key_update_test.go
@@ -5,9 +5,10 @@ import (
 	"fmt"
 	"io/ioutil"
 	"net"
-	"os"
 
 	quic "github.com/lucas-clemente/quic-go"
+	"github.com/lucas-clemente/quic-go/internal/handshake"
+
 	. "github.com/onsi/ginkgo"
 	. "github.com/onsi/gomega"
 )
@@ -35,7 +36,7 @@ var _ = Describe("Key Update tests", func() {
 
 	BeforeEach(func() {
 		// update keys as frequently as possible
-		os.Setenv("QUIC_GO_KEY_UPDATE_INTERVAL", "1")
+		handshake.KeyUpdateInterval = 1
 		runServer()
 	})
 
diff --git a/internal/handshake/updatable_aead.go b/internal/handshake/updatable_aead.go
index 65e0042fb..aadf0309a 100644
--- a/internal/handshake/updatable_aead.go
+++ b/internal/handshake/updatable_aead.go
@@ -6,8 +6,6 @@ import (
 	"crypto/tls"
 	"encoding/binary"
 	"fmt"
-	"os"
-	"strconv"
 	"time"
 
 	"github.com/lucas-clemente/quic-go/internal/protocol"
@@ -17,30 +15,9 @@ import (
 	"github.com/lucas-clemente/quic-go/logging"
 )
 
-// By setting this environment variable, the key update interval can be adjusted.
-// This is not needed in production, but useful for integration and interop testing.
-// Note that no matter what value is set, a key update is only initiated once it is
-// permitted (i.e. once an ACK for a packet sent at the current key phase has been received).
-const keyUpdateEnv = "QUIC_GO_KEY_UPDATE_INTERVAL"
-
-var keyUpdateInterval uint64
-
-func init() {
-	setKeyUpdateInterval()
-}
-
-func setKeyUpdateInterval() {
-	env := os.Getenv(keyUpdateEnv)
-	if env == "" {
-		keyUpdateInterval = protocol.KeyUpdateInterval
-		return
-	}
-	interval, err := strconv.ParseUint(env, 10, 64)
-	if err != nil {
-		panic(fmt.Sprintf("Cannot parse %s: %s", keyUpdateEnv, err))
-	}
-	keyUpdateInterval = interval
-}
+// KeyUpdateInterval is the maximum number of packets we send or receive before initiating a key update.
+// It's a package-level variable to allow modifying it for testing purposes.
+var KeyUpdateInterval uint64 = protocol.KeyUpdateInterval
 
 type updatableAEAD struct {
 	suite *qtls.CipherSuiteTLS13
@@ -92,7 +69,7 @@ func newUpdatableAEAD(rttStats *utils.RTTStats, tracer logging.ConnectionTracer,
 		largestAcked:            protocol.InvalidPacketNumber,
 		firstRcvdWithCurrentKey: protocol.InvalidPacketNumber,
 		firstSentWithCurrentKey: protocol.InvalidPacketNumber,
-		keyUpdateInterval:       keyUpdateInterval,
+		keyUpdateInterval:       KeyUpdateInterval,
 		rttStats:                rttStats,
 		tracer:                  tracer,
 		logger:                  logger,
diff --git a/internal/handshake/updatable_aead_test.go b/internal/handshake/updatable_aead_test.go
index 96c65bf52..ccb785ccf 100644
--- a/internal/handshake/updatable_aead_test.go
+++ b/internal/handshake/updatable_aead_test.go
@@ -4,7 +4,6 @@ import (
 	"crypto/rand"
 	"crypto/tls"
 	"fmt"
-	"os"
 	"time"
 
 	"github.com/golang/mock/gomock"
@@ -452,29 +451,6 @@ var _ = Describe("Updatable AEAD", func() {
 							Expect(err).ToNot(HaveOccurred())
 						})
 					})
-
-					Context("reading the key update env", func() {
-						AfterEach(func() {
-							os.Setenv(keyUpdateEnv, "")
-							setKeyUpdateInterval()
-						})
-
-						It("uses the default value if the env is not set", func() {
-							setKeyUpdateInterval()
-							Expect(keyUpdateInterval).To(BeEquivalentTo(protocol.KeyUpdateInterval))
-						})
-
-						It("uses the env", func() {
-							os.Setenv(keyUpdateEnv, "1337")
-							setKeyUpdateInterval()
-							Expect(keyUpdateInterval).To(BeEquivalentTo(1337))
-						})
-
-						It("panics when it can't parse the env", func() {
-							os.Setenv(keyUpdateEnv, "foobar")
-							Expect(setKeyUpdateInterval).To(Panic())
-						})
-					})
 				})
 			})
 		})
diff --git a/internal/protocol/params.go b/internal/protocol/params.go
index 1fdfab82f..4799b4b4f 100644
--- a/internal/protocol/params.go
+++ b/internal/protocol/params.go
@@ -160,7 +160,7 @@ const MaxAckDelay = 25 * time.Millisecond
 // This is the value that should be advertised to the peer.
 const MaxAckDelayInclGranularity = MaxAckDelay + TimerGranularity
 
-// KeyUpdateInterval is the maximum number of packets we send or receive before initiating a key udpate.
+// KeyUpdateInterval is the maximum number of packets we send or receive before initiating a key update.
 const KeyUpdateInterval = 100 * 1000
 
 // Max0RTTQueueingDuration is the maximum time that we store 0-RTT packets in order to wait for the corresponding Initial to be received.