From 0f16e08e1401797e81bbda4a4f9ade3075bec543 Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Thu, 27 Jun 2019 09:27:27 +0800
Subject: [PATCH] introduce an environment variable to set the key update
 interval

QUIC_GO_KEY_UPDATE_INTERVAL determines the maximum number of packets
that are sent or received using a key. The default value should be safe
for use in production, but setting it to a lower value is useful for
integration and interop testing.
---
 internal/handshake/updatable_aead.go      | 30 ++++++++++++++++++++++-
 internal/handshake/updatable_aead_test.go | 24 ++++++++++++++++++
 2 files changed, 53 insertions(+), 1 deletion(-)

diff --git a/internal/handshake/updatable_aead.go b/internal/handshake/updatable_aead.go
index c05e5dae..4fe8767b 100644
--- a/internal/handshake/updatable_aead.go
+++ b/internal/handshake/updatable_aead.go
@@ -4,6 +4,9 @@ import (
 	"crypto"
 	"crypto/cipher"
 	"encoding/binary"
+	"fmt"
+	"os"
+	"strconv"
 
 	"github.com/lucas-clemente/quic-go/internal/qerr"
 	"github.com/lucas-clemente/quic-go/internal/utils"
@@ -12,6 +15,31 @@ import (
 	"github.com/marten-seemann/qtls"
 )
 
+// 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 mattter 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
+}
+
 type updatableAEAD struct {
 	suite cipherSuite
 
@@ -51,7 +79,7 @@ func newUpdatableAEAD(logger utils.Logger) *updatableAEAD {
 		largestAcked:            protocol.InvalidPacketNumber,
 		firstRcvdWithCurrentKey: protocol.InvalidPacketNumber,
 		firstSentWithCurrentKey: protocol.InvalidPacketNumber,
-		keyUpdateInterval:       protocol.KeyUpdateInterval,
+		keyUpdateInterval:       keyUpdateInterval,
 		logger:                  logger,
 	}
 }
diff --git a/internal/handshake/updatable_aead_test.go b/internal/handshake/updatable_aead_test.go
index 64de2234..314d048d 100644
--- a/internal/handshake/updatable_aead_test.go
+++ b/internal/handshake/updatable_aead_test.go
@@ -5,6 +5,7 @@ import (
 	"crypto/aes"
 	"crypto/cipher"
 	"crypto/rand"
+	"os"
 
 	"github.com/lucas-clemente/quic-go/internal/protocol"
 	"github.com/lucas-clemente/quic-go/internal/utils"
@@ -207,6 +208,29 @@ var _ = Describe("Updatable AEAD", func() {
 					Expect(server.KeyPhase()).To(Equal(protocol.KeyPhaseOne))
 				})
 			})
+
+			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())
+				})
+			})
 		})
 	})
 })