Add plugins (#27)

plugins/wasm-cpp/extensions/key_rate_limit/BUILD

@@ -0,0 +1,59 @@
+load("@proxy_wasm_cpp_sdk//bazel/wasm:wasm.bzl", "wasm_cc_binary")
+load("//bazel:wasm.bzl", "declare_wasm_image_targets")
+
+wasm_cc_binary(
+    name = "key_rate_limit.wasm",
+    srcs = [
+        "plugin.cc",
+        "plugin.h",
+        "bucket.h",
+        "bucket.cc",
+    ],
+    deps = [
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/time",
+        "//common:json_util",
+        "@proxy_wasm_cpp_sdk//:proxy_wasm_intrinsics",
+        "//common:http_util",
+        "//common:rule_util",
+    ],
+)
+
+cc_library(
+    name = "key_rate_limit_lib",
+    srcs = [
+        "plugin.cc",
+        "bucket.cc",
+    ],
+    hdrs = [
+        "plugin.h",
+        "bucket.h",
+    ],
+    copts = ["-DNULL_PLUGIN"],
+    deps = [
+        "@com_google_absl//absl/strings",
+        "//common:json_util",
+        "@proxy_wasm_cpp_host//:lib",
+        "//common:http_util",
+        "//common:rule_util",
+    ],
+)
+
+cc_test(
+    name = "key_rate_limit_test",
+    srcs = [
+        "plugin_test.cc",
+    ],
+    copts = ["-DNULL_PLUGIN"],
+    deps = [
+        ":key_rate_limit_lib",
+        "@com_google_googletest//:gtest",
+        "@com_google_googletest//:gtest_main",
+        "@proxy_wasm_cpp_host//:lib",
+    ],
+)
+
+declare_wasm_image_targets(
+    name = "key_rate_limit",
+    wasm_file = ":key_rate_limit.wasm",
+)

plugins/wasm-cpp/extensions/key_rate_limit/bucket.cc

@@ -0,0 +1,177 @@
+// Copyright (c) 2022 Alibaba Group Holding Ltd.
+//
+// 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.
+
+#include "extensions/key_rate_limit/bucket.h"
+
+#include <string>
+#include <unordered_map>
+
+#include "absl/strings/str_cat.h"
+#include "absl/strings/str_join.h"
+
+namespace {
+
+const int maxGetTokenRetry = 20;
+
+// Key-prefix for token bucket shared data.
+std::string tokenBucketPrefix = "mse.token_bucket";
+
+// Key-prefix for token bucket last updated time.
+std::string lastRefilledPrefix = "mse.last_refilled";
+
+}  // namespace
+
+bool getToken(int rule_id, const std::string &key) {
+  WasmDataPtr token_bucket_data;
+  uint32_t cas;
+  std::string tokenBucketKey =
+      std::to_string(rule_id) + tokenBucketPrefix + key;
+  for (int i = 0; i < maxGetTokenRetry; i++) {
+    if (WasmResult::Ok !=
+        getSharedData(tokenBucketKey, &token_bucket_data, &cas)) {
+      return false;
+    }
+    uint64_t token_left =
+        *reinterpret_cast<const uint64_t *>(token_bucket_data->data());
+    if (token_left == 0) {
+      return false;
+    }
+    token_left -= 1;
+    auto res = setSharedData(
+        tokenBucketKey,
+        {reinterpret_cast<const char *>(&token_left), sizeof(token_left)}, cas);
+    if (res == WasmResult::Ok) {
+      return true;
+    }
+    if (res == WasmResult::CasMismatch) {
+      continue;
+    }
+    return false;
+  }
+
+  LOG_WARN("get token failed with cas mismatch");
+  return true;
+}
+
+void refillToken(const std::vector<std::pair<int, LimitItem>> &rules) {
+  uint32_t last_update_cas;
+  WasmDataPtr last_update_data;
+  for (const auto &rule : rules) {
+    auto id = std::to_string(rule.first);
+    std::string lastRefilledKey = id + lastRefilledPrefix + rule.second.key;
+    std::string tokenBucketKey = id + tokenBucketPrefix + rule.second.key;
+    auto result =
+        getSharedData(lastRefilledKey, &last_update_data, &last_update_cas);
+    if (result != WasmResult::Ok) {
+      LOG_WARN(
+          absl::StrCat("failed to get last update time of the local rate limit "
+                       "token bucket ",
+                       toString(result)));
+      continue;
+    }
+    uint64_t last_update =
+        *reinterpret_cast<const uint64_t *>(last_update_data->data());
+    uint64_t now = getCurrentTimeNanoseconds();
+    if (now - last_update < rule.second.refill_interval_nanosec) {
+      continue;
+    }
+    // Otherwise, try set last updated time. If updated failed because of cas
+    // mismatch, the bucket is going to be refilled by other VMs.
+    auto res = setSharedData(
+        lastRefilledKey, {reinterpret_cast<const char *>(&now), sizeof(now)},
+        last_update_cas);
+    if (res == WasmResult::CasMismatch) {
+      continue;
+    }
+    do {
+      if (WasmResult::Ok !=
+          getSharedData(tokenBucketKey, &last_update_data, &last_update_cas)) {
+        LOG_WARN("failed to get current local rate limit token bucket");
+        break;
+      }
+      uint64_t token_left =
+          *reinterpret_cast<const uint64_t *>(last_update_data->data());
+      // Refill tokens, and update bucket with cas. If update failed because of
+      // cas mismatch, retry refilling.
+      token_left += rule.second.tokens_per_refill;
+      if (token_left > rule.second.max_tokens) {
+        token_left = rule.second.max_tokens;
+      }
+      if (WasmResult::CasMismatch ==
+          setSharedData(
+              tokenBucketKey,
+              {reinterpret_cast<const char *>(&token_left), sizeof(token_left)},
+              last_update_cas)) {
+        continue;
+      }
+      break;
+    } while (true);
+  }
+}
+
+bool initializeTokenBucket(
+    const std::vector<std::pair<int, LimitItem>> &rules) {
+  uint32_t last_update_cas;
+  WasmDataPtr last_update_data;
+  uint64_t initial_value = 0;
+  for (const auto &rule : rules) {
+    auto id = std::to_string(rule.first);
+    std::string lastRefilledKey = id + lastRefilledPrefix + rule.second.key;
+    std::string tokenBucketKey = id + tokenBucketPrefix + rule.second.key;
+    auto res =
+        getSharedData(lastRefilledKey, &last_update_data, &last_update_cas);
+    if (res == WasmResult::NotFound) {
+      setSharedData(lastRefilledKey,
+                    {reinterpret_cast<const char *>(&initial_value),
+                     sizeof(initial_value)});
+      setSharedData(tokenBucketKey,
+                    {reinterpret_cast<const char *>(&rule.second.max_tokens),
+                     sizeof(uint64_t)});
+      continue;
+    }
+    // reconfigure
+    do {
+      if (WasmResult::Ok !=
+          getSharedData(lastRefilledKey, &last_update_data, &last_update_cas)) {
+        LOG_WARN("failed to get lastRefilled");
+        return false;
+      }
+      if (WasmResult::CasMismatch ==
+          setSharedData(lastRefilledKey,
+                        {reinterpret_cast<const char *>(&initial_value),
+                         sizeof(initial_value)},
+                        last_update_cas)) {
+        continue;
+      }
+      break;
+    } while (true);
+    do {
+      if (WasmResult::Ok !=
+          getSharedData(tokenBucketKey, &last_update_data, &last_update_cas)) {
+        LOG_WARN("failed to get tokenBucket");
+        return false;
+      }
+      if (WasmResult::CasMismatch ==
+          setSharedData(
+              tokenBucketKey,
+              {reinterpret_cast<const char *>(&rule.second.max_tokens),
+               sizeof(uint64_t)},
+              last_update_cas)) {
+        continue;
+      }
+      break;
+    } while (true);
+  }
+  return true;
+}

plugins/wasm-cpp/extensions/key_rate_limit/bucket.h

@@ -0,0 +1,40 @@
+/*
+ * Copyright (c) 2022 Alibaba Group Holding Ltd.
+ *
+ * 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.
+ */
+
+#pragma once
+
+#ifndef NULL_PLUGIN
+
+#include "proxy_wasm_intrinsics.h"
+
+#else
+
+#include "include/proxy-wasm/null_plugin.h"
+using namespace proxy_wasm::null_plugin;
+using proxy_wasm::WasmResult;
+
+#endif
+
+struct LimitItem {
+  std::string key;
+  uint64_t tokens_per_refill;
+  uint64_t refill_interval_nanosec;
+  uint64_t max_tokens;
+};
+
+bool getToken(int rule_id, const std::string& key);
+void refillToken(const std::vector<std::pair<int, LimitItem>>& rules);
+bool initializeTokenBucket(const std::vector<std::pair<int, LimitItem>>& rules);

plugins/wasm-cpp/extensions/key_rate_limit/plugin.cc

@@ -0,0 +1,231 @@
+// Copyright (c) 2022 Alibaba Group Holding Ltd.
+//
+// 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.
+
+#include "extensions/key_rate_limit/plugin.h"
+
+#include <array>
+#include <vector>
+
+#include "absl/strings/str_cat.h"
+#include "absl/strings/str_split.h"
+#include "common/json_util.h"
+
+using ::nlohmann::json;
+using ::Wasm::Common::JsonArrayIterate;
+using ::Wasm::Common::JsonGetField;
+using ::Wasm::Common::JsonObjectIterate;
+using ::Wasm::Common::JsonValueAs;
+
+#ifdef NULL_PLUGIN
+
+namespace proxy_wasm {
+namespace null_plugin {
+namespace key_rate_limit {
+
+PROXY_WASM_NULL_PLUGIN_REGISTRY
+
+#endif
+
+static RegisterContextFactory register_KeyRateLimit(
+    CONTEXT_FACTORY(PluginContext), ROOT_FACTORY(PluginRootContext));
+
+namespace {
+
+constexpr uint64_t second_nano = 1000 * 1000 * 1000;
+constexpr uint64_t minute_nano = 60 * second_nano;
+constexpr uint64_t hour_nano = 60 * minute_nano;
+constexpr uint64_t day_nano = 24 * hour_nano;
+
+// tooManyRequest returns a 429 response code.
+void tooManyRequest() {
+  sendLocalResponse(429, "Too many requests", "rate_limited", {});
+}
+
+}  // namespace
+
+bool PluginRootContext::parsePluginConfig(const json& configuration,
+                                          KeyRateLimitConfigRule& rule) {
+  if (!JsonArrayIterate(
+          configuration, "limit_keys", [&](const json& item) -> bool {
+            std::string key =
+                Wasm::Common::JsonGetField<std::string>(item, "key").value();
+            uint64_t qps =
+                Wasm::Common::JsonGetField<uint64_t>(item, "query_per_second")
+                    .value_or(0);
+            if (qps > 0) {
+              rule.limit_keys.emplace(key, LimitItem{
+                                               key,
+                                               qps,
+                                               second_nano,
+                                               qps,
+                                           });
+              return true;
+            }
+            uint64_t qpm =
+                Wasm::Common::JsonGetField<uint64_t>(item, "query_per_minute")
+                    .value_or(0);
+            if (qpm > 0) {
+              rule.limit_keys.emplace(key, LimitItem{
+                                               key,
+                                               qpm,
+                                               minute_nano,
+                                               qpm,
+                                           });
+              return true;
+            }
+            uint64_t qph =
+                Wasm::Common::JsonGetField<uint64_t>(item, "query_per_hour")
+                    .value_or(0);
+            if (qph > 0) {
+              rule.limit_keys.emplace(key, LimitItem{
+                                               key,
+                                               qph,
+                                               hour_nano,
+                                               qph,
+                                           });
+              return true;
+            }
+            uint64_t qpd =
+                Wasm::Common::JsonGetField<uint64_t>(item, "query_per_day")
+                    .value_or(0);
+            if (qpd > 0) {
+              rule.limit_keys.emplace(key, LimitItem{
+                                               key,
+                                               qpd,
+                                               day_nano,
+                                               qpd,
+                                           });
+              return true;
+            }
+            LOG_WARN(
+                "one of 'query_per_second', 'query_per_minute', "
+                "'query_per_hour' or 'query_per_day' must be set");
+            return false;
+          })) {
+    LOG_WARN("failed to parse configuration for limit_keys.");
+    return false;
+  }
+  if (rule.limit_keys.empty()) {
+    LOG_WARN("no limit keys found in configuration");
+    return false;
+  }
+  auto it = configuration.find("limit_by_header");
+  if (it != configuration.end()) {
+    auto limit_by_header = JsonValueAs<std::string>(it.value());
+    if (limit_by_header.second != Wasm::Common::JsonParserResultDetail::OK) {
+      LOG_WARN("cannot parse limit_by_header");
+      return false;
+    }
+    rule.limit_by_header = limit_by_header.first.value();
+  }
+  it = configuration.find("limit_by_param");
+  if (it != configuration.end()) {
+    auto limit_by_param = JsonValueAs<std::string>(it.value());
+    if (limit_by_param.second != Wasm::Common::JsonParserResultDetail::OK) {
+      LOG_WARN("cannot parse limit_by_param");
+      return false;
+    }
+    rule.limit_by_param = limit_by_param.first.value();
+  }
+  auto emptyHeader = rule.limit_by_header.empty();
+  auto emptyParam = rule.limit_by_param.empty();
+  if ((emptyHeader && emptyParam) || (!emptyHeader && !emptyParam)) {
+    LOG_WARN("only one of 'limit_by_param' and 'limit_by_header' can be set");
+    return false;
+  }
+  return true;
+}
+
+bool PluginRootContext::checkPlugin(int rule_id,
+                                    const KeyRateLimitConfigRule& config) {
+  const auto& headerKey = config.limit_by_header;
+  const auto& paramKey = config.limit_by_param;
+  std::string key;
+  if (!headerKey.empty()) {
+    GET_HEADER_VIEW(headerKey, header);
+    key = header;
+  } else {
+    // use paramKey which must not be empty
+    GET_HEADER_VIEW(":path", path);
+    const auto& params = Wasm::Common::Http::parseQueryString(path);
+    auto it = params.find(paramKey);
+    if (it != params.end()) {
+      key = it->second;
+    }
+  }
+  const auto& limit_keys = config.limit_keys;
+  if (limit_keys.find(key) == limit_keys.end()) {
+    return true;
+  }
+  if (!getToken(rule_id, key)) {
+    tooManyRequest();
+    return false;
+  }
+  return true;
+}
+
+void PluginRootContext::onTick() { refillToken(limits_); }
+
+bool PluginRootContext::onConfigure(size_t size) {
+  // Parse configuration JSON string.
+  if (size > 0 && !configure(size)) {
+    LOG_WARN("configuration has errors initialization will not continue.");
+    setInvalidConfig();
+    return true;
+  }
+  const auto& rules = getRules();
+  for (const auto& rule : rules) {
+    for (auto& keyItem : rule.second.get().limit_keys) {
+      limits_.emplace_back(rule.first, keyItem.second);
+    }
+  }
+  initializeTokenBucket(limits_);
+  proxy_set_tick_period_milliseconds(1000);
+  return true;
+}
+
+bool PluginRootContext::configure(size_t configuration_size) {
+  auto configuration_data = getBufferBytes(WasmBufferType::PluginConfiguration,
+                                           0, configuration_size);
+  // Parse configuration JSON string.
+  auto result = ::Wasm::Common::JsonParse(configuration_data->view());
+  if (!result.has_value()) {
+    LOG_WARN(absl::StrCat("cannot parse plugin configuration JSON string: ",
+                          configuration_data->view()));
+    return false;
+  }
+  if (!parseRuleConfig(result.value())) {
+    LOG_WARN(absl::StrCat("cannot parse plugin configuration JSON string: ",
+                          configuration_data->view()));
+    return false;
+  }
+  return true;
+}
+
+FilterHeadersStatus PluginContext::onRequestHeaders(uint32_t, bool) {
+  auto* rootCtx = rootContext();
+  return rootCtx->checkRuleWithId([rootCtx](auto rule_id, const auto& config) {
+    return rootCtx->checkPlugin(rule_id, config);
+  })
+             ? FilterHeadersStatus::Continue
+             : FilterHeadersStatus::StopIteration;
+}
+
+#ifdef NULL_PLUGIN
+
+}  // namespace key_rate_limit
+}  // namespace null_plugin
+}  // namespace proxy_wasm
+
+#endif

plugins/wasm-cpp/extensions/key_rate_limit/plugin.h

@@ -0,0 +1,89 @@
+/*
+ * Copyright (c) 2022 Alibaba Group Holding Ltd.
+ *
+ * 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.
+ */
+
+#include <assert.h>
+
+#include <cstdint>
+#include <string>
+#include <unordered_map>
+
+#include "bucket.h"
+#include "common/http_util.h"
+#include "common/route_rule_matcher.h"
+#define ASSERT(_X) assert(_X)
+
+#ifndef NULL_PLUGIN
+
+#include "proxy_wasm_intrinsics.h"
+
+#else
+
+#include "include/proxy-wasm/null_plugin.h"
+
+namespace proxy_wasm {
+namespace null_plugin {
+namespace key_rate_limit {
+
+#endif
+
+struct KeyRateLimitConfigRule {
+  std::unordered_map<std::string, LimitItem> limit_keys;
+  std::string limit_by_header;
+  std::string limit_by_param;
+};
+
+// PluginRootContext is the root context for all streams processed by the
+// thread. It has the same lifetime as the worker thread and acts as target for
+// interactions that outlives individual stream, e.g. timer, async calls.
+class PluginRootContext : public RootContext,
+                          public RouteRuleMatcher<KeyRateLimitConfigRule> {
+ public:
+  PluginRootContext(uint32_t id, std::string_view root_id)
+      : RootContext(id, root_id) {}
+  ~PluginRootContext() {}
+  bool onConfigure(size_t) override;
+  void onTick() override;
+  bool checkPlugin(int, const KeyRateLimitConfigRule&);
+  bool configure(size_t);
+
+ private:
+  bool parsePluginConfig(const json&, KeyRateLimitConfigRule&) override;
+
+  std::vector<std::pair<int, LimitItem>> limits_;
+
+  friend class KeyRateLimitTest_Config_Test;
+  friend class KeyRateLimitTest_RuleConfig_Test;
+};
+
+// Per-stream context.
+class PluginContext : public Context {
+ public:
+  explicit PluginContext(uint32_t id, RootContext* root) : Context(id, root) {}
+  FilterHeadersStatus onRequestHeaders(uint32_t, bool) override;
+
+ private:
+  inline PluginRootContext* rootContext() {
+    return dynamic_cast<PluginRootContext*>(this->root());
+  }
+};
+
+#ifdef NULL_PLUGIN
+
+}  // namespace key_rate_limit
+}  // namespace null_plugin
+}  // namespace proxy_wasm
+
+#endif

plugins/wasm-cpp/extensions/key_rate_limit/plugin_test.cc

@@ -0,0 +1,206 @@
+// Copyright (c) 2022 Alibaba Group Holding Ltd.
+//
+// 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.
+
+#include "extensions/key_rate_limit/plugin.h"
+
+#include "absl/strings/str_join.h"
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+#include "include/proxy-wasm/context.h"
+#include "include/proxy-wasm/null.h"
+
+namespace proxy_wasm {
+namespace null_plugin {
+namespace key_rate_limit {
+
+NullPluginRegistry* context_registry_;
+RegisterNullVmPluginFactory register_key_rate_limit_plugin(
+    "key_rate_limit", []() {
+      return std::make_unique<NullPlugin>(key_rate_limit::context_registry_);
+    });
+
+class MockContext : public proxy_wasm::ContextBase {
+ public:
+  MockContext(WasmBase* wasm) : ContextBase(wasm) {}
+
+  MOCK_METHOD(BufferInterface*, getBuffer, (WasmBufferType));
+  MOCK_METHOD(WasmResult, log, (uint32_t, std::string_view));
+  MOCK_METHOD(WasmResult, getHeaderMapValue,
+              (WasmHeaderMapType /* type */, std::string_view /* key */,
+               std::string_view* /*result */));
+  MOCK_METHOD(WasmResult, sendLocalResponse,
+              (uint32_t /* response_code */, std::string_view /* body */,
+               Pairs /* additional_headers */, uint32_t /* grpc_status */,
+               std::string_view /* details */));
+  MOCK_METHOD(WasmResult, getProperty, (std::string_view, std::string*));
+};
+
+class KeyRateLimitTest : public ::testing::Test {
+ protected:
+  KeyRateLimitTest() {
+    // Initialize test VM
+    test_vm_ = createNullVm();
+    wasm_base_ = std::make_unique<WasmBase>(
+        std::move(test_vm_), "test-vm", "", "",
+        std::unordered_map<std::string, std::string>{},
+        AllowedCapabilitiesMap{});
+    wasm_base_->load("key_rate_limit");
+    wasm_base_->initialize();
+
+    // Initialize host side context
+    mock_context_ = std::make_unique<MockContext>(wasm_base_.get());
+    current_context_ = mock_context_.get();
+
+    ON_CALL(*mock_context_, log(testing::_, testing::_))
+        .WillByDefault([](uint32_t, std::string_view m) {
+          std::cerr << m << "\n";
+          return WasmResult::Ok;
+        });
+
+    ON_CALL(*mock_context_, getHeaderMapValue(WasmHeaderMapType::RequestHeaders,
+                                              testing::_, testing::_))
+        .WillByDefault([&](WasmHeaderMapType, std::string_view header,
+                           std::string_view* result) {
+          if (header == ":authority") {
+            *result = authority_;
+          }
+          return WasmResult::Ok;
+        });
+
+    ON_CALL(*mock_context_,
+            getHeaderMapValue(WasmHeaderMapType::ResponseHeaders, testing::_,
+                              testing::_))
+        .WillByDefault([&](WasmHeaderMapType, std::string_view header,
+                           std::string_view* result) {
+          if (header == ":status") {
+            *result = status_code_;
+          }
+          return WasmResult::Ok;
+        });
+
+    ON_CALL(*mock_context_, getProperty(testing::_, testing::_))
+        .WillByDefault([&](std::string_view path, std::string* result) {
+          *result = route_name_;
+          return WasmResult::Ok;
+        });
+
+    // Initialize Wasm sandbox context
+    root_context_ = std::make_unique<PluginRootContext>(0, "");
+    context_ = std::make_unique<PluginContext>(1, root_context_.get());
+  }
+  ~KeyRateLimitTest() override {}
+
+  std::unique_ptr<WasmBase> wasm_base_;
+  std::unique_ptr<WasmVm> test_vm_;
+  std::unique_ptr<MockContext> mock_context_;
+
+  std::unique_ptr<PluginRootContext> root_context_;
+  std::unique_ptr<PluginContext> context_;
+
+  std::string authority_;
+  std::string route_name_;
+  std::string status_code_;
+};
+
+TEST_F(KeyRateLimitTest, Config) {
+  std::string configuration = R"(
+{
+    "limit_by_header": "x-api-key",
+    "limit_keys": [
+      {
+        "key": "a",
+        "query_per_second": 1
+      },
+      {
+        "key": "b",
+        "query_per_minute": 1
+      },
+      {
+        "key": "c",
+        "query_per_hour": 1
+      },
+      {
+        "key": "d",
+        "query_per_day": 1
+      }
+    ],
+    "_rules_" : [
+      {
+        "_match_route_":["test"],
+        "limit_by_param": "apikey",
+        "limit_keys": [
+          {
+            "key": "a",
+            "query_per_second": 10
+          }
+        ]
+      }
+    ]
+})";
+
+  BufferBase buffer;
+  buffer.set({configuration.data(), configuration.size()});
+
+  EXPECT_CALL(*mock_context_, getBuffer(WasmBufferType::PluginConfiguration))
+      .WillOnce([&buffer](WasmBufferType) { return &buffer; });
+  EXPECT_TRUE(root_context_->onConfigure(configuration.size()));
+  EXPECT_EQ(root_context_->limits_.size(), 5);
+  EXPECT_EQ(root_context_->limits_[0].first, 0);
+  EXPECT_EQ(root_context_->limits_[1].first, 0);
+  EXPECT_EQ(root_context_->limits_[2].first, 0);
+  EXPECT_EQ(root_context_->limits_[3].first, 0);
+  EXPECT_EQ(root_context_->limits_[4].first, 1);
+}
+
+TEST_F(KeyRateLimitTest, RuleConfig) {
+  std::string configuration = R"(
+{
+    "_rules_" : [
+      {
+        "_match_route_":["test"],
+        "limit_by_param": "apikey",
+        "limit_keys": [
+          {
+            "key": "a",
+            "query_per_second": 10
+          }
+        ]
+      },
+      {
+        "_match_route_":["abc"],
+        "limit_by_param": "apikey",
+        "limit_keys": [
+          {
+            "key": "a",
+            "query_per_second": 100
+          }
+        ]
+      }
+    ]
+})";
+
+  BufferBase buffer;
+  buffer.set({configuration.data(), configuration.size()});
+
+  EXPECT_CALL(*mock_context_, getBuffer(WasmBufferType::PluginConfiguration))
+      .WillOnce([&buffer](WasmBufferType) { return &buffer; });
+  EXPECT_TRUE(root_context_->onConfigure(configuration.size()));
+  EXPECT_EQ(root_context_->limits_.size(), 2);
+  EXPECT_EQ(root_context_->limits_[0].first, 1);
+  EXPECT_EQ(root_context_->limits_[1].first, 2);
+}
+
+}  // namespace key_rate_limit
+}  // namespace null_plugin
+}  // namespace proxy_wasm