fuzztest/fuzzing_bit_gen.h

// Copyright 2025 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef FUZZTEST_FUZZTEST_FUZZING_BIT_GEN_H_
#define FUZZTEST_FUZZTEST_FUZZING_BIT_GEN_H_

#include <cstddef>
#include <cstdint>
#include <cstring>
#include <limits>

#include "absl/base/attributes.h"
#include "absl/base/fast_type_id.h"
#include "absl/numeric/int128.h"
#include "absl/random/mocking_access.h"
#include "absl/types/span.h"

namespace fuzztest {

/// FuzzingBitGen is a BitGen instance which uses the Abseil mock mechanisms
/// to return distribution-specific variates based on the underlying control
/// and data streams. The specific sequence generated by a FuzzingBitGen may
/// vary due to the underlying code paths and whether implementation details
/// change, such as adding support for new distributions, etc.
///
/// The `control stream` is a sequence of control bytes which modifies the
/// behavior of the mocked distribution functions, such as returning min,
/// max, mean, or alternate values. When the control stream is exhausted,
/// it is not reused; it behaves as though it were zero-padded.
///
/// The `data_stream` is used to provide the random values for the mocked
/// distribution functions, as well as an internal URBG used for non-mocked
/// functions, such as std::shuffle(...).  When the data stream is exhausted,
/// it is *not* reused (wrapped around), instead the internal URBG is used as
/// the source of random variates.
///
/// The `seed_value` is used to seed the internal URBG.
///
/// NOTE: After both control and data streams are exhausted, FuzzingBitGen falls
/// back to an internal LCG PRNG, which should avoid infinite loops in fuzzed
/// functions.
///
/// This type is thread-compatible, but not thread-safe.
class FuzzingBitGen {
 public:
  // Create a FuzzingBitGen from an unowned fuzzed `data_stream` source, an
  // optional `control_stream` source, and an optional `seed_value`.
  // Both streams must outlive the FuzzingBitGen.
  explicit FuzzingBitGen(
      absl::Span<const uint8_t> data_stream ABSL_ATTRIBUTE_LIFETIME_BOUND,
      absl::Span<const uint8_t> control_stream ABSL_ATTRIBUTE_LIFETIME_BOUND,
      uint64_t seed_value);
  explicit FuzzingBitGen(
      absl::Span<const uint8_t> data_stream ABSL_ATTRIBUTE_LIFETIME_BOUND);

  FuzzingBitGen(const FuzzingBitGen&) = delete;
  FuzzingBitGen& operator=(const FuzzingBitGen&) = delete;
  FuzzingBitGen(FuzzingBitGen&&) = default;
  FuzzingBitGen& operator=(FuzzingBitGen&&) = default;

  using result_type = uint64_t;

  static constexpr result_type(min)() {
    return (std::numeric_limits<result_type>::min)();
  }

  static constexpr result_type(max)() {
    return (std::numeric_limits<result_type>::max)();
  }

  void seed(result_type seed_value = 0);

  result_type operator()();

 private:
  // Consumes up to x bytes from the data stream, or if there are no remaining
  // bytes, returns bytes from the internal LCG PRNG.
  // `result` points to a buffer that is pre-initialized as the entire buffer
  // may not be overwritten.
  void DataStreamFn(bool use_lcg, void* result, size_t result_size);

  // InvokeMock meets the requirements of absl::BitGenRef::InvokeMock.
  // This method detects whether the key has been registered as supported,
  // and, if so, returns a value derived from `data_stream_`.
  bool InvokeMock(absl::FastTypeIdType key_id, void* args_tuple, void* result);

  absl::uint128 urbg_state_ = 0;
  absl::Span<const uint8_t> control_stream_;
  absl::Span<const uint8_t> data_stream_;

  friend class ::absl::RandomMockingAccess;  // for InvokeMock
};

}  // namespace fuzztest

#endif  // FUZZTEST_FUZZTEST_FUZZING_BIT_GEN_H_