g3/rate_limiter.h

#ifndef THIRD_PARTY_MILOTIC_INTERNAL_CC_BMCWEB_RATE_LIMITER_H_
#define THIRD_PARTY_MILOTIC_INTERNAL_CC_BMCWEB_RATE_LIMITER_H_

#include <memory>
#include <optional>
#include <string>

#include "absl/base/thread_annotations.h"
#include "absl/container/flat_hash_set.h"
#include "absl/random/random.h"
#include "absl/strings/string_view.h"
#include "absl/synchronization/mutex.h"
#include "absl/time/time.h"
#include "time/clock.h"
#include "central_config.pb.h"

namespace crow {

struct RateLimiterDropInfo {
  double bucket_level = 0.0;
  double bucket_capacity = 0.0;
  double leak_rate_per_sec = 0.0;
};

namespace internal {
class LeakyBucketWithRed {
 public:
  explicit LeakyBucketWithRed(
      const milotic_tlbmc::RedfishRateLimiterModule& config,
      ecclesia::Clock* clock = ecclesia::Clock::RealClock());

  std::optional<RateLimiterDropInfo> Admit();
  void UpdateConfig(const milotic_tlbmc::RedfishRateLimiterModule& config);

 private:
  void Leak() ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_);

  absl::Mutex mu_;
  ecclesia::Clock* clock_;
  double leak_rate_ ABSL_GUARDED_BY(mu_);  // units per second
  double capacity_ ABSL_GUARDED_BY(mu_);
  double level_ ABSL_GUARDED_BY(mu_);
  absl::Time last_leak_time_ ABSL_GUARDED_BY(mu_);

  bool enable_red_ ABSL_GUARDED_BY(mu_);
  double red_min_th_ ABSL_GUARDED_BY(mu_);
  double red_max_th_ ABSL_GUARDED_BY(mu_);
  double red_max_p_ ABSL_GUARDED_BY(mu_);
  absl::BitGen bitgen_ ABSL_GUARDED_BY(mu_);
};
}  // namespace internal

// The rate limiter protects the server from DoS attacks by limiting the rate
// of incoming requests using a leaky bucket algorithm with optional Random
// Early Detection (RED). The class is thread-safe.
class RateLimiter {
 public:
  explicit RateLimiter(const milotic_tlbmc::RedfishRateLimiterModule& config,
                       ecclesia::Clock* clock = ecclesia::Clock::RealClock());

  void UpdateConfig(const milotic_tlbmc::RedfishRateLimiterModule& config);
  ~RateLimiter();

  RateLimiter(const RateLimiter&) = delete;
  RateLimiter& operator=(const RateLimiter&) = delete;
  RateLimiter(RateLimiter&&) = delete;
  RateLimiter& operator=(RateLimiter&&) = delete;

  // If request is admitted, returns std::nullopt. If request should be
  // dropped, returns RateLimiterDropInfo with bucket stats.
  std::optional<RateLimiterDropInfo> AdmitRequest(absl::string_view uri);

 private:
  std::unique_ptr<internal::LeakyBucketWithRed> bucket_;
  absl::Mutex mutex_;
  bool active_ ABSL_GUARDED_BY(mutex_);
  absl::flat_hash_set<std::string> bypass_uris_ ABSL_GUARDED_BY(mutex_);
};

}  // namespace crow

#endif  // THIRD_PARTY_MILOTIC_INTERNAL_CC_BMCWEB_RATE_LIMITER_H_