Merge pull request #171 from qor5/optimize-perm-role

optimize perm

perm/builder.go

@@ -1,12 +1,12 @@
 package perm
 
 import (
-	"crypto/md5"
 	"errors"
 	"fmt"
 	"net/http"
 	"reflect"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/iancoleman/strcase"
@@ -57,11 +57,12 @@ func ToPermissionRN(v interface{}) []string {
 }
 
 type Builder struct {
-	policies      []*PolicyBuilder
-	ladon         *ladon.Ladon
-	subjectsFunc  SubjectsFunc
-	contextFunc   ContextFunc
-	dbPolicyModel DBPolicy
+	m            sync.Mutex
+	policies     []*PolicyBuilder
+	ladon        *ladon.Ladon
+	subjectsFunc SubjectsFunc
+	contextFunc  ContextFunc
+	dbPolicy     *DBPolicyBuilder
 }
 
 func New() *Builder {
@@ -74,17 +75,63 @@ func New() *Builder {
 }
 
 func (b *Builder) Policies(ps ...*PolicyBuilder) (r *Builder) {
-	b.policies = ps
-	for _, p := range b.policies {
-		if p.policy.ID == "" {
-			p.policy.ID = fmt.Sprintf("%x", md5.Sum(p.Json()))
+	b.DeletePolicies(b.policies...)
+	b.policies = make([]*PolicyBuilder, 0)
+	b.CreatePolicies(ps...)
+	return b
+}
+
+func (b *Builder) createPolicy(p *PolicyBuilder) {
+	p.setIDIfEmpty()
+	err := b.ladon.Manager.Create(p.policy)
+	if err != nil {
+		panic(err)
+	}
+	b.policies = append(b.policies, p)
+}
+
+func (b *Builder) updatePolicy(p *PolicyBuilder) {
+	i := b.findPolicyIndex(p.GetID())
+	if i < 0 {
+		return
+	}
+
+	err := b.ladon.Manager.Update(p.policy)
+	if err != nil {
+		panic(err)
+	}
+	b.policies[i] = p
+}
+
+func (b *Builder) findPolicyIndex(id string) int {
+	for i, p := range b.policies {
+		if p.GetID() == id {
+			return i
 		}
-		err := b.ladon.Manager.Create(p.policy)
-		if err != nil {
-			panic(err)
+	}
+	return -1
+}
+
+func (b *Builder) updateOrCreatePolicy(p *PolicyBuilder) {
+	i := b.findPolicyIndex(p.GetID())
+	if i < 0 {
+		b.createPolicy(p)
+	} else {
+		b.updatePolicy(p)
+	}
+}
+
+func (b *Builder) deletePolicy(p *PolicyBuilder) {
+	for i, ep := range b.policies {
+		if ep.GetID() == p.GetID() {
+			err := b.ladon.Manager.Delete(p.GetID())
+			if err != nil {
+				panic(err)
+			}
+			b.policies = append(b.policies[:i], b.policies[i+1:]...)
+			break
 		}
 	}
-	return b
 }
 
 func (b *Builder) SubjectsFunc(v SubjectsFunc) (r *Builder) {
@@ -92,38 +139,62 @@ func (b *Builder) SubjectsFunc(v SubjectsFunc) (r *Builder) {
 	return b
 }
 
+func (b *Builder) GetSubjectsFunc() SubjectsFunc {
+	return b.subjectsFunc
+}
+
 func (b *Builder) ContextFunc(v ContextFunc) (r *Builder) {
 	b.contextFunc = v
 	return b
 }
 
-func (b *Builder) EnableDBPolicy(db *gorm.DB, dbPolicyModel DBPolicy, loadDuration time.Duration) (r *Builder) {
-	if dbPolicyModel == nil {
-		b.dbPolicyModel = DefaultDBPolicy{}
-	} else {
-		b.dbPolicyModel = dbPolicyModel
-	}
+func (b *Builder) GetContextFunc() ContextFunc {
+	return b.contextFunc
+}
 
-	go b.loopLoadDBPolicies(db, loadDuration)
+func (b *Builder) DBPolicy(dpb *DBPolicyBuilder) (r *Builder) {
+	b.dbPolicy = dpb
+
+	go b.loopLoadDBPolicies(dpb.db, dpb.loadFrequency)
 	return b
 }
 
+func (b *Builder) CreatePolicies(ps ...*PolicyBuilder) {
+	b.m.Lock()
+	defer b.m.Unlock()
+	for _, p := range ps {
+		b.createPolicy(p)
+	}
+}
+
 func (b *Builder) UpdatePolicies(toUpdate ...*PolicyBuilder) {
+	b.m.Lock()
+	defer b.m.Unlock()
+	for _, p := range toUpdate {
+		b.updatePolicy(p)
+	}
+}
+
+func (b *Builder) UpdateOrCreatePolicies(toUpdate ...*PolicyBuilder) {
+	b.m.Lock()
+	defer b.m.Unlock()
 	for _, p := range toUpdate {
-		b.ladon.Manager.Update(p.policy)
+		b.updateOrCreatePolicy(p)
 	}
 }
 
 func (b *Builder) DeletePolicies(toDelete ...*PolicyBuilder) {
+	b.m.Lock()
+	defer b.m.Unlock()
 	for _, p := range toDelete {
-		b.ladon.Manager.Delete(p.GetID())
+		b.deletePolicy(p)
 	}
 }
 
 func (b *Builder) LoadDBPoliciesToMemory(db *gorm.DB, startFrom *time.Time) {
-	toUpdate, toDelete := b.dbPolicyModel.LoadDBPolicies(db, startFrom)
+	toUpdateOrCreate, toDelete := b.dbPolicy.model.LoadDBPolicies(db, startFrom)
 	b.DeletePolicies(toDelete...)
-	b.UpdatePolicies(toUpdate...)
+	b.UpdateOrCreatePolicies(toUpdateOrCreate...)
 	if Verbose {
 		b.printPolices()
 	}
@@ -141,6 +212,7 @@ func (b *Builder) loopLoadDBPolicies(db *gorm.DB, duration time.Duration) {
 
 func (b *Builder) printPolices() {
 	allp, _ := b.ladon.Manager.GetAll(100, 0)
+	fmt.Printf("all permission policies: \n")
 	for _, p := range allp {
 		fmt.Printf("%+v \n", p)
 	}

perm/db_policy.go

@@ -1,13 +1,38 @@
 package perm
 
 import (
-	"github.com/lib/pq"
-	"gorm.io/gorm"
 	"strconv"
 	"strings"
 	"time"
+
+	"github.com/lib/pq"
+	"gorm.io/gorm"
 )
 
+type DBPolicyBuilder struct {
+	db            *gorm.DB
+	model         DBPolicy
+	loadFrequency time.Duration
+}
+
+func NewDBPolicy(db *gorm.DB) *DBPolicyBuilder {
+	return &DBPolicyBuilder{
+		db:            db,
+		model:         DefaultDBPolicy{},
+		loadFrequency: time.Minute,
+	}
+}
+
+func (dpb *DBPolicyBuilder) Model(m DBPolicy) *DBPolicyBuilder {
+	dpb.model = m
+	return dpb
+}
+
+func (dpb *DBPolicyBuilder) LoadFrequency(d time.Duration) *DBPolicyBuilder {
+	dpb.loadFrequency = d
+	return dpb
+}
+
 type DefaultDBPolicy struct {
 	gorm.Model
 
@@ -18,7 +43,7 @@ type DefaultDBPolicy struct {
 	Resources pq.StringArray `gorm:"type:text[]"`
 }
 
-func (p DefaultDBPolicy) LoadDBPolicies(db *gorm.DB, startFrom *time.Time) (toUpdate []*PolicyBuilder, toDelete []*PolicyBuilder) {
+func (p DefaultDBPolicy) LoadDBPolicies(db *gorm.DB, startFrom *time.Time) (toUpdateOrCreate []*PolicyBuilder, toDelete []*PolicyBuilder) {
 	var ps []DefaultDBPolicy
 	if startFrom == nil || startFrom.IsZero() {
 		db.Find(&ps)
@@ -30,7 +55,7 @@ func (p DefaultDBPolicy) LoadDBPolicies(db *gorm.DB, startFrom *time.Time) (toUp
 		if p.DeletedAt.Valid {
 			toDelete = append(toDelete, p.ToPolicy())
 		} else {
-			toUpdate = append(toUpdate, p.ToPolicy())
+			toUpdateOrCreate = append(toUpdateOrCreate, p.ToPolicy())
 		}
 	}
 	return

perm/policy.go

@@ -1,7 +1,9 @@
 package perm
 
 import (
+	"crypto/md5"
 	"encoding/json"
+	"fmt"
 	"strings"
 
 	"github.com/ory/ladon"
@@ -59,9 +61,16 @@ func (b *PolicyBuilder) Given(conditions Conditions) (r *PolicyBuilder) {
 	return b
 }
 
-func (b PolicyBuilder) Json() []byte {
-	str, _ := json.Marshal(b.policy)
-	return str
+func (b *PolicyBuilder) setIDIfEmpty() {
+	if b.policy.ID != "" {
+		return
+	}
+
+	bs, err := json.Marshal(b.policy)
+	if err != nil {
+		panic(err)
+	}
+	b.policy.ID = fmt.Sprintf("%x", md5.Sum(bs))
 }
 
 func (b *PolicyBuilder) GetID() string {