Add IP restriction on Rack app for specific accounts

safelist_per_api_key.rb

# frozen_string_literal: true

# Read the blog:

require "bundler/inline"

gemfile(true) do
  source "https://rubygems.org"

  gem "rack-attack", "~> 6.7.0"
  gem "rack-test", "~> 2.2.0"
  gem "rspec", "~> 3.13.0"
  gem "redis", "~> 5.4.0"
  gem "connection_pool", "~> 2.5.0"
end

require "redis"
require "connection_pool"
require "digest"
require "securerandom"
require "json"

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 and these requests will fail for authenticated routes
    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

    !JSON.parse(allowed_ips).include?(req.ip)
  end

  private

  # Prefer not storing raw API keys in cache for security
  def cache_key
    Digest::SHA1.hexdigest(api_key)
  end

  def api_key
    req.params["api_key"]
  end
end

#### TESTING ####

require "rack/attack"
require "rspec"
require "rack/test"

RSpec.configure do |config|
  config.include Rack::Test::Methods
end

Rack::Attack.blocklist("whitelist IPs") do |req|
  RackAttackHelper.new(req).blocklisted?
end

# For the sake of this minimal script without Rails,
# I create a dummy rack-application and use rack-test for testing it
# The main purpose is to test the rack-attack config.
describe "Application" do
  let(:app) do
    Rack::Builder.new do
      use Rack::Attack

      # Sample rack application
      map "/" do
        run lambda { |env| [200, {"content-type" => "text/plain"}, ["All good!"]] }
      end
    end
  end

  let(:api_key) { SecureRandom.uuid }
  let(:cache_key) { Digest::SHA1.hexdigest(api_key) }
  let(:allowed_ips) { ["10.0.0.1", "10.0.0.106"] }

  context "when no safelisted IPs present" do
    it "returns success" do
      get "/", { api_key: }
      expect(last_response.body).to eq("All good!")
    end
  end

  context "when safelisted IPs present" do
    before { allow_any_instance_of(Rack::Attack::Request).to receive(:ip).and_return(request_ip) }

    around do |example|
      REDIS.with { |conn| conn.set(cache_key, allowed_ips) }
      Rack::Attack.enabled = true
      example.run
      REDIS.with { |conn| conn.del(cache_key) }
      Rack::Attack.enabled = false
    end

    context "when IP exists in safelist" do
      let(:request_ip) { "10.0.0.106" }

      it "returns success when IP exists in safelist" do
        get "/", { api_key: }
        expect(last_response.body).to eq("All good!")
      end
    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
end