Adding IP restriction to Rack app for specific accounts
Exploring an interesting requirement of adding IP restrictions on a Rails app.
For a global IP rule on the entire application, my go-to solution is putting a CDN/infrastructure level whitelist rule - minimal, blazing fast, effective and secure.
But this requirement was to allow corporate customers to whitelist specific IPs for their accounts. In other words, their account (and only their specific account) should be accessible the IPs they specify and be restricted for the outside world.
This makes things complicated interesting. Let us explore the possibilities.
The Rails way
Since the whitelist will be account-specific and optional, it will be stored in database. In Rails, the first solution that comes to my mind is to use a before_action
:
class ApplicationController < ActionController::Base
before_action :require_whitelisted_ip
def require_whitelisted_ip
return unless params[:api_key]
whitelisted_ips = Account.find_by(api_key: params[:api_key])&.whitelisted_ips
return unless whitelisted_ips.present? # whitelisting is optional, default is to allow everything
whitelisted_ips.include?(request.remote_ip)
end
end
This is simple, elegant and works as per requirement - but has a downside of having to query the database on every request, even for invalid ones. For large-scale applications, it is not a good idea to add a global level query for every request. The performance hits can be massive on the database, not to mention the slowdown in response times.
Caching
We can fix this by caching the whitelisted IPs in redis/memcache.
class ApplicationController < ActionController::Base
before_action :require_whitelisted_ip
private
def require_whitelisted_ip
return unless params[:api_key]
whitelisted_ips = Rails.cache.fetch(cache_key(params[:api_key])) do
Account.find_by(api_key: params[:api_key])&.whitelisted_ips
end
return unless whitelisted_ips
whitelisted_ips.include?(request.remote_ip)
end
end
This significantly reduces the overhead and is highly scalable. We can keep the cache in sync with IPs in database easily since they don't change very often.
For most applications this is a fine solution. But let us take it up a notch by restricting the invalid requests even before they hit our Rails controllers!
Rack ⚡
Rails is built on Rack - a modular interface for web applications. We can add a rack middleware to validate IPs.
Speaking of rack middlewares, my earlier blog post about unit testing rack middlewares is a fun read.
Creating our own middleware should be fairly easy, but why reinvent the wheel when an extremely minimal tool tailored for allowing/restricting rack requests exists: rack-attack.
# config/initializers/rack_attack.rb
# Rails.cache is not available here: initialize redis connection
REDIS = ConnectionPool.new(size: 5) do
Redis.new(host: "localhost", port: "6379")
end
class RackAttackHelper
attr_reader :req
def initialize(req)
@req = req
end
def blocklisted?
# allow requests without API key - API key will be validated later
return false unless api_key
allowed_ips = REDIS.with { |conn| conn.get(cache_key) }
# Whitelisting is optional
# allow if whitelist is not configured for account
return false unless allowed_ips
# Redis returns a string, use JSON parse to get an array
!JSON.parse(allowed_ips).include?(req.ip)
end
private
# Prefer not storing plain API keys in cache for security
def cache_key
Digest::SHA1.hexdigest(api_key)
end
def api_key
req.params["api_key"]
end
end
Rack::Attack.blocklist("whitelist IPs") do |req|
# Requests are blocked if the return value is truthy
RackAttackHelper.new(req).blocklisted?
end
Rack::Attack.enabled = true
Let us unpack above code: Firstly we store the whitelisted IPs against account's API key. But we don't have access to Rails.cache
at rack level, so we need to initialize the redis connection. If you've read my blog on redis connection pool, you know it is always better to use a connection pool instead of a single redis connection, especially for high-traffic applications.
We've extracted the logic to check IP in a separate class. We hash the API keys for more security when storing in cache. We configure rack-attack
to block requests if request IP does not exist in whitelist.
Finally we can test this entire setup:
describe "IP whitelisting" do
let(:api_key) { SecureRandom.uuid }
let(:cache_key) { Digest::SHA1.hexdigest(api_key) }
let(:allowed_ips) { ["10.0.0.1", "10.0.0.106"] }
around do |example|
CACHE.set(cache_key, allowed_ips)
Rack::Attack.enabled = true
example.run
CACHE.del(cache_key)
Rack::Attack.enabled = false
end
before do
allow_any_instance_of(Rack::Attack::Request)
.to receive(:ip).and_return(request_ip)
end
context "when IP does not exist in safelist" do
let(:request_ip) { "10.0.0.200" }
it "returns success when IP exists in safelist" do
get "/", { api_key: }
expect(last_response.status).to eq(403)
expect(last_response.body).to eq("Forbidden\n")
end
end
end
Sidenote: For plain rack apps without Rails, I found rack-test a nifty tool for testing.
Was this rack solution a bit overkill? Probably. Did I have fun exploring this solution? Definitely.
All code along with tests can be found in this executable gist.