From bc02d030b676dbde24538328caedacd9bdf7a874 Mon Sep 17 00:00:00 2001 From: Ted Johansson Date: Mon, 3 Apr 2023 13:46:35 +0800 Subject: [PATCH] DEV: Add configurable rate limit for Data Explorer API query runs (#238) Data Explorer can run arbitrary SQL queries which can be costly for us if over-used. Because of that we want to add the ability to rate limit the query run endpoint, in particular when requested programmatically using API. This commit introduces a rate limit to the `QueryController#run` endpoint. It heavily leans on the existing `RateLimiter` implementation, and the ability of `ApplicationController` to turn rate limit exceptions into nicely formatted JSON responses. The rate limit (per 10 seconds) can be configured through the global setting `max_data_explorer_api_reqs_per_10_seconds`, and defaults to 2. Handling can be configured through `max_data_explorer_api_req_mode`, and can be set to warn, block, or both warn and block. We will default to warn for now and monitor the logs for a while. --- .../query_controller.rb | 18 ++++ plugin.rb | 8 ++ spec/requests/query_controller_spec.rb | 91 +++++++++++++++++++ 3 files changed, 117 insertions(+) diff --git a/app/controllers/discourse_data_explorer/query_controller.rb b/app/controllers/discourse_data_explorer/query_controller.rb index 9e50869..2ecfc34 100644 --- a/app/controllers/discourse_data_explorer/query_controller.rb +++ b/app/controllers/discourse_data_explorer/query_controller.rb @@ -135,6 +135,8 @@ module ::DiscourseDataExplorer # explain - string. (Optional - pass explain=true in the request) Postgres query plan, UNIX newlines. # rows - array of array of strings. Results of the query. In the same order as 'columns'. def run + rate_limit_query_runs! + check_xhr unless params[:download] query = Query.find(params[:id].to_i) @@ -227,6 +229,22 @@ module ::DiscourseDataExplorer private + def rate_limit_query_runs! + return if !is_api? && !is_user_api? + + RateLimiter.new( + nil, + "api-query-run-10-sec", + GlobalSetting.max_data_explorer_api_reqs_per_10_seconds, + 10.seconds, + ).performed! + rescue RateLimiter::LimitExceeded => e + if GlobalSetting.max_data_explorer_api_req_mode.include?("warn") + Discourse.warn("Query run 10 second rate limit exceeded", query_id: params[:id]) + end + raise e if GlobalSetting.max_data_explorer_api_req_mode.include?("block") + end + def set_group @group = Group.find_by(name: params["group_name"]) end diff --git a/plugin.rb b/plugin.rb index 13042ca..beec18c 100644 --- a/plugin.rb +++ b/plugin.rb @@ -38,6 +38,14 @@ after_initialize do require_relative "lib/discourse_data_explorer/queries" require_relative "lib/discourse_data_explorer/query_group_bookmarkable" + GlobalSetting.add_default(:max_data_explorer_api_reqs_per_10_seconds, 2) + + # Available options: + # - warn + # - warn+block + # - block + GlobalSetting.add_default(:max_data_explorer_api_req_mode, "warn") + add_to_class(:guardian, :user_is_a_member_of_group?) do |group| return false if !current_user return true if current_user.admin? diff --git a/spec/requests/query_controller_spec.rb b/spec/requests/query_controller_spec.rb index 93e1cd9..20029b2 100644 --- a/spec/requests/query_controller_spec.rb +++ b/spec/requests/query_controller_spec.rb @@ -170,6 +170,97 @@ describe DiscourseDataExplorer::QueryController do expect(response_json["errors"].first).to match(/ValidationError/) end + context "when rate limited" do + def unlimited_request(query_id, headers = {}) + post "/admin/plugins/explorer/queries/#{query_id}/run.json", + params: { + params: {}.to_json, + }, + headers: headers + + expect(response.status).to eq(200) + end + + def limited_request(query_id, headers = {}) + post "/admin/plugins/explorer/queries/#{query_id}/run.json", + params: { + params: {}.to_json, + }, + headers: headers + + expect(response.status).to eq(429) + expect(response.parsed_body["extras"]).to eq( + { "wait_seconds" => 9, "time_left" => "9 seconds" }, + ) + end + + it "limits query runs from API when using block mode" do + global_setting :max_data_explorer_api_reqs_per_10_seconds, 1 + global_setting :max_data_explorer_api_req_mode, "block" + + RateLimiter.enable + RateLimiter.clear_all! + + admin = Fabricate(:admin) + api_key = Fabricate(:api_key, user: admin) + + query = make_query("SELECT 23 as my_value") + + headers = { HTTP_API_KEY: api_key.key, HTTP_API_USERNAME: admin.username } + + now = Time.now + freeze_time(now) + + unlimited_request(query.id, headers) + + freeze_time(now + 1.second) + + limited_request(query.id, headers) + + freeze_time(now + 10.seconds) + + unlimited_request(query.id, headers) + end + + it "does not limit query runs from API when using warn mode" do + global_setting :max_data_explorer_api_reqs_per_10_seconds, 1 + global_setting :max_data_explorer_api_req_mode, "warn" + + RateLimiter.enable + RateLimiter.clear_all! + + admin = Fabricate(:admin) + api_key = Fabricate(:api_key, user: admin) + + query = make_query("SELECT 23 as my_value") + + headers = { HTTP_API_KEY: api_key.key, HTTP_API_USERNAME: admin.username } + + freeze_time + + unlimited_request(query.id, headers) + + Discourse.expects(:warn).once + + unlimited_request(query.id, headers) + end + + it "does not limit query runs from UI" do + global_setting :max_data_explorer_api_reqs_per_10_seconds, 1 + global_setting :max_data_explorer_api_req_mode, "block" + + RateLimiter.enable + RateLimiter.clear_all! + + query = make_query("SELECT 23 as my_value") + + freeze_time + + unlimited_request(query.id) + unlimited_request(query.id) + end + end + it "doesn't allow you to modify the database #1" do p = create_post