Bug fixing

src/server.cpp

@@ -32,38 +32,50 @@ bool Server::init_db() {
 
 bool Server::validate_write_request(const httplib::Request &req, httplib::Response &res, const std::vector<std::string> &required_params, std::map<std::string, std::string> &params)
 {
+    std::string client_ip = req.remote_addr;
+
+    // Check if the client is already over the limit (do NOT increment)
+    if (auth_rate_limiter_->is_over_limit(client_ip)) {
+        res.status = 429;
+        nlohmann::json response = {{"result", "error"}, {"error", "Too many authentication attempts. Please try again later."}};
+        res.set_content(response.dump(), "application/json");
+        return false;
+    }
+
     // Get token from Authorization header
     std::string token;
     if (req.has_header("Authorization")) {
         const auto& auth_header = req.get_header_value("Authorization");
-        // Check if it's a Bearer token
         if (auth_header.substr(0, 7) == "Bearer ") {
             token = auth_header.substr(7);
         }
     }
 
     if (token.empty()) {
+        // Only count failed attempt (increment the limiter)
+        auth_rate_limiter_->is_allowed(client_ip); // This will increment the count
         res.status = 401;
         nlohmann::json response = {{"result", "error"}, {"error", "Missing or invalid Authorization header"}};
         res.set_content(response.dump(), "application/json");
         return false;
     }
 
-    // Check if token is valid
     bool write_token_valid = std::find(config_.write_tokens.begin(), config_.write_tokens.end(), token) != config_.write_tokens.end();
     if (!write_token_valid) {
+        // Only count failed attempt (increment the limiter)
+        auth_rate_limiter_->is_allowed(client_ip); // This will increment the count
         res.status = 403;
         nlohmann::json response = {{"result", "error"}, {"error", "Invalid write token"}};
         res.set_content(response.dump(), "application/json");
         return false;
     }
 
-    // Get other parameters from query params
+    // If authentication is successful, do not increment rate limiter
+
     for (const auto& param : req.params) {
         params[param.first] = param.second;
     }
 
-    // Check for required parameters
     for (const auto& param : required_params) {
         if (!req.has_param(param)) {
             res.status = 400;
@@ -92,6 +104,12 @@ Server::Server(const ServerConfig& config)
 
     // Initialize the put handler
     put_handler_ = std::make_unique<PutHandler>(*this);
+
+    // Initialize rate limiter
+    auth_rate_limiter_ = std::make_unique<RateLimiter>(
+        config_.auth_rate_limit,
+        std::chrono::seconds(config_.auth_window_seconds)
+    );
 }
 
 Server::~Server() {