Fix ASAN error revealed by a new fuzz test.

The QUIC-LB spec allows nonce lengths of 18 bytes, although this is an extreme case. For convenience, the LoadBalancerEncoder implementation limits it to 16 bytes so that it fits in a 128-bit integer.

The fuzz test revealed that if the nonce is 16 bytes long, the code that confines the 128-bit integer to the configured length has a bit overflow (i.e. 1 << 128 is out of range).

In the refactoring to fix this, the code switches from starting the nonce at seed_ and decrementing, and instead incrementing. One test of multiple consecutive connection IDs therefore has to change.

This part of the code is not in production. Even when QUIC-LB is deployed, we will not use 16 byte nonces. Nevertheless, the OSS implementation should be as flexible as possible.

PiperOrigin-RevId: 963480631
diff --git a/quiche/quic/load_balancer/load_balancer_encoder.cc b/quiche/quic/load_balancer/load_balancer_encoder.cc
index 65a8a26..32e559b 100644
--- a/quiche/quic/load_balancer/load_balancer_encoder.cc
+++ b/quiche/quic/load_balancer/load_balancer_encoder.cc
@@ -26,11 +26,6 @@
 
 namespace {
 
-// Returns the number of nonces given a certain |nonce_len|.
-absl::uint128 NumberOfNonces(uint8_t nonce_len) {
-  return (static_cast<absl::uint128>(1) << (nonce_len * 8));
-}
-
 // Writes the |size| least significant bytes from |in| to |out| in host byte
 // order. Returns false if |out| does not have enough space.
 bool WriteUint128(const absl::uint128 in, uint8_t size, QuicDataWriter &out) {
@@ -90,10 +85,16 @@
   }
   config_ = config;
   server_id_ = server_id;
+  if (config.nonce_len() >= 16) {
+    nonce_mask_ = absl::Uint128Max();
+  } else {
+    nonce_mask_ =
+        (static_cast<absl::uint128>(1) << (config.nonce_len() * 8)) - 1;
+  }
+  seed_ = absl::MakeUint128(random_.RandUint64(), random_.RandUint64()) &
+          nonce_mask_;
+  next_nonce_ = seed_;
 
-  seed_ = absl::MakeUint128(random_.RandUint64(), random_.RandUint64()) %
-          NumberOfNonces(config.nonce_len());
-  num_nonces_left_ = NumberOfNonces(config.nonce_len());
   connection_id_lengths_[config.config_id()] = config.total_len();
   return true;
 }
@@ -104,13 +105,15 @@
   }
   config_.reset();
   server_id_.reset();
-  num_nonces_left_ = 0;
 }
 
 QuicConnectionId LoadBalancerEncoder::GenerateConnectionId() {
-  absl::Cleanup cleanup = [&] {
-    if (num_nonces_left_ == 0) {
-      DeleteConfig();
+  absl::Cleanup cleanup = [this] {
+    if (config_.has_value()) {
+      next_nonce_ = (next_nonce_ + 1) & nonce_mask_;
+      if (next_nonce_ == seed_) {
+        DeleteConfig();
+      }
     }
   };
   uint8_t config_id = config_.has_value() ? config_->config_id()
@@ -137,10 +140,8 @@
   QuicDataWriter writer(length, reinterpret_cast<char *>(result),
                         quiche::HOST_BYTE_ORDER);
   writer.WriteUInt8(first_byte);
-  absl::uint128 next_nonce =
-      (seed_ + num_nonces_left_--) % NumberOfNonces(config_->nonce_len());
   writer.WriteBytes(server_id_->data().data(), server_id_->length());
-  if (!WriteUint128(next_nonce, config_->nonce_len(), writer)) {
+  if (!WriteUint128(next_nonce_, config_->nonce_len(), writer)) {
     return QuicConnectionId();
   }
   if (!config_->IsEncrypted()) {
diff --git a/quiche/quic/load_balancer/load_balancer_encoder.h b/quiche/quic/load_balancer/load_balancer_encoder.h
index b6625c2..56f575e 100644
--- a/quiche/quic/load_balancer/load_balancer_encoder.h
+++ b/quiche/quic/load_balancer/load_balancer_encoder.h
@@ -109,10 +109,6 @@
   // on.
   virtual void DeleteConfig();
 
-  // Returns the number of additional connection IDs that can be generated with
-  // the current config, or 0 if there is no current config.
-  absl::uint128 num_nonces_left() const { return num_nonces_left_; }
-
   // Functions below are declared virtual to enable mocking.
   // Returns true if there is an active configuration.
   virtual bool IsEncoding() const { return config_.has_value(); }
@@ -157,7 +153,7 @@
   LoadBalancerEncoderVisitorInterface* const visitor_;
 
   std::optional<LoadBalancerConfig> config_;
-  absl::uint128 seed_, num_nonces_left_ = 0;
+  absl::uint128 seed_, next_nonce_, nonce_mask_;
   std::optional<LoadBalancerServerId> server_id_;
   uint8_t connection_id_lengths_[kNumLoadBalancerConfigs + 1];
 };
diff --git a/quiche/quic/load_balancer/load_balancer_encoder_test.cc b/quiche/quic/load_balancer/load_balancer_encoder_test.cc
index 14112c9..fbc84c8 100644
--- a/quiche/quic/load_balancer/load_balancer_encoder_test.cc
+++ b/quiche/quic/load_balancer/load_balancer_encoder_test.cc
@@ -28,9 +28,12 @@
 
 class LoadBalancerEncoderPeer {
  public:
-  static void SetNumNoncesLeft(LoadBalancerEncoder &encoder,
-                               uint64_t nonces_remaining) {
-    encoder.num_nonces_left_ = absl::uint128(nonces_remaining);
+  static void SetOffsetFromSeed(LoadBalancerEncoder& encoder, uint64_t offset) {
+    encoder.next_nonce_ =
+        absl::uint128(encoder.seed_ + offset) & encoder.nonce_mask_;
+  }
+  static absl::uint128 GetOffsetFromSeed(LoadBalancerEncoder& encoder) {
+    return (encoder.next_nonce_ - encoder.seed_) & encoder.nonce_mask_;
   }
 };
 
@@ -188,9 +191,8 @@
     random_.AddNextValues(kNonceHigh, kNonceLow);
     auto encoder = LoadBalancerEncoder::Create(random_, nullptr, true, 8);
     EXPECT_TRUE(encoder->UpdateConfig(test.config, test.server_id));
-    absl::uint128 nonces_left = encoder->num_nonces_left();
     EXPECT_EQ(encoder->GenerateConnectionId(), test.connection_id);
-    EXPECT_EQ(encoder->num_nonces_left(), nonces_left - 1);
+    EXPECT_EQ(LoadBalancerEncoderPeer::GetOffsetFromSeed(*encoder), 1);
   }
 }
 
@@ -274,11 +276,12 @@
   EXPECT_TRUE(
       encoder->UpdateConfig(*config, MakeServerId(kServerId, server_id_len)));
   EXPECT_EQ(visitor.num_adds(), 1u);
-  LoadBalancerEncoderPeer::SetNumNoncesLeft(*encoder, 2);
-  EXPECT_EQ(encoder->num_nonces_left(), 2);
+  LoadBalancerEncoderPeer::SetOffsetFromSeed(*encoder, UINT32_MAX - 1);
+  EXPECT_EQ(LoadBalancerEncoderPeer::GetOffsetFromSeed(*encoder),
+            UINT32_MAX - 1);
   EXPECT_EQ(encoder->GenerateConnectionId(),
-            QuicConnectionId({0x07, 0x29, 0xd8, 0xc2, 0x17, 0xce, 0x2d, 0x92}));
-  EXPECT_EQ(encoder->num_nonces_left(), 1);
+            QuicConnectionId({0x07, 0x60, 0xa7, 0x5a, 0xea, 0x67, 0x6a, 0xd6}));
+  EXPECT_EQ(LoadBalancerEncoderPeer::GetOffsetFromSeed(*encoder), UINT32_MAX);
   encoder->GenerateConnectionId();
   EXPECT_EQ(encoder->IsEncoding(), false);
   // No retire_calls except for the initial UpdateConfig.
@@ -289,7 +292,6 @@
   random_.AddNextValues(0x83, kNonceHigh);
   auto encoder = LoadBalancerEncoder::Create(random_, nullptr, false);
   ASSERT_TRUE(encoder.has_value());
-  EXPECT_EQ(encoder->num_nonces_left(), 0);
   auto connection_id = encoder->GenerateConnectionId();
   // The first byte is the config_id (0xe0) xored with (0x83 & 0x1f).
   // The remaining bytes are random, and therefore match kNonceHigh.
@@ -319,12 +321,11 @@
   auto encoder = LoadBalancerEncoder::Create(random_, &visitor, true);
   EXPECT_TRUE(encoder->UpdateConfig(*config, MakeServerId(kServerId, 3)));
   EXPECT_EQ(visitor.num_adds(), 1u);
-  absl::uint128 left = encoder->num_nonces_left();
-  EXPECT_EQ(left, (0x1ull << 32));
+  EXPECT_EQ(LoadBalancerEncoderPeer::GetOffsetFromSeed(*encoder), 0);
   EXPECT_TRUE(encoder->IsEncoding());
   EXPECT_FALSE(encoder->IsEncrypted());
   encoder->GenerateConnectionId();
-  EXPECT_EQ(encoder->num_nonces_left(), left - 1);
+  EXPECT_EQ(LoadBalancerEncoderPeer::GetOffsetFromSeed(*encoder), 1);
   EXPECT_EQ(visitor.num_deletes(), 0u);
 }
 
@@ -354,7 +355,6 @@
   EXPECT_EQ(visitor.num_deletes(), 1u);
   EXPECT_FALSE(encoder->IsEncoding());
   EXPECT_FALSE(encoder->IsEncrypted());
-  EXPECT_EQ(encoder->num_nonces_left(), 0);
 }
 
 TEST_F(LoadBalancerEncoderTest, DeleteConfigNoVisitor) {
@@ -365,7 +365,6 @@
   encoder->DeleteConfig();
   EXPECT_FALSE(encoder->IsEncoding());
   EXPECT_FALSE(encoder->IsEncrypted());
-  EXPECT_EQ(encoder->num_nonces_left(), 0);
 }
 
 TEST_F(LoadBalancerEncoderTest, MaybeReplaceConnectionIdReturnsNoChange) {
@@ -454,6 +453,20 @@
   EXPECT_EQ(encoder->ConnectionIdLength(0x09), config_0_len);
 }
 
+// Test that the encoder doesn't crash if the nonce is max length.
+TEST_F(LoadBalancerEncoderTest, NonceIs16Bytes) {
+  auto config = LoadBalancerConfig::CreateUnencrypted(0, 1, 16);
+  ASSERT_TRUE(config.has_value());
+  auto encoder = LoadBalancerEncoder::Create(random_, nullptr, true);
+  uint8_t server_id[] = {0x3e};
+  EXPECT_TRUE(encoder->UpdateConfig(*config, MakeServerId(server_id, 1)));
+  EXPECT_TRUE(encoder
+                  ->MaybeReplaceConnectionId(TestConnectionId(1),
+                                             ParsedQuicVersion::RFCv1())
+                  .has_value());
+  EXPECT_EQ(encoder->IsEncoding(), true);
+}
+
 }  // namespace
 
 }  // namespace test