Implement Edge Rate Limiting

.github/workflows/ci.yml

@@ -37,7 +37,7 @@ jobs:
       - name: Install WASI SDK
         uses: konsumer/install-wasi-sdk@v1
         with:
-          version: "25"
+          version: "30"
       - name: Run clippy
         run: cargo clippy --all -- -D warnings
 
@@ -49,7 +49,7 @@ jobs:
       - name: Install WASI SDK
         uses: konsumer/install-wasi-sdk@v1
         with:
-          version: "25"
+          version: "30"
       - name: Check C++ Format
         run: /opt/wasi-sdk/bin/clang-format --dry-run --Werror include/**/*.h src/**/*.h src/**/*.cpp
 
@@ -63,17 +63,22 @@ jobs:
         uses: actions/checkout@v4
 
       - uses: actions-rust-lang/setup-rust-toolchain@v1
+        with:
+          toolchain: stable
+
       - name: Install Viceroy
         run: cargo install viceroy --locked
 
       - name: Install WASI SDK
         uses: konsumer/install-wasi-sdk@v1
         with:
-          version: "25"
+          version: "30"
 
       - uses: extractions/setup-just@v3
 
       - uses: actions-rust-lang/setup-rust-toolchain@v1
+        with:
+          target: wasm32-wasip1
 
       - name: Build WebAssembly module
         run: just

Cargo.toml

@@ -5,7 +5,7 @@ authors = ["Kat Marchán <kat.marchan@fastly.com>", "Sy Brand <sy.brand@fastly.c
 edition = "2024"
 
 [dependencies]
-cxx = { version = "1.0.158", features = ["c++17"] }
+cxx = { version = "=1.0.158", features = ["c++17"] }
 fastly = "0.11.9"
 fastly-shared = "0.11.9"
 http = "1.3.1"
@@ -16,7 +16,7 @@ esi = "0.6.1"
 quick-xml = "0.38.3"
 
 [build-dependencies]
-cxx-build = "1.0"
+cxx-build = "=1.0.158"
 
 [lib]
 crate-type = ["staticlib"]

examples/erl.cpp

@@ -0,0 +1,32 @@
+//! @example erl.cpp
+#include "fastly/sdk.h"
+
+// Note that this example will return "Welcome!" unconditionally on Viceroy; to
+// see the expected behavior, you need to run it on a Compute service.
+
+int main() {
+  auto req{fastly::Request::from_client()};
+  auto client = req.get_client_ip_addr().value();
+
+  fastly::erl::RateCounter rate_counter("mycounter");
+  fastly::erl::PenaltyBox penalty_box("mypenaltybox");
+
+  auto erl = fastly::erl::ERL(rate_counter, penalty_box);
+  auto check_result = erl.check_rate(
+      client, // Use the client IP address as the entry to check in the rate
+              // counter
+      1,      // How many requests to count this as
+      fastly::erl::RateWindow::SixtySecs, // The window to check the rate over
+      5, // The maximum allowed rate for the client over the window
+      std::chrono::minutes(
+          1)); // The duration to penalize the client if they exceed the rate
+  if (check_result.has_value() && *check_result) {
+    fastly::Response::from_body("You are blocked!").send_to_client();
+  } else {
+    std::string message = "Welcome! Your current rate is ";
+    message += std::to_string(
+        rate_counter.lookup_rate(client, fastly::erl::RateWindow::SixtySecs)
+            .value());
+    fastly::Response::from_body(message).send_to_client();
+  }
+}

include/fastly/erl.h

@@ -0,0 +1,101 @@
+#ifndef FASTLY_ERL_H
+#define FASTLY_ERL_H
+
+#include <chrono>
+#include <fastly/expected.h>
+
+namespace fastly::erl {
+/// Errors that can arise during ERL operations.
+class ERLError {
+public:
+  enum Code {
+    Unexpected,
+    Unknown,
+    InvalidArgument,
+  };
+  explicit ERLError(Code code) : code_(code) {}
+  Code code() const { return code_; }
+
+private:
+  Code code_;
+};
+
+/// A penalty box that can be used with the edge rate limiter or stand alone for
+/// adding and checking if some entry is in the data set.
+class PenaltyBox {
+public:
+  explicit PenaltyBox(std::string name) : name_(std::move(name)) {};
+  /// Add entry to a the penaltybox for the duration of ttl. Valid ttl span is
+  /// 1m to 1h.
+  tl::expected<void, ERLError> add(std::string_view entry,
+                                   std::chrono::minutes ttl);
+  /// Check if entry is in the penaltybox.
+  tl::expected<bool, ERLError> has(std::string_view entry) const;
+  std::string_view name() const { return name_; }
+
+private:
+  std::string name_;
+};
+
+/// To be used for picking the duration in a rate counter `lookup_count` call
+enum class CounterDuration {
+  TenSec = 10,
+  TwentySecs = 20,
+  ThirtySecs = 30,
+  FortySecs = 40,
+  FiftySecs = 50,
+  SixtySecs = 60
+};
+
+/// To be used for picking the window in a rate counter `lookup_rate` or a ERL
+/// `check_rate` call.
+enum class RateWindow {
+  OneSec = 1,
+  TenSecs = 10,
+  SixtySecs = 60,
+};
+
+/// A rate counter that can be used with an edge rate limiter or stand alone for
+/// counting and rate calculations
+class RateCounter {
+public:
+  explicit RateCounter(std::string name) : name_(std::move(name)) {}
+  /// Increment an entry in the ratecounter by delta.
+  tl::expected<void, ERLError> increment(std::string_view entry,
+                                         std::uint32_t delta);
+  /// Lookup the current rate for entry in the rate counter for a window.
+  tl::expected<std::uint32_t, ERLError> lookup_rate(std::string_view entry,
+                                                    RateWindow window) const;
+  /// Lookup the current count for entry in the rate counter for a duration.
+  tl::expected<std::uint32_t, ERLError>
+  lookup_count(std::string_view entry, CounterDuration duration) const;
+  std::string_view name() const { return name_; }
+
+private:
+  std::string name_;
+};
+
+class ERL {
+public:
+  ERL(RateCounter rate_counter, PenaltyBox penalty_box)
+      : rate_counter_(std::move(rate_counter)),
+        penalty_box_(std::move(penalty_box)) {}
+
+  /// Increment an entry in a rate counter and check if the client has exceeded
+  /// some average number
+  /// of requests per second (RPS) over the window. If the client is over the
+  /// rps limit for the window, add to the penaltybox for ttl. Valid ttl span is
+  /// 1m to 1h.
+  tl::expected<bool, ERLError>
+  check_rate(std::string_view entry, std::uint32_t delta, RateWindow window,
+             std::uint32_t limit, std::chrono::minutes ttl) const;
+
+  const RateCounter &rate_counter() const { return rate_counter_; }
+  const PenaltyBox &penalty_box() const { return penalty_box_; }
+
+private:
+  RateCounter rate_counter_;
+  PenaltyBox penalty_box_;
+};
+} // namespace fastly::erl
+#endif

rust-toolchain.toml

@@ -1,3 +1,5 @@
 [toolchain]
+# We have to keep this at a version that syncs up with our llvm version.
+# wask-sdk as of 30.0 is using LLVM 19, but Rust 1.88 onwards is on LLVM 21
 channel = "1.86.0"
 targets = [ "wasm32-wasip1" ]

src/cpp/erl.cpp

@@ -0,0 +1,84 @@
+#include "fastly.h"
+#include <fastly/erl.h>
+#include <iostream>
+
+namespace {
+fastly::erl::ERLError from_status(const fastly::Status &status) {
+  using Code = fastly::Status::Code;
+  switch (status.code()) {
+  case Code::InvalidArgument:
+    return fastly::erl::ERLError(fastly::erl::ERLError::Code::InvalidArgument);
+  case Code::GenericError:
+    return fastly::erl::ERLError(fastly::erl::ERLError::Code::Unexpected);
+  default:
+    return fastly::erl::ERLError(fastly::erl::ERLError::Code::Unknown);
+  }
+}
+} // namespace
+
+#define ERL_TRY(expr)                                                          \
+  do {                                                                         \
+    auto status = (expr);                                                      \
+    if (!status.is_ok()) {                                                     \
+      return tl::unexpected(from_status(status));                              \
+    }                                                                          \
+  } while (0)
+
+namespace fastly::erl {
+tl::expected<void, ERLError> PenaltyBox::add(std::string_view entry,
+                                             std::chrono::minutes ttl) {
+  // The host expects the TTL in seconds, even though it's truncated to minutes.
+  std::chrono::seconds ttl_seconds =
+      std::chrono::duration_cast<std::chrono::seconds>(ttl);
+  ERL_TRY(fastly::penaltybox_add(name_.c_str(), name_.size(), entry.data(),
+                                 entry.size(), ttl_seconds.count()));
+  return {};
+}
+tl::expected<bool, ERLError> PenaltyBox::has(std::string_view entry) const {
+  alignas(4) bool has_out;
+  ERL_TRY(fastly::penaltybox_has(name_.c_str(), name_.size(), entry.data(),
+                                 entry.size(), &has_out));
+  return has_out;
+}
+
+tl::expected<void, ERLError> RateCounter::increment(std::string_view entry,
+                                                    std::uint32_t delta) {
+  ERL_TRY(fastly::ratecounter_increment(name_.c_str(), name_.size(),
+                                        entry.data(), entry.size(), delta));
+  return {};
+}
+
+tl::expected<std::uint32_t, ERLError>
+RateCounter::lookup_rate(std::string_view entry, RateWindow window) const {
+  std::uint32_t rate_out;
+  ERL_TRY(fastly::ratecounter_lookup_rate(
+      name_.c_str(), name_.size(), entry.data(), entry.size(),
+      static_cast<std::uint32_t>(window), &rate_out));
+  return rate_out;
+}
+
+tl::expected<std::uint32_t, ERLError>
+RateCounter::lookup_count(std::string_view entry,
+                          CounterDuration duration) const {
+  std::uint32_t count_out;
+  ERL_TRY(fastly::ratecounter_lookup_count(
+      name_.c_str(), name_.size(), entry.data(), entry.size(),
+      static_cast<std::uint32_t>(duration), &count_out));
+  return count_out;
+}
+
+tl::expected<bool, ERLError>
+ERL::check_rate(std::string_view entry, std::uint32_t delta, RateWindow window,
+                std::uint32_t limit, std::chrono::minutes ttl) const {
+  // The host expects the TTL in seconds, even though it's truncated to minutes.
+  std::chrono::seconds ttl_seconds =
+      std::chrono::duration_cast<std::chrono::seconds>(ttl);
+  alignas(4) bool blocked_out;
+  ERL_TRY(fastly::check_rate(
+      rate_counter_.name().data(), rate_counter_.name().size(), entry.data(),
+      entry.size(), delta, static_cast<std::uint32_t>(window), limit,
+      penalty_box_.name().data(), penalty_box_.name().size(),
+      static_cast<std::uint32_t>(ttl_seconds.count()), &blocked_out));
+  return blocked_out;
+}
+} // namespace fastly::erl

src/cpp/fastly.h

@@ -0,0 +1,71 @@
+#ifndef FASTLY_H
+#define FASTLY_H
+
+#include <cstddef>
+#include <cstdint>
+#include <type_traits>
+
+#define WASM_IMPORT(module, name)                                              \
+  __attribute__((import_module(module), import_name(name)))
+
+namespace fastly {
+class Status {
+public:
+  enum Code : std::uint32_t {
+    Ok = 0,
+    GenericError = 1,
+    InvalidArgument = 2,
+    BadHandle = 3,
+    BufferLen = 4,
+    Unsupported = 5,
+    BadAlign = 6,
+    HttpInvalid = 7,
+    HttpUser = 8,
+    HttpIncomplete = 9,
+    OptionalNone = 10,
+    HttpHeadTooLarge = 11,
+    HttpInvalidStatus = 12,
+    LimitExceeded = 13,
+  };
+  bool is_ok() const { return code_ == Ok; }
+  explicit operator bool() const { return is_ok(); }
+  Code code() const { return code_; }
+  Status() = default;
+  Status(Code code) : code_(code) {}
+  Status(const Status &) = default;
+  Status &operator=(const Status &) = default;
+
+private:
+  Code code_;
+};
+static_assert(std::is_trivial_v<Status>, "Status must be trivial");
+
+WASM_IMPORT("fastly_erl", "check_rate")
+Status check_rate(const char *rc, size_t rc_len, const char *entry,
+                  size_t entry_len, uint32_t delta, uint32_t window,
+                  uint32_t limit, const char *pb, size_t pb_len, uint32_t ttl,
+                  bool *blocked_out);
+
+WASM_IMPORT("fastly_erl", "ratecounter_increment")
+Status ratecounter_increment(const char *rc, size_t rc_len, const char *entry,
+                             size_t entry_len, uint32_t delta);
+
+WASM_IMPORT("fastly_erl", "ratecounter_lookup_rate")
+Status ratecounter_lookup_rate(const char *rc, size_t rc_len, const char *entry,
+                               size_t entry_len, uint32_t window,
+                               uint32_t *rate_out);
+
+WASM_IMPORT("fastly_erl", "ratecounter_lookup_count")
+Status ratecounter_lookup_count(const char *rc, size_t rc_len,
+                                const char *entry, size_t entry_len,
+                                uint32_t duration, uint32_t *count_out);
+
+WASM_IMPORT("fastly_erl", "penaltybox_add")
+Status penaltybox_add(const char *pb, size_t pb_len, const char *entry,
+                      size_t entry_len, uint32_t ttl);
+
+WASM_IMPORT("fastly_erl", "penaltybox_has")
+Status penaltybox_has(const char *pb, size_t pb_len, const char *entry,
+                      size_t entry_len, bool *has_out);
+} // namespace fastly
+#endif // FASTLY_H

src/http/request.rs

@@ -43,6 +43,7 @@ pub mod request {
         }
     }
 
+    #[allow(clippy::large_enum_variant)]
     pub enum PollResult {
         Pending(PendingRequest),
         Response(Response),