From 14e613da954eebc55b193cc24e50e88a7128d6f7 Mon Sep 17 00:00:00 2001
From: Ryan Clarke <faceless.saint@gmail.com>
Date: Thu, 15 Mar 2018 15:13:57 -0400
Subject: [PATCH] [feature] Add policy filtering

Allow for filtered policy loading, so that the loaded policy is
a targeted subset of the policy in storage. This allows the policy
enforcement to more effectively scale for very large policies in
a multi-tenant environment by only loading the policies relevant
to the requested tenant.

To protect the full policy from accidental corruption, a filtered
policy cannot be saved back to storage.
---
 README.md       | 15 +++++++++++++++
 adapter.go      | 30 +++++++++++++++++++++++++++++-
 adapter_test.go | 45 +++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 89 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 7a434e3..23a4194 100644
--- a/README.md
+++ b/README.md
@@ -45,6 +45,21 @@ func main() {
 }
 ```
 
+## Filtered Policies
+
+```go
+import "gopkg.in/mgo.v2/bson"
+
+// This adapter also implements the FilteredAdapter interface. This allows for
+// efficent, scalable enforcement of very large policies:
+filter := &bson.M{"v0": "alice"}
+e.LoadFilteredPolicy(filter)
+
+// The loaded policy is now a subset of the policy in storage, containing only
+// the policy lines that match the provided filter. This filter should be a
+// valid MongoDB selector using BSON. A filtered policy cannot be saved.
+```
+
 ## Getting Help
 
 - [Casbin](https://github.com/casbin/casbin)
diff --git a/adapter.go b/adapter.go
index 7b999cd..6252f80 100644
--- a/adapter.go
+++ b/adapter.go
@@ -15,6 +15,7 @@
 package mongodbadapter
 
 import (
+	"errors"
 	"runtime"
 
 	"github.com/casbin/casbin/model"
@@ -38,6 +39,7 @@ type adapter struct {
 	url        string
 	session    *mgo.Session
 	collection *mgo.Collection
+	filtered   bool
 }
 
 // finalizer is the destructor for adapter.
@@ -59,6 +61,13 @@ func NewAdapter(url string) persist.Adapter {
 	return a
 }
 
+// NewFilteredAdapter is the constructor for FilteredAdapter. Behavior is
+// otherwise indentical to the NewAdapter function.
+func NewFilteredAdapter(url string) persist.FilteredAdapter {
+	// The adapter already supports the new interface, it just needs to be retyped.
+	return NewAdapter(url).(*adapter)
+}
+
 func (a *adapter) open() {
 	dI, err := mgo.ParseURL(a.url)
 	if err != nil {
@@ -156,8 +165,19 @@ LineEnd:
 
 // LoadPolicy loads policy from database.
 func (a *adapter) LoadPolicy(model model.Model) error {
+	return a.LoadFilteredPolicy(model, nil)
+}
+
+// LoadFilteredPolicy loads matching policy lines from database. If not nil,
+// the filter must be a valid MongoDB selector.
+func (a *adapter) LoadFilteredPolicy(model model.Model, filter interface{}) error {
+	if filter == nil {
+		a.filtered = false
+	} else {
+		a.filtered = true
+	}
 	line := CasbinRule{}
-	iter := a.collection.Find(nil).Iter()
+	iter := a.collection.Find(filter).Iter()
 	for iter.Next(&line) {
 		loadPolicyLine(line, model)
 	}
@@ -165,6 +185,11 @@ func (a *adapter) LoadPolicy(model model.Model) error {
 	return iter.Close()
 }
 
+// IsFiltered returns true if the loaded policy has been filtered.
+func (a *adapter) IsFiltered() bool {
+	return a.filtered
+}
+
 func savePolicyLine(ptype string, rule []string) CasbinRule {
 	line := CasbinRule{
 		PType: ptype,
@@ -194,6 +219,9 @@ func savePolicyLine(ptype string, rule []string) CasbinRule {
 
 // SavePolicy saves policy to database.
 func (a *adapter) SavePolicy(model model.Model) error {
+	if a.filtered {
+		return errors.New("cannot save a filtered policy")
+	}
 	if err := a.dropTable(); err != nil {
 		return err
 	}
diff --git a/adapter_test.go b/adapter_test.go
index 08e8f00..76f6fff 100644
--- a/adapter_test.go
+++ b/adapter_test.go
@@ -20,6 +20,7 @@ import (
 
 	"github.com/casbin/casbin"
 	"github.com/casbin/casbin/util"
+	"gopkg.in/mgo.v2/bson"
 )
 
 var testDbURL = os.Getenv("TEST_MONGODB_URL")
@@ -137,6 +138,50 @@ func TestAdapter(t *testing.T) {
 	testGetPolicy(t, e, [][]string{})
 }
 
+func TestFilteredAdapter(t *testing.T) {
+	// Now the DB has policy, so we can provide a normal use case.
+	// Create an adapter and an enforcer.
+	// NewEnforcer() will load the policy automatically.
+	a := NewAdapter(getDbURL())
+	e := casbin.NewEnforcer("examples/rbac_model.conf", a)
+
+	// Load filtered policies from the database.
+	e.AddPolicy("alice", "data1", "write")
+	e.AddPolicy("bob", "data2", "write")
+	// Reload the filtered policy from the storage.
+	filter := &bson.M{"v0": "bob"}
+	if err := e.LoadFilteredPolicy(filter); err != nil {
+		t.Errorf("Expected LoadFilteredPolicy() to be successful; got %v", err)
+	}
+	// Only bob's policy should have been loaded
+	testGetPolicy(t, e, [][]string{{"bob", "data2", "write"}})
+
+	// Verify that alice's policy remains intact in the database.
+	filter = &bson.M{"v0": "alice"}
+	if err := e.LoadFilteredPolicy(filter); err != nil {
+		t.Errorf("Expected LoadFilteredPolicy() to be successful; got %v", err)
+	}
+	// Only alice's policy should have been loaded,
+	testGetPolicy(t, e, [][]string{{"alice", "data1", "write"}})
+
+	// Test safe handling of SavePolicy when using filtered policies.
+	if err := e.SavePolicy(); err == nil {
+		t.Errorf("Expected SavePolicy() to fail for a filtered policy")
+	}
+	if err := e.LoadPolicy(); err != nil {
+		t.Errorf("Expected LoadPolicy() to be successful; got %v", err)
+	}
+	if err := e.SavePolicy(); err != nil {
+		t.Errorf("Expected SavePolicy() to be successful; got %v", err)
+	}
+
+	e.RemoveFilteredPolicy(2, "write")
+	if err := e.LoadPolicy(); err != nil {
+		t.Errorf("Expected LoadPolicy() to be successful; got %v", err)
+	}
+	testGetPolicy(t, e, [][]string{})
+}
+
 func TestNewAdapterWithInvalidURL(t *testing.T) {
 	defer func() {
 		if r := recover(); r == nil {