Custom Rego Function by Example

Open Policy Agent and its own language Rego are rich in opportunities and at the same time easy to learn. However, sometimes realities start to demand more and more as the project grows and developers need to design solutions. Fortunately, Rego allows us to create custom built-in functions and now I’ll tell you how to do it by live example.

I’ve been implementing OPA for authorization clients in the API. It was required, based on the user’s association with groups (teams, units, etc.), to allow or not some action (in reality, much more factors were involved in politics, but this is irrelevant to the current topic).

So let’s say we have some kind of group structure and users that can be in more than one group at a time. And the groups themselves also have relations "many-to-many".

As you can see, one of the developers is a member of two teams at once. He initially worked for the billing team, but had the devops skills and wanted to grow as a devops engineer. There is also a SWAT team that is designed to put out fires and help other teams to remove the required volumes in the right time, so it is part of two other groups.

And there are two actions in the system:

• cordon - cordon node
• deploy - deploy the application in the production

Actions are group-bound and also apply to all child groups.

The task is to write a rule for the OPA.

The solution first requires an org structure in json format to be used in rego rules. About this:

{
    "users": [
        {
            "id": "oliver",
            "groups": ["billing"]
        },
        {
            "id": "liam",
            "groups": ["swat"]
        },
    ],
    "groups": [
        {
            "id": "infra",
            "parent": []
        },
        {
            "id": "devops",
            "parent": ["infra"]
        },
        {
            "id": "admin",
            "parent": ["infra"]
        },
        {
            "id": "dev",
            "parent": []
        },
        {
            "id": "search",
            "parent": ["dev"]
        },
        {
            "id": "billing",
            "parent": ["dev"]
        },
        {
            "id": "swat",
            "parent": ["infra", "billing"]
        },
    ],
}

Since the rules are assigned to groups, it is enough to check whether the user is in the desired group or not. And whether a user group is a subgroup of another group (at any level of nesting). This is where the problem arose. If recursion is done on Rego (I’m not sure it is possible), then readability will be terrible. There are no ready-made functions either. But luckily there is a possibility to extend Rego with its functions.

Here is the implementation of the custom function itself:

// search all group parents
func RegoGroupParentsFunction() func(*rego.Rego) {
	return rego.Function2(&rego.Function{
		Name: "group_parents",
		Decl: types.NewFunction(types.Args(types.S, types.A), types.A),
	},
		func(bctx rego.BuiltinContext, groupID, groupsData *ast.Term) (*ast.Term, error) {
			groups := []Group{}
			if err := json.Unmarshal([]byte(groupsData.Value.String()), &groups); err != nil {
				return nil, fmt.Errorf("unmarshal rego groups data: %v", err)
			}

			mappedGroups := map[string]*Group{}
			for k, group := range groups {
				mappedGroups[group.ID] = &groups[k]
			}

			gID := trimDoubleQuotes(groupID.Value.String())
			parentGroups := []*Group{}
			SearchParentsRecursive(gID, mappedGroups, &parentGroups)

			values := []*ast.Term{}
			for _, v := range parentGroups {
				val, err := ast.InterfaceToValue(v)
				if err != nil {
					return nil, fmt.Errorf("convert group to rego value: %v", err)
				}
				values = append(values, ast.NewTerm(val))
			}

			return ast.ArrayTerm(values...), nil
		})
}

func SearchParentsRecursive(groupID string, groups map[string]*Group, result *[]*Group) {
	group, ok := groups[groupID]
	if !ok {
		return
	}
	if len(group.Parents) == 0 {
		return
	}
	for _, parentID := range group.Parents {
		parentGroup, ok := groups[parentID]
		if !ok {
			continue
		}
		SearchParentsRecursive(parentID, groups, result)
		*result = append(*result, parentGroup)
	}
}

I can describe what’s going on here:

return rego.Function2(&rego.Function{
		Name: "group_parents",
		Decl: types.NewFunction(types.Args(types.S, types.A), types.A),
	},
		func(bctx rego.BuiltinContext, groupID, groupsData *ast.Term) (*ast.Term, error) {

A new group_parents function is defined, and has two arguments: the group id that you are looking for and the group structure described above. Accordingly for rego functions with three arguments there is a function Function3, with four Function4, etc.

Next, parse json with the structure and search recursively for all parents takes place. Difficulties arise when the result is returned, because opa-sdk interface is not very obvious.

values := []*ast.Term{}
for _, v := range parentGroups {
	val, err := ast.InterfaceToValue(v)
	if err != nil {
		return nil, fmt.Errorf("convert group to rego value: %v", err)
	}
	values = append(values, ast.NewTerm(val))
}

return ast.ArrayTerm(values...), nil

The function must return an array with all parent groups to go through the rules themselves. Therefore, we return ast.ArrayTerm. This is an array of rego values (Term) that must be previously created with ast.NewTerm. Before that you will also need to convert string values from the id of the ast.InterfaceToValue.

Now group_parents is available in rego policies. I’ll show you how it works:

package group_search

import future.keywords.in

default parent_groups_is_ok = false

# Проверяем что группа SWAT состоит во всех вышестоящих группах billing и devops
parent_groups_is_ok {
   groups := group_parents("swat", data.groups)

   groups[0].name == "billing"
   groups[1].name == "devops"
}

# Убеждаемся что в лишних группах swat нет
parent_groups_not_exists {
   groups := group_parents("swat", data.groups)
   groups[_].name != "search"
}

Now I can test:

func TestGroupParentsOk(t *testing.T) {
   query, err := rego.New(
      rego.Query("data.group_search.parent_groups_is_ok"),
      RegoGroupParentsFunction()
      rego.LoadBundle("testdata"),
   ).PrepareForEval(context.Background())
   if err != nil {
      t.Fatalf("prepare rego query: %v", err)
   }

   resultSet, err := query.Eval(context.Background())
   if err != nil {
      t.Fatalf("eval rego query: %v", err)
   }

   if len(resultSet) == 0 {
      t.Error("undefined result")
   }

   assert.True(t, resultSet.Allowed())
}

func TestGroupParentsNotExist(t *testing.T) {
   query, err := rego.New(
      rego.Query("data.group_search.runtime_parent_groups_not_exists"),
      RegoGroupParentsFunction(),
      rego.LoadBundl("testdata"),
   ).PrepareForEval(context.Background())
   if err != nil {
      t.Fatalf("prepare rego query: %v", err)
   }

   resultSet, err := query.Eval(context.Background())
   if err != nil {
      t.Fatalf("eval rego query: %v", err)
   }

   if len(resultSet) == 0 {
      t.Error("undefined result")
   }

   assert.True(t, resultSet.Allowed())
}

Thus it is possible to extend the rego as much as you need. For example, add work with some external sources: databases, headlines, etc. Or how in my case to implement complex logic.