From 9dd758224978b9afca7f94337401c28f772ff2c8 Mon Sep 17 00:00:00 2001
From: Greg Harris <gharris@angrylogic.net>
Date: Tue, 18 Sep 2018 10:12:30 -0400
Subject: [PATCH] Enable wildcard with prefix.

Use case:

We use the Aurora scheduler in Mesos which assign tasks names using a
jobkey which consists of: "role/environment/jobname".  We often end up
in a scenario where we want all jobs running under a given "role" or
"role/environment" combination to have access to a Vault role.  Under
the current implementation this would require us to iterate all of the
possible combination and install policy files.

Previous versions of Vault gatekeeper allowed us to do wildcards with
prefixes as they used a simple glob match.  Moving to the radix tree
implementation broke this behaviour.

This patch changes the matching behaviour slightly.  Whenever we detect
a wildcard match we flag the policy as such.

When performing policy matching for wildcards we do a prefix search,
rather than looking for the last character being ":". Policy matching
for non wildcard policies is unchanged.

No changes to the radix tree structure are required.
---
 policy/policy.go      | 9 ++++++++-
 policy/policy_test.go | 8 ++++++++
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/policy/policy.go b/policy/policy.go
index ad2007a..32d41c5 100644
--- a/policy/policy.go
+++ b/policy/policy.go
@@ -13,6 +13,7 @@ type Policy struct {
 	Roles         []string `json:"roles"`
 	NumUses       int      `json:"num_uses"`
 	strictestPath []byte
+	wildcard      bool
 }
 
 func (p *Policy) merge(path []byte, other Policy) {
@@ -51,6 +52,9 @@ func LoadPoliciesFromJson(data []byte) (*Policies, error) {
 		tree := iradix.New()
 		txn := tree.Txn()
 		for k, v := range pol {
+			if strings.HasSuffix(k, "*") {
+				v.wildcard = true
+			}
 			if strings.HasSuffix(k, ":") {
 				return nil, errors.New("Invalid key name '" + k + "'. Keys must not end with a ':'")
 			}
@@ -79,7 +83,10 @@ func (p *Policies) Get(path string) (*Policy, bool) {
 
 	walkFn := func(k []byte, _v interface{}) bool {
 		v := _v.(Policy)
-		if bytes.Equal(k, []byte(path)) || k[len(k)-1] == ':' {
+		if v.wildcard && bytes.HasPrefix([]byte(path), k) {
+			ret.merge(k, v)
+			foundPolicy = true
+		} else if bytes.Equal(k, []byte(path)) {
 			ret.merge(k, v)
 			foundPolicy = true
 		}
diff --git a/policy/policy_test.go b/policy/policy_test.go
index 4188ce5..81fa5d8 100644
--- a/policy/policy_test.go
+++ b/policy/policy_test.go
@@ -36,6 +36,10 @@ const samplePolicy = `{
 	"mesos:framework:task2":{
 		"roles":["mesos_framework_task2"],
 		"num_uses":1
+	},
+	"mesos:framework:service/*": {
+	    "roles":["mesos_framework_service"],
+	    "num_uses":1
 	}
 }`
 
@@ -79,6 +83,10 @@ func TestSamplePolicy(t *testing.T) {
 			t.Fatalf("Test of '%s' failed. Expected: %v Had: %v", "mesos:jamp", expected, actual)
 		}
 
+		if pass, expected, actual := shouldContainAll(mustGet(pols.Get("mesos:framework:service/instance-1")), "wildcard", "mesos_child", "mesos_framework_child", "mesos_framework_service"); !pass {
+			t.Fatalf("Test of '%s' failed. Expected: %v Had: %v", "mesos:framework:service/instance-1", expected, actual)
+		}
+
 		if pass, _, actual := shouldContainAll(mustGet(pols.Get("mesos:framework:task2")), "mesos_framework_task"); pass {
 			t.Fatalf("Test of '%s' failed. 'task2' should not conatain permission of 'task'. Had: %v", "mesos:framework:task", actual)
 		}