Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Action based rate limiting. #1917

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Action based rate limiting. #1917

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
43 changes: 43 additions & 0 deletions spec/lucky/rate_limit_spec.cr
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,43 @@
require "../spec_helper"

include ContextHelper

class RateLimitRoutes::Index < TestAction
include Lucky::RateLimit

get "/rate_limit" do
plain_text "hello"
end

private def rate_limit : NamedTuple(to: Int32, within: Time::Span)
{to: 1, within: 1.minute}
end
Comment on lines +12 to +14
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is actually a pretty slick interface. It feels consistent with some other Lucky modules.

end

describe Lucky::RateLimit do
describe "RateLimit" do
it "when request count is less than the rate limit" do
headers = HTTP::Headers.new
headers["X_FORWARDED_FOR"] = "127.0.0.1"
request = HTTP::Request.new("GET", "/rate_limit", body: "", headers: headers)
context = build_context(request)

route = RateLimitRoutes::Index.new(context, params).call
route.context.response.status.should eq(HTTP::Status::OK)
end

it "when request count is over the rate limit" do
headers = HTTP::Headers.new
headers["X_FORWARDED_FOR"] = "127.0.0.1"
request = HTTP::Request.new("GET", "/rate_limit", body: "", headers: headers)
context = build_context(request)

10.times do
RateLimitRoutes::Index.new(context, params).call
end

route = RateLimitRoutes::Index.new(context, params).call
route.context.response.status.should eq(HTTP::Status::TOO_MANY_REQUESTS)
end
end
end
4 changes: 4 additions & 0 deletions spec/spec_helper.cr
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@

Avram.configure do |settings|
settings.database_to_migrate = UnusedDatabase
end

Check failure on line 41 in spec/spec_helper.cr

View workflow job for this annotation

GitHub Actions / specs (shard.yml, 1.10.0, false) 

undefined constant LuckyCache

Check failure on line 41 in spec/spec_helper.cr

View workflow job for this annotation

GitHub Actions / specs (shard.yml, latest, false) 

undefined constant LuckyCache

Check failure on line 41 in spec/spec_helper.cr

View workflow job for this annotation

GitHub Actions / specs (shard.edge.yml, latest, true) 

undefined constant LuckyCache

Check failure on line 41 in spec/spec_helper.cr

View workflow job for this annotation

GitHub Actions / specs (shard.override.yml, nightly, true) 

undefined constant LuckyCache

Lucky::ErrorHandler.configure do |settings|
settings.show_debug_output = false
Expand All @@ -48,4 +48,8 @@
settings.enabled = true
end

LuckyCache.configure do |settings|
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not a huge fan of this but I'm not sure how best to configure the storage at the individual spec level.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is where you'd use temp_config and put the spec inside of that block. Here's an example:

lucky/spec/tasks/exec_spec.cr

Lines 16 to 20 in bddb10b

private def with_test_template(&)
Lucky::Exec.temp_config(template_path: "spec/support/exec_template.cr.template") do
yield
end
end

It could just be added to some helper method.

settings.storage = LuckyCache::MemoryStore.new
end

Habitat.raise_if_missing_settings!
41 changes: 41 additions & 0 deletions src/lucky/rate_limit.cr
View file
Open in desktop
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks so much for this !

I'd love to get a macro one liner with keeping the possibility to have a method coputing values at runtime like you proposed.

Example :

class RateLimitRoutes::Index < TestAction
  include Lucky::RateLimit
  rate_limit to: 1, within: 1.minute

  get "/rate_limit" do
    plain_text "hello"
  end
end

The dynamic setup still works as before :

class ComputedRateLimitRoutes::Index < TestAction
  include Lucky::RateLimit

  get "/computed_rate_limit" do
    plain_text "hello"
  end

  private def computed_rate_limit : NamedTuple(to: Int32, within: Time::Span)
    {to: 1, within: 1.minute}
  end
end

Here's what I got working locally :

module Lucky::RateLimit
  macro included
    before enforce_rate_limit
  end

  macro rate_limit(**tuple)
    private def computed_rate_limit : NamedTuple(to: Int32, within: Time::Span)
      {{tuple}}
    end
  end

  abstract def computed_rate_limit : NamedTuple(to: Int32, within: Time::Span)

  private def enforce_rate_limit
    cache = LuckyCache.settings.storage
    count = cache.fetch(rate_limit_key, as: Int32, expires_in: computed_rate_limit["within"]) { 0 }
    cache.write(rate_limit_key, expires_in: computed_rate_limit["within"]) { count + 1 }

    if count > computed_rate_limit["to"]
      context.response.status = HTTP::Status::TOO_MANY_REQUESTS
      context.response.headers["Retry-After"] = computed_rate_limit["within"].to_s
      plain_text("Rate limit exceeded")
    else
      continue
    end
  end

  private def rate_limit_key : String
    klass = self.class.to_s.downcase.gsub("::", ":")
    "ratelimit:#{klass}:#{rate_limit_identifier}"
  end

  private def rate_limit_identifier : Socket::Address | Nil
    request = context.request

    if x_forwarded = request.headers["X_FORWARDED_FOR"]?.try(&.split(',').first?).presence
      begin
        Socket::IPAddress.new(x_forwarded, 0)
      rescue Socket::Error
        # if the x_forwarded is not a valid ip address we fallback to request.remote_address
        request.remote_address
      end
    else
      request.remote_address
    end
  end
end

I'm sure there are still improvements macro wise but that's the idea.

Original file line number Diff line number Diff line change
@@ -0,0 +1,41 @@
module Lucky::RateLimit
macro included
before enforce_rate_limit
end

abstract def rate_limit : NamedTuple(to: Int32, within: Time::Span)

private def enforce_rate_limit
cache = LuckyCache.settings.storage
count = cache.fetch(rate_limit_key, as: Int32, expires_in: rate_limit["within"]) { 0 }
cache.write(rate_limit_key, expires_in: rate_limit["within"]) { count + 1 }

if count > rate_limit["to"]
context.response.status = HTTP::Status::TOO_MANY_REQUESTS
context.response.headers["Retry-After"] = rate_limit["within"].to_s
plain_text("Rate limit exceeded")
else
continue
end
end

private def rate_limit_key : String
klass = self.class.to_s.downcase.gsub("::", ":")
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
klass = self.class.to_s.downcase.gsub("::", ":")
klass = {{ @type.stringify.downcase.gsub(/::/, ":") }}

I'm not sure if this matters here, but if we build it at compile time, then it won't need to compute on each request.

"ratelimit:#{klass}:#{rate_limit_identifier}"
end

private def rate_limit_identifier : Socket::Address | Nil
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This code is a copy of the RemoteIP code. You can override the method in your action, but the default needs something and I think the IP makes sense. The specs pass in a HTTP::Request so I don't get context.remote_ip.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think you should have access to request.remote_ip

lucky/src/charms/request_extensions.cr

Lines 1 to 8 in bddb10b

class HTTP::Request
# This is an alternative to `remote_address`
# since that casts to `Socket::Address`, and not all
# subclasses have an `address` method to give you the value.
# ```
# request.remote_address.as?(Socket::IPAddress).try(&.address)
# ```
property remote_ip : String = ""

This is actually patched in Lucky. context doesn't have a remote_ip method, but you may be thinking of context.request_id which is something different.

Now, with that said, it brings up the point that if you don't have the RemoteIpHandler in your middleware stack, that could cause issues... I wonder if there's a way that we could say you must have that handler in your stack in order to use this module? 🤔

request = context.request

if x_forwarded = request.headers["X_FORWARDED_FOR"]?.try(&.split(',').first?).presence
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe this can pull the header name from Lucky::RemoteIpHandler.settings.ip_header_name ?

begin
Socket::IPAddress.new(x_forwarded, 0)
rescue Socket::Error
# if the x_forwarded is not a valid ip address we fallback to request.remote_address
request.remote_address
end
else
request.remote_address
end
end
end
Loading