From 27acc9f9f0836394403176ed160c9dbde48cef49 Mon Sep 17 00:00:00 2001 From: Your Name Date: Sun, 25 May 2025 15:05:01 +1200 Subject: [PATCH] Bug fixing --- README.md | 10 +++---- src/config.cpp | 11 ++++++++ src/config.hpp | 3 +++ src/rate_limiter.hpp | 62 ++++++++++++++++++++++++++++++++++++++++++++ src/server.cpp | 26 ++++++++++++++++--- src/server.hpp | 2 ++ test.sh | 20 ++++++++++++++ 7 files changed, 125 insertions(+), 9 deletions(-) create mode 100644 src/rate_limiter.hpp diff --git a/README.md b/README.md index ca5f114..f35c8cb 100644 --- a/README.md +++ b/README.md @@ -84,17 +84,17 @@ The server can be configured by creating a JSON configuration file at `~/.config "port": 8080, "storage_path": "/path/to/storage", "write_tokens": ["your-secret-token"], -} -``` - -Optionally, you can modify the CORS configuration. Defaults: -```json "cors": { "allowed_origins": ["*"], "allowed_methods": ["GET", "PUT", "POST", "DELETE", "OPTIONS"], "allowed_headers": ["Authorization", "Content-Type"], "allow_credentials": false + }, + "rate_limiting": { + "auth_rate_limit": 5, // Maximum number of auth attempts + "auth_window_seconds": 10 // Time window in seconds (10 seconds) } +} ``` ## API Endpoints diff --git a/src/config.cpp b/src/config.cpp index 6ddd1db..7aacf47 100644 --- a/src/config.cpp +++ b/src/config.cpp @@ -59,6 +59,17 @@ bool load_config(const std::string& config_path, ServerConfig& config) { } } + // Parse rate limiting configuration + if (j.contains("rate_limiting")) { + const auto& rate_limit = j["rate_limiting"]; + if (rate_limit.contains("auth_rate_limit")) { + config.auth_rate_limit = rate_limit["auth_rate_limit"].get(); + } + if (rate_limit.contains("auth_window_seconds")) { + config.auth_window_seconds = rate_limit["auth_window_seconds"].get(); + } + } + return true; } catch (const std::exception& e) { std::cerr << "Error parsing config file: " << e.what() << std::endl; diff --git a/src/config.hpp b/src/config.hpp index 4704437..4a2d36f 100644 --- a/src/config.hpp +++ b/src/config.hpp @@ -17,6 +17,9 @@ struct ServerConfig { std::vector allowed_methods = {"GET", "PUT", "POST", "DELETE", "OPTIONS"}; std::vector allowed_headers = {"Authorization", "Content-Type"}; bool allow_credentials = false; + // Rate limiting configuration + int auth_rate_limit = 5; // Maximum number of auth attempts + int auth_window_seconds = 10; // Time window in seconds (10 seconds) }; bool load_config(const std::string& config_path, ServerConfig& config); diff --git a/src/rate_limiter.hpp b/src/rate_limiter.hpp new file mode 100644 index 0000000..f8de2d0 --- /dev/null +++ b/src/rate_limiter.hpp @@ -0,0 +1,62 @@ +#ifndef RATE_LIMITER_HPP +#define RATE_LIMITER_HPP + +#include +#include +#include +#include +#include + +namespace simple_object_storage { + +class RateLimiter { +public: + RateLimiter(int max_requests, std::chrono::seconds window) + : max_requests_(max_requests), window_(window) {} + + bool is_allowed(const std::string& key) { + std::lock_guard lock(mutex_); + + auto now = std::chrono::steady_clock::now(); + auto& bucket = buckets_[key]; + + // Clean up old requests + while (!bucket.empty() && (now - bucket.front()) > window_) { + bucket.pop(); + } + + // Check if we're under the limit + if (bucket.size() < max_requests_) { + bucket.push(now); + return true; + } + + return false; + } + + // New method: check if over the limit WITHOUT incrementing + bool is_over_limit(const std::string& key) { + std::lock_guard lock(mutex_); + auto now = std::chrono::steady_clock::now(); + auto& bucket = buckets_[key]; + while (!bucket.empty() && (now - bucket.front()) > window_) { + bucket.pop(); + } + return bucket.size() >= max_requests_; + } + + void reset(const std::string& key) { + std::lock_guard lock(mutex_); + buckets_.erase(key); + } + +private: + int max_requests_; + std::chrono::seconds window_; + std::unordered_map> buckets_; + std::mutex mutex_; +}; + +} // namespace simple_object_storage + +#endif \ No newline at end of file diff --git a/src/server.cpp b/src/server.cpp index ff4ca7c..2235bce 100644 --- a/src/server.cpp +++ b/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 &required_params, std::map ¶ms) { + 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(*this); + + // Initialize rate limiter + auth_rate_limiter_ = std::make_unique( + config_.auth_rate_limit, + std::chrono::seconds(config_.auth_window_seconds) + ); } Server::~Server() { diff --git a/src/server.hpp b/src/server.hpp index 6f1990d..b281970 100644 --- a/src/server.hpp +++ b/src/server.hpp @@ -7,6 +7,7 @@ #include "json.hpp" #include "database.hpp" #include "config.hpp" +#include "rate_limiter.hpp" namespace simple_object_storage { @@ -43,6 +44,7 @@ private: httplib::Server server_; bool running_; std::unique_ptr put_handler_; + std::unique_ptr auth_rate_limiter_; }; } // namespace simple_object_storage \ No newline at end of file diff --git a/test.sh b/test.sh index 1655e4e..8127359 100755 --- a/test.sh +++ b/test.sh @@ -267,4 +267,24 @@ fi curl -s -H "Authorization: Bearer ${WRITE_TOKEN}" "${BASE_URL}/deleteobject?hash=${FIRST_HASH}" > /dev/null curl -s -H "Authorization: Bearer ${WRITE_TOKEN}" "${BASE_URL}/deleteobject?hash=${SECOND_HASH}" > /dev/null +# Test 3: Verify rate limiting behavior +title "Testing rate limiting behavior" + +# Use a known invalid token +INVALID_TOKEN="invalid_token" + +# Make 5 requests with an invalid token +for i in {1..5}; do + echo "Attempt $i with invalid token" + RESPONSE=$(curl -s -X PUT -H "Authorization: Bearer ${INVALID_TOKEN}" -F "file=@${SCRIPT_DIR}/${SCRIPT_NAME}" -F "metadata={\"labeltags\":[\"test:latest\"]}" "${BASE_URL}/upload") + echo "Response: ${RESPONSE}" +done + +# Now try a request with a valid token - should be rate limited +echo "Attempting request with valid token (should be rate limited)" +RESPONSE=$(curl -s -X PUT -H "Authorization: Bearer ${WRITE_TOKEN}" -F "file=@${SCRIPT_DIR}/${SCRIPT_NAME}" -F "metadata={\"labeltags\":[\"test:latest\"]}" "${BASE_URL}/upload") +if ! echo "${RESPONSE}" | jq -r '.error' | grep -q "Too many authentication attempts"; then + die "Expected rate limit error, got: ${RESPONSE}" +fi + title "ALL TESTS PASSED"