How to create middleware in your Rails application

Creating your own middleware in a Rails application is not something a lot of developers do or have to do. But when you want to add custom headers, authentication or some kind of rate limiting, you should definitely consider using middleware. I want to show you how you can implement a small rate limiter using your own middleware.

Middleware in Rails

Middleware in Rails is actually middleware in Rack. Rack connects a web framework (e.g. Rails) with a web server (e.g. Puma). It is responsible for wrapping http data to present it in a single method call where you can work with requests and responses. That single method call is where you put the middleware logic. You can consider middleware to be the entry and exit of you Rails application.

Each piece of middleware receives the current environment from the previous middleware (if any) and reads and/or changes this environment to pass it on to the next middleware (again, if any).

There is actually already a lot of default middleware in a vanilla Rails application created with rails new. You can list that middleware by running ./bin/rails middleware. For a new Rails 7 application it looks like this:

use ActionDispatch::HostAuthorization
use Rack::Sendfile
use ActionDispatch::Static
use ActionDispatch::Executor
use ActionDispatch::ServerTiming
use ActiveSupport::Cache::Strategy::LocalCache::Middleware
use Rack::Runtime
use Rack::MethodOverride
use ActionDispatch::RequestId
use ActionDispatch::RemoteIp
use Sprockets::Rails::QuietAssets
use Rails::Rack::Logger
use ActionDispatch::ShowExceptions
use Sentry::Rails::CaptureExceptions
use WebConsole::Middleware
use ActionDispatch::DebugExceptions
use Sentry::Rails::RescuedExceptionInterceptor
use ActionDispatch::ActionableExceptions
use ActionDispatch::Reloader
use ActionDispatch::Callbacks
use ActiveRecord::Migration::CheckPending
use ActionDispatch::Cookies
use ActionDispatch::Session::CookieStore
use ActionDispatch::Flash
use ActionDispatch::ContentSecurityPolicy::Middleware
use ActionDispatch::PermissionsPolicy::Middleware
use Rack::Head
use Rack::ConditionalGet
use Rack::ETag
use Rack::TempfileReaper
use ActionDispatch::Static
run MyApp::Application.routes

Let's look at small piece of default middleware called ActionDispatch::RequestId (source). It already contains a good explaination of the inner workings in the comments of that file.

# https://github.com/rails/rails/blob/81c5c9971abe7a42a53ddbfede2683081a67e9d1/actionpack/lib/action_dispatch/middleware/request_id.rb

require "securerandom"
require "active_support/core_ext/string/access"

module ActionDispatch
  # Makes a unique request id available to the +action_dispatch.request_id+ env variable (which is then accessible
  # through ActionDispatch::Request#request_id or the alias ActionDispatch::Request#uuid) and sends
  # the same id to the client via the X-Request-Id header.
  #
  # The unique request id is either based on the X-Request-Id header in the request, which would typically be generated
  # by a firewall, load balancer, or the web server, or, if this header is not available, a random uuid. If the
  # header is accepted from the outside world, we sanitize it to a max of 255 chars and alphanumeric and dashes only.
  #
  # The unique request id can be used to trace a request end-to-end and would typically end up being part of log files
  # from multiple pieces of the stack.
  class RequestId
    def initialize(app, header:)
      @app = app
      @header = header
    end

    def call(env)
      req = ActionDispatch::Request.new env
      req.request_id = make_request_id(req.headers[@header])
      @app.call(env).tap { |_status, headers, _body| headers[@header] = req.request_id }
    end

    private
      def make_request_id(request_id)
        if request_id.presence
          request_id.gsub(/[^\w\-@]/, "").first(255)
        else
          internal_request_id
        end
      end

      def internal_request_id
        SecureRandom.uuid
      end
  end
end

The methods initialize and call are the key methods. The initialize method is called only once at booting the app and can be used to configure the middleware and set up data that is shared with all requests. The most important part is the call method definition. It is called at every request a client makes to a Rails app. The call method receives the current Rack environment as Ruby hash and should return a three element array with the status code, response headers and optionally a body. That array is generated by running @app.call(env). Of the two methods, you'll probably work most of the time with call.

In the RequestId middleware above, the environment (env) is changed by adding a header with a request id. That id is determined in the private methods of the middleware class. The header name is taken from the header keyword argument in initalize. This makes sense: the name of the header doesn't change during the app lifetime but the contents of the header do change each request. That's why the header name is set in initalize (runs once at boot time) and the value is set in call (runs each new request).

Our own middleware

Now that we understand a little more about middleware, let's create our own. We will create middleware that implements a simple rate limiter using Kredis. Kredis is a nice wrapper around Redis. Other than that, I assume you have some basic knowledge of Rails.

Basic setup

Make sure you have a working Redis server available. Run redis-cli in your shell to check if Redis is running. Something like 127.0.0.1:6379> should pop up.

Start by creating a new Rails app and add the Kredis gem.

$ rails new myratelimiter
$ cd myratelimiter
$ ./bin/bundle add kredis
$ ./bin/rails kredis:install

Check config/redis/shared.yml to make sure the configuration matches your Redis server. When you have a default Redis setup you probably don't have to change anything.

We need something to request from our Rails application so create a controller and model with ./bin/rails g scaffold User name email and run .bin/rails db:migrate.

Create five users by running 5.times {|i| User.create(name: "Name #{i}", email: "user#{i}@example.com")} in the Rails console.

Start the Rails server and run curl http://localhost:3000/users.json. You should see some json containing the five users.

Rate limiter

We have a basic api Rails app in place that shows users through /users.json. The problem is that consumers of our api can keep on sending requests so it's time to implement rate limiting.

The requirements of our rate limiter are simple for the sake of simplicity:

the consumer of the api cannot send more than 50 requests per 5 minutes
the consumer receives an http status 429 - Too many requests when the limit is reached
the server always responds with 3 headers:
- Rate-Limit-Reached: a boolean which tells the consumer if the rate limit is reached
- Rate-Limit-Requests-Left: a number which tells how many requests are left in the current time window
- Rate-Limit-Requests-Reset: a datetime which tells when the rate limiter will reset

Let's start with a few tests for the first two requirements:

# test/controllers/users_controller_test.rb
require 'test_helper'

class UsersControllerTest < ActionDispatch::IntegrationTest
  setup do
    Kredis.redis.flushall
    @user = users(:one)
  end

  test 'should not rate limit normal use' do
    49.times do
      get users_url
      assert_response :success
    end
  end

  test 'should rate limit abnormal use' do
    50.times do
      get users_url
      assert_response :success
    end

    get users_url
    assert_response :too_many_requests
  end
end

If you run the test suite you'll see that the second tests fails. Of course that's correct since we haven't implemented the rate limiter yet so let's go ahead with that.

Create lib/middleware/rate_limiter.rb with this class:

# lib/middleware/rate_limiter.rb - first draft
class RateLimiter
  def initialize(app)
    @app = app
  end

  def call(env)
    @app.call(env)
  end
end

Add this middleware to your app in config/application.rb:

...
require_relative '../lib/middleware/rate_limiter'
...
module Myratelimiter
  class Application < Rails::Application
    ...
    config.middleware.insert_before 0, RateLimiter
  end
end

That's it. You now have your own middleware in place. It doesn't do a lot at this point and running the test suite still yields the same failing tests.
Time to implement the first two requirements. I'm going to show you the complete middleware class for those requirements and go through it step by step.

# lib/middleware/rate_limiter.rb - second draft
class RateLimiter
  MAX_PER_WINDOW = 50
  WINDOW_SIZE = 1.minute

  def initialize(app)
    @app = app
  end

  def call(env)
    @req = ActionDispatch::Request.new(env)
    rate_limited? ? response_limit_reached : response_normal
  end

  private

  def rate_limited?
    request_counter.value >= MAX_PER_WINDOW
  end

  def kredis_key
    "rate_limiter:#{remote_ip}"
  end

  def request_counter
    # Only set the expires_in when the key is created
    # for the first time. Otherwise expires_in is
    # reset each time the key is accessed.

    if key_exists?
      Kredis.counter(kredis_key)
    else
      Kredis.counter(kredis_key, expires_in: WINDOW_SIZE)
    end
  end

  def key_exists?
    Kredis.redis.exists(kredis_key).positive?
  end

  def remote_ip
    # No need to re-invent logic to calculate the remote IP. It's already
    # available to use in ActionDispatch.

    ActionDispatch::RemoteIp::GetIp.new(@req, false, []).calculate_ip
  end

  def response_normal
    # Give back a normal response after incrementing the counter

    request_counter.increment
    @app.call(@req.env)
  end

  def response_limit_reached
    # We can also just put 429 here but this is more explicit.
    status_code = Rack::Utils::SYMBOL_TO_STATUS_CODE[:too_many_requests]

    [status_code, {}, []]
  end
end

If you run the test suite now, it will pass. If you want to test it yourself, set MAX_PER_WINDOW to 3 and run curl -v http://localhost:3000/users.json four times. The first three times the server responds as usual. The fourth time it returns an empty body with status code 429 - Too many requests.

What's happening in our newly created middleware? On each request the method call is called where an ActionDispatch::Request class is instantiated. With that object we can work easier with the data in env. Then we check if the rate limit is reached and return an appropriate response.

In rate_limited? the method request_counter is called which brings us to the part where Kredis is used. We use Kredis to initialize a counter in Redis. Kredis 'instantiates' the value from Redis. In other words, when you call Kredis.counter("mykey") we have an object that points to a Redis value under mykey. On that object we can call #increment which increments the current value in Redis. As you can see we check if the Redis key exists so that we can decide to use the call with expires_in. Each time you call #counter with expires_in, the expire timer resets. We don't want that because then the key will never expire. Checkout the Kredis docs for more information about Kredis.

We need some way to identify the consumer (or visitor if you will). A session is not an option since we don't have access to it yet and it's easy to circumvent by simple deleting cookies client side. Maybe an IP address? It's not really easy to spoof and you have to go through greater lengths to change the IP each n requests. On the other hand, a public IP is often shared amongst many users behind some router. For now, we go with an IP address but in the real world you'll probably identify a user with a token. In the method kredis_key the key is composed from remote_ip. In remote_ip we use logic that Rails already provides from ActionDispatch.

If the rate limit is not reached, response_normal is called where the request counter in Redis is incremented and the control is given back to Rails with @app.call(@req.env). But if the rate limit is reached, response_limit_reached is called. In that method we return a three element array containing the http status code, extra headers and a body. The last two are empty since the status code 'too many requests' is enough to inform the consumer the rate limit is reached.

You now have a working rate limiter as middleware. It might be too basic for a real world application but it does the job.

Extra headers

We still need to fulfill the third requirement: three headers containing information about the rate limiter.

Let's start with the first header called Rate-Limit-Reached. First we write a new test:

class UsersControllerTest < ActionDispatch::IntegrationTest
  ...
  test 'should return correct Rate-Limit-Reached header' do
    get users_url

    # You may be tempted to use `refute` but we want to make sure it really is false
    # instead of falsey like nil.
    assert_equal false, response.headers['Rate-Limit-Reached']

    50.times { get users_url } # Trigger rate limiter
    assert_equal true, response.headers['Rate-Limit-Reached']
  end
  ...
end

Run the test suite and it should fail. Add/replace the following methods:

class RateLimiter
  ...
  def rate_limit_headers
    {
      'Rate-Limit-Reached' => rate_limited?
    }
  end

  def response_normal
    # Give back a normal response after incrementing the counter

    request_counter.increment
    @app.call(@req.env).tap do |_status, headers, _body|
      rate_limit_headers.each { |key, value| headers[key] = value }
    end
  end

  def response_limit_reached
    # We can also just put 429 here but this is more explicit.
    status_code = Rack::Utils::SYMBOL_TO_STATUS_CODE[:too_many_requests]

    [status_code, rate_limit_headers, []]
  end
  ...
end

Now the test suite passes and we have implemented the first header requirement.

On to the next header called Rate-Limit-Requests-Left. Append the following to the test suite:

class UsersControllerTest < ActionDispatch::IntegrationTest
  ...
  test 'should return correct Rate-Limit-Left header' do
    get users_url

    assert_equal 49, response.headers['Rate-Limit-Left']

    50.times { get users_url } # Trigger rate limiter
    assert_equal 0, response.headers['Rate-Limit-Left']
  end
  ...
end

Run the test suite again to make sure it fails and add/replace the following methods:

class RateLimiter
  ...
  def rate_limit_headers
    {
      'Rate-Limit-Reached' => rate_limited?,
      'Rate-Limit-Left' => requests_left
    }
  end

  def requests_left
    MAX_PER_WINDOW - request_counter.value
  end
  ...
end

This will yield a passing test suite. We now have implemented the first two required headers.

Let's finish up with the third header called Rate-Limit-Requests-Reset. Again we start with a new test:

class UsersControllerTest < ActionDispatch::IntegrationTest
  ...
  test 'should return correct Rate-Limit-Reset header' do
    get users_url

    seconds_left = Time.parse(response.headers['Rate-Limit-Reset']) - Time.now.utc
    assert seconds_left.between?(1.minute - 5.seconds, 1.minute)

    50.times { get users_url } # Trigger rate limiter
    seconds_left = Time.parse(response.headers['Rate-Limit-Reset']) - Time.now.utc
    assert seconds_left.between?(1.minute - 5.seconds, 1.minute)
  end
  ...
end

And the logic to add/replace in the middleware:

class RateLimiter
  ...
  def rate_limit_headers
    {
      'Rate-Limit-Reached' => rate_limited?,
      'Rate-Limit-Left' => requests_left,
      'Rate-Limit-Reset' => reset_time
    }
  end

  def reset_time
    ttl = Kredis.redis.ttl(kredis_key) # Ask Redis how long the key has left to live
    (ttl >= 0 ? ttl.seconds.from_now : Time.zone.now).iso8601 # Create a datetime from TTL
  end
  ...
end

Run the test suite again, all tests should pass.

Wrapping it up

And now you have created some middleware to implement a rate limiter with a test suite. You can find the source of the above Rails application here.

The main purpose of this article was to show you how to develop middleware and not how to develop the world's most sophisticated rate limiter. But I hope it gave you some inspiration!

Notes

The test suite is kind of slow because we run the n.times { get users_url } blocks a lot... I'm going to let you figure out a refactor of the test suite to increase the performance. A small pointer: think about running two response cycles (normal and abnormal) and caching the results so each test can use the cached responses.
There is one important case not tested in the suite and that is the moment the request limit has been reset. Normally you can use travel but the TTL is set in Redis and not in the Rails app. You can refactor the middleware to accept a configuration value for the window size and set it to a few seconds so that it can be tested.

Blog