Rails Model Firewall Mixin Tuesday, August 26, 2008

At my company, Pharos, we're about to launch a new product which will contain sensitive data for multiple firms in a single database. This is essentially a lightweight version of our flagship product, which was built for a single client.

Of course, as a result, I had to refactor like crazy to get rid of the implicit "one-firm" assumption that was built into the code and database schemas.

The essential task was to add "firm_id" to each of the private table schemas, and then make sure that all the code that accesses the model specifies the firm in the query. The two access idioms that were being widely used (unsurprisingly):

results = ClassName.find(:all, :conditions => [....])

and

results = ClassName.find_by_entity_id_and_hour(...)

I was able to make minimal changes to the code by supporting the following new idioms through a mixin (the mixin code is at the end of the article):

results = ClassName.find_in_firm_scope(firm_id, :all, :conditions => [....])

results = ClassName.with_firm_scope(firm_id) do |klass|
  klass.find_by_entity_id_and_hour(...)
end

(The second idiom I found easier to make (and the diff easier to read) than:

ClassName.find_by_firm_id_and_entity_id_and_hour(firm_id, ...)

but really, that's a matter of taste.)

But I was still nervous. What if I missed an instance of a database lookup that wasn't specifying firm, and as a result one client saw another client's records? That would be a Really Bad Thing ^TM, and I want to explicitly make sure that can't happen. But how?

After a half hour of poking around and futzing, I came up with a find()-and-friends implementation that will check with_scope conditions as well as the :conditions parameter to the find() call:

>> My::PrivateModel.find_by_entity_id(1)
RuntimeError: My::PrivateModel PrivateRecord find() did not specify firm_id

Without further ado, here's the mixin:

# lib/private_record.rb
module PrivateRecord
  def self.included(base)
    base.validates_presence_of :firm_id
    base.extend PrivateRecordClassExtendor
  end
end

module PrivateRecordClassExtendor

  def find_every(*args)
    check_for_firm_id(*args)
    super(*args)
  end

  # the DRY idiom here is: results = ClassName.with_firm_scope(firm) {|klass| klass.find(...) }
  def with_firm_scope(firm, &block)
    with_scope(:find => {:conditions => "firm_id = #{firm}"}, :create => {:firm_id => firm}) do
      yield self
    end
  end

  def find_in_firm_scope(firm, *args)
    with_firm_scope(firm) do
      find(*args)
    end
  end

private
  FIRM_ID_RE = /firm_id =/
  def check_for_firm_id(*args)
    ok = false
    if scoped_methods
      scoped_methods.each do |j|
        if j[:find] && j[:find][:conditions] && j[:find][:conditions] =~ FIRM_ID_RE
          ok = true 
          break
        end
      end
    end
    if !ok
      args.each do |j|
        if j.is_a?(Hash) && j[:conditions]
          if (j[:conditions].is_a?(String) && j[:conditions] =~ FIRM_ID_RE) \
             or (j[:conditions].is_a?(Hash) && j[:conditions][:firm_id])
            ok = true 
            break
          end
        end
      end
    end
    raise "#{self} PrivateRecord find() did not specify firm_id" if !ok
  end
end

The magic is all in the check_for_firm_id() method. To use this, simply:

include PrivateRecord

and go to town.

Oh, and lest ye be skeptical, here are the test cases:

require File.dirname(__FILE__) + '/../test_helper'

class PrivateModelTest < ActiveSupport::TestCase

  fixtures :isone_da_schedules

  def test_privaterecord_disallow_find_requirement
    assert_raises(RuntimeError) { My::PrivateModel.find(1) }
    assert_raises(RuntimeError) { My::PrivateModel.find_by_entity_id(1) }
    assert_raises(RuntimeError) { My::PrivateModel.find_all_by_entity_id(1) }
    assert_raises(RuntimeError) { My::PrivateModel.find(:all, :conditions => 'entity_id = 1') }
    assert_raises(RuntimeError) { My::PrivateModel.find(:first, :conditions => 'entity_id = 1') }
  end

  def test_privaterecord_allow_find_requirement
    assert_nothing_thrown { My::PrivateModel.find_in_firm_scope(1, 1) }
    assert_nothing_thrown { My::PrivateModel.with_firm_scope(1) {|k| k.find_by_entity_id(1) } }
    assert_nothing_thrown { My::PrivateModel.with_firm_scope(1) {|k| k.find_all_by_entity_id(1) } }
    assert_nothing_thrown { My::PrivateModel.find_in_firm_scope(1, :all, :conditions => 'entity_id = 0') }
    assert_nothing_thrown { My::PrivateModel.find_in_firm_scope(1, :first, :conditions => 'entity_id = 0') }
    assert_nothing_thrown { My::PrivateModel.find_by_firm_id_and_entity_id(1, 1) }
  end

end

Let me know in the comments if you found this at all useful! Keep coding.