Bläddra i källkod

feat(auth): per-key fixed-window RateLimiter (singleton)

Process-singleton mutex-guarded counter keyed by API-key id; fixed
60-second windows; 0 limit means unlimited; injectable clock for
deterministic unit tests (no real sleeps).
Fszontagh 1 månad sedan
förälder
incheckning
20696e2991
5 ändrade filer med 65 tillägg och 0 borttagningar
  1. 1 0
      src/CMakeLists.txt
  2. 21 0
      src/rate_limiter.cpp
  3. 20 0
      src/rate_limiter.hpp
  4. 4 0
      tests/CMakeLists.txt
  5. 19 0
      tests/test_rate_limiter.cpp

+ 1 - 0
src/CMakeLists.txt

@@ -11,6 +11,7 @@ add_library(vectorapi_core STATIC
     collection_registry.cpp
     embeddings.cpp
     embedding_cache.cpp
+    rate_limiter.cpp
     auth.cpp
     server.cpp
     handlers/meta.cpp

+ 21 - 0
src/rate_limiter.cpp

@@ -0,0 +1,21 @@
+#include "rate_limiter.hpp"
+#include <chrono>
+namespace svapi {
+RateLimiter::RateLimiter()
+    : now_([]{ return std::chrono::duration_cast<std::chrono::seconds>(
+          std::chrono::system_clock::now().time_since_epoch()).count(); }) {}
+bool RateLimiter::allow(const std::string& keyId, uint32_t limitPerMin) {
+    if (limitPerMin == 0) return true;
+    std::lock_guard<std::mutex> lk(mu_);
+    int64_t minute = now_() / 60;
+    auto& w = windows_[keyId];
+    if (w.minute != minute) { w.minute = minute; w.count = 0; }
+    if (w.count >= limitPerMin) return false;
+    ++w.count;
+    return true;
+}
+void RateLimiter::setClockForTesting(std::function<int64_t()> fn) {
+    std::lock_guard<std::mutex> lk(mu_); now_ = std::move(fn);
+}
+RateLimiter& rateLimiter() { static RateLimiter inst; return inst; }
+}

+ 20 - 0
src/rate_limiter.hpp

@@ -0,0 +1,20 @@
+#pragma once
+#include <cstdint>
+#include <functional>
+#include <mutex>
+#include <string>
+#include <unordered_map>
+namespace svapi {
+class RateLimiter {
+public:
+    RateLimiter();
+    bool allow(const std::string& keyId, uint32_t limitPerMin); // true if under limit; 0 = unlimited
+    void setClockForTesting(std::function<int64_t()> fn);
+private:
+    struct Window { int64_t minute = -1; uint32_t count = 0; };
+    std::unordered_map<std::string, Window> windows_;
+    std::function<int64_t()> now_;   // epoch seconds
+    std::mutex mu_;
+};
+RateLimiter& rateLimiter();  // process singleton
+}

+ 4 - 0
tests/CMakeLists.txt

@@ -45,6 +45,10 @@ add_executable(test_auth test_auth.cpp)
 target_link_libraries(test_auth PRIVATE vectorapi_core GTest::gtest_main)
 gtest_discover_tests(test_auth)
 
+add_executable(test_rate_limiter test_rate_limiter.cpp)
+target_link_libraries(test_rate_limiter PRIVATE vectorapi_core GTest::gtest_main)
+gtest_discover_tests(test_rate_limiter)
+
 add_executable(test_api_integration test_api_integration.cpp)
 target_link_libraries(test_api_integration PRIVATE vectorapi_core GTest::gtest_main)
 target_compile_definitions(test_api_integration PRIVATE SVAPI_API_DOCS_DIR="${CMAKE_SOURCE_DIR}/api")

+ 19 - 0
tests/test_rate_limiter.cpp

@@ -0,0 +1,19 @@
+#include <gtest/gtest.h>
+#include "rate_limiter.hpp"
+using namespace svapi;
+TEST(RateLimiter, AllowsUpToLimitThenBlocks) {
+    RateLimiter rl; int64_t t = 1000; rl.setClockForTesting([&]{ return t; });
+    EXPECT_TRUE(rl.allow("k", 2));
+    EXPECT_TRUE(rl.allow("k", 2));
+    EXPECT_FALSE(rl.allow("k", 2));      // 3rd in same minute blocked
+    t += 60;                              // next window
+    EXPECT_TRUE(rl.allow("k", 2));       // resets
+}
+TEST(RateLimiter, ZeroMeansUnlimited) {
+    RateLimiter rl; for (int i=0;i<1000;i++) EXPECT_TRUE(rl.allow("k", 0));
+}
+TEST(RateLimiter, KeysAreIndependent) {
+    RateLimiter rl; int64_t t=1000; rl.setClockForTesting([&]{ return t; });
+    EXPECT_TRUE(rl.allow("a",1)); EXPECT_FALSE(rl.allow("a",1));
+    EXPECT_TRUE(rl.allow("b",1));
+}