From 0be05eb47cb825901b5481a4e730b4a5747638c9 Mon Sep 17 00:00:00 2001
From: olegdemchenko <oleg@oleg.name>
Date: Wed, 28 Sep 2016 21:21:07 +0300
Subject: [PATCH] policies implementation

---
 policy.go | 97 +++++++++++++++++++++++++++++++++++++++++++++++++++++++
 router.go | 10 +++++-
 2 files changed, 106 insertions(+), 1 deletion(-)
 create mode 100644 policy.go

diff --git a/policy.go b/policy.go
new file mode 100644
index 00000000..2b91fdcc
--- /dev/null
+++ b/policy.go
@@ -0,0 +1,97 @@
+// Copyright 2016 beego authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package beego
+
+import (
+	"strings"
+
+	"github.com/astaxie/beego/context"
+)
+
+// PolicyFunc defines a policy function which is invoked before the controller handler is executed.
+type PolicyFunc func(*context.Context)
+
+// FindRouter Find Router info for URL
+func (p *ControllerRegister) FindPolicy(cont *context.Context) []PolicyFunc {
+	var urlPath = cont.Input.URL()
+	if !BConfig.RouterCaseSensitive {
+		urlPath = strings.ToLower(urlPath)
+	}
+	httpMethod := cont.Input.Method()
+	isWildcard := false
+	// Find policy for current method
+	t, ok := p.policies[httpMethod]
+	// If not found - find policy for whole controller
+	if !ok {
+		t, ok = p.policies["*"]
+		isWildcard = true
+	}
+	if ok {
+		runObjects := t.Match(urlPath, cont)
+		if r, ok := runObjects.([]PolicyFunc); ok {
+			return r
+		} else if !isWildcard {
+			// If no policies found and we checked not for "*" method - try to find it
+			t, ok = p.policies["*"]
+			if ok {
+				runObjects = t.Match(urlPath, cont)
+				if r, ok = runObjects.([]PolicyFunc); ok {
+					return r
+				}
+			}
+		}
+	}
+	return nil
+}
+
+func (p *ControllerRegister) addToPolicy(method, pattern string, r ...PolicyFunc) {
+	method = strings.ToUpper(method)
+	p.enablePolicy = true
+	if !BConfig.RouterCaseSensitive {
+		pattern = strings.ToLower(pattern)
+	}
+	if t, ok := p.policies[method]; ok {
+		t.AddRouter(pattern, r)
+	} else {
+		t := NewTree()
+		t.AddRouter(pattern, r)
+		p.policies[method] = t
+	}
+}
+
+// Register new policy in beego
+func Policy(pattern, method string, policy ...PolicyFunc) {
+	BeeApp.Handlers.addToPolicy(method, pattern, policy...)
+}
+
+// Find policies and execute if were found
+func (p *ControllerRegister) execPolicy(cont *context.Context, urlPath string) (started bool) {
+	if !p.enablePolicy {
+		return false
+	}
+	// Find Policy for method
+	policyList := p.FindPolicy(cont)
+	if len(policyList) > 0 {
+		// Run policies
+		for _, runPolicy := range policyList {
+			runPolicy(cont)
+			if cont.ResponseWriter.Started {
+				return true
+			}
+		}
+		return false
+	}
+	return false
+}
diff --git a/router.go b/router.go
index 97b0edba..4e065a84 100644
--- a/router.go
+++ b/router.go
@@ -114,6 +114,8 @@ type controllerInfo struct {
 // ControllerRegister containers registered router rules, controller handlers and filters.
 type ControllerRegister struct {
 	routers      map[string]*Tree
+	enablePolicy bool
+	policies     map[string]*Tree
 	enableFilter bool
 	filters      [FinishRouter + 1][]*FilterRouter
 	pool         sync.Pool
@@ -122,7 +124,8 @@ type ControllerRegister struct {
 // NewControllerRegister returns a new ControllerRegister.
 func NewControllerRegister() *ControllerRegister {
 	cr := &ControllerRegister{
-		routers: make(map[string]*Tree),
+		routers:  make(map[string]*Tree),
+		policies: make(map[string]*Tree),
 	}
 	cr.pool.New = func() interface{} {
 		return beecontext.NewContext()
@@ -711,6 +714,11 @@ func (p *ControllerRegister) ServeHTTP(rw http.ResponseWriter, r *http.Request)
 		goto Admin
 	}
 
+	//check policies
+	if p.execPolicy(context, urlPath) {
+		goto Admin
+	}
+
 	if routerInfo != nil {
 		if BConfig.RunMode == DEV {
 			//store router pattern into context