Automatic source locations with Rego

SnykSec - Feb 13 - - Dev Community

rules written in Rego, and customers can add their own custom rules as well.

We recently released a series of improvements to Snyk IaC, and in this blog post, we’re taking a technical dive into a particularly interesting feature — automatic source code locations for rule violations.

When checking IaC files against known issues, the updated snyk iac test command will show accurate file, line, and column information for each rule violation. This works even for custom rules, without the user doing any work.


But before we provide a standalone proof-of-concept for this technique, we’ll need to make some simplifications. The full implementation of this is available in our unified policy engine.

Let’s start by looking at a CloudFormation example. While our IaC engine supports many formats, with a strong focus on Terraform, CloudFormation is a good example since we can parse it without too many dependencies (it’s just YAML, after all).

Resources:
  Vpc:
    Type: AWS::EC2::VPC
    Properties:
      CidrBlock: 10.0.0.0/16
  PublicSubnet:
    Type: AWS::EC2::Subnet
    Properties:
      VpcId:
        Ref: Vpc
      CidrBlock: 10.0.0.0/24
  PrivateSubnet:
    Type: AWS::EC2::Subnet
    Properties:
      VpcId:
        Ref: Vpc
      CidrBlock: 10.0.128.0/20
Enter fullscreen mode Exit fullscreen mode

We want to ensure that no subnets use a CIDR block larger than /24, so let’s write a Rego policy to do just that:

package policy

deny[resourceId] {
    resource := input.Resources[resourceId]
    resource.Type == "AWS::EC2::Subnet"
    [_, mask] = split(resource.Properties.CidrBlock, "/")
    to_number(mask) < 24
}
Enter fullscreen mode Exit fullscreen mode

This way, deny will produce a set of denied resources. We won’t go into the details of how Rego works, but if you want to learn more, we recommend the excellent OPA by Example course.

We can subdivide the problem into two parts:

  1. We’ll want to infer that our policy uses the CidrBlock attribute
  2. Then, we’ll retrieve the source code location

Let’s start with (2) since it provides a good way to familiarize ourselves with the code.

Source location retrieval

A source location looks like this:

type Location struct {
    File   string
    Line   int
    Column int
}
func (loc Location) String() string {
    return fmt.Sprintf("%s:%d:%d", loc.File, loc.Line, loc.Column)
}
Enter fullscreen mode Exit fullscreen mode

We will also introduce an auxiliary type to represent paths in YAML. In YAML, there are two kinds of nested documents — arrays and objects.

some_array:
- hello
- world
some_object:
  foo: bar
Enter fullscreen mode Exit fullscreen mode

If we wanted to be able to refer to any subdocument, we could use something akin to JSON paths. In the example above, ["some\_array", 1] would then point to "word". But since we won’t support arrays in our proof-of-concept, we can get by just using an array of strings.

type Path []string
Enter fullscreen mode Exit fullscreen mode

One example of a path would be something like:

Path{"Resources", "PrivateSubnet", "Properties", "CidrBlock"}
Enter fullscreen mode Exit fullscreen mode

Now we can provide a convenience type to load YAML and tell us the Location of certain Paths.

type Source struct {
    file string
    root *yaml.Node
}
func NewSource(file string) (*Source, error) {
    bytes, err := ioutil.ReadFile(file)
    if err != nil {
        return nil, err
    }

    var root yaml.Node
    if err := yaml.Unmarshal(bytes, &root); err != nil {
        return nil, err
    }

    return &Source{file: file, root: &root}, nil
}
Enter fullscreen mode Exit fullscreen mode

Finding the source location of a Path comes down to walking a tree of YAML nodes:

func (source *Source) Location(path Path) *Location {
    cursor := source.root
    for len(path) > 0 {
        switch cursor.Kind {
        // Ignore multiple docs in our PoC.
        case yaml.DocumentNode:
            cursor = cursor.Content[0]
        case yaml.MappingNode:
            // Objects are stored as an array.
            // Content[2 * n] holds to the key and
            // Content[2 * n + 1] to the value.
            for i := 0; i < len(cursor.Content); i += 2 {
                if cursor.Content[i].Value == path[0] {
                    cursor = cursor.Content[i+1]
                    path = path[1:]
                }
            }
        }
    }
    return &Location{
        File:   source.file,
        Line:   cursor.Line,
        Column: cursor.Column,
    }
}
Enter fullscreen mode Exit fullscreen mode

Sets and trees of paths

With that out of the way, we’ve reduced the problem from automatically inferring source locations that are used in a policy to automatically inferring attribute paths.

This is also significant for other reasons — for example, Snyk can apply the same policies to IaC resources as well as resources discovered through cloud scans, the latter of which don’t really have meaningful source locations, but they do have meaningful attribute paths!

So, we want to define sets of attribute paths. Since paths are backed by arrays, we unfortunately can’t use something like map[Path]struct{} as a set in Go.

Instead, we will need to store these in a recursive tree.

type PathTree map[string]PathTree
Enter fullscreen mode Exit fullscreen mode

This representation has other advantages. In general, we only care about the longest paths that a policy uses, since they are more specific. Our example policy is using Path{"Resources", "PrivateSubnet", "Properties"} as well as `Path{"Resources", "PrivateSubnet", "Properties", "CidrBlock"} — we only care about the latter.

We’ll define a recursive method to insert a Path into our tree:


func (tree PathTree) Insert(path Path) {
if len(path) > 0 {
if _, ok := tree[path[0]]; !ok {
tree[path[0]] = map[string]PathTree{}
}
tree[path[0]].Insert(path[1:])
}
}

…as well as a way to get a list of Paths back out. This does a bit of unnecessary allocation, but we can live with that.


func (tree PathTree) List() []Path {
if len(tree) == 0 {
// Return the empty path.
return []Path{{}}
} else {
out := []Path{}
for k, child := range tree {
// Prepend
kto every child path.
for _, childPath := range child.List() {
path := Path{k}
path = append(path, childPath...)
out = append(out, path)
}
}
return out
}
}

We now have a way to nicely store the Paths that were used by a policy, and we have a way to convert those into source locations.

Static vs runtime analysis

The next question is to figure out which Paths in a given input are used by a policy, and then Insert those into the tree.

This is not an easy question, as the code may manipulate the input in different ways before using the paths. We’ll need to look through user-defined ( has\_bad\_subnet) as well as built-in functions ( object.get), just to illustrate one of the possible obstacles:

`
has_bad_subnet(props) {
[_, mask] = split(props.CidrBlock, "/")
to_number(mask) < 24
}

deny[resourceId] {
resource := input.Resources[resourceId]
resource.Type == "AWS::EC2::Subnet"
has_bad_subnet(object.get(resource, "Properties", {}))
}
`

Fortunately, we are not alone in this since people have been curious about what programs do since the first program was written. There are generally two ways of answering a question like that about a piece of code:

  • Static analysis: Try to answer by looking at the syntax tree, types, and other static information that we can retrieve from (or add to) the OPA interpreter. The advantage is that we don’t need to run this policy, which is great if we don’t trust the policy authors. The downside is that static analysis techniques will usually result in some false negatives as well as false positives.
  • Runtime analysis: Trace the execution of specific policies, and infer from what Paths are being used by looking at runtime information. The downside here is that we actually need to run the policy, and adding this analysis may slow down policy evaluation.

We tried both approaches but decided to go with the latter since we found it much easier to implement reliably, and the performance overhead was negligible. It’s also worth mentioning that this is not a binary choice — you could take a hybrid approach and combine the two.

OPA provides a Tracer interface that can be used to receive events about what the interpreter is doing. A common use case for tracers is to send metrics or debug information to some centralized log. Today, we’ll be using it for something else, though.


type locationTracer struct {
tree PathTree
}
func newLocationTracer() *locationTracer {
return &locationTracer{tree: PathTree{}}
}
func (tracer *locationTracer) Enabled() bool {
return true
}

Tracing usage of terms

Rego is an expressive language. Even though some desugaring happens to reduce it to a simpler format for the interpreter, there are still a fair number of events.

We are only interested in two of them. We consider a value used if:

  1. It is unified (you can think of this as assigned, we won’t go in detail) against another expression, such as:


x = input.Foo

This also covers == and :=. Since this is a test that can fail, we can state we used the left-hand side as well as the right-hand side.

  1. It is used as an argument to a built-in function, like:


regex.match("/24$", input.cidr)

While Rego borrows some concepts from lazy languages, arguments to built-in functions are always completely grounded before the built-in is invoked. Therefore, we can say we used all arguments supplied to the built-in.

  1. It is used as a standalone expression, such as:


volume.encrypted

This is commonly used to evaluate booleans and check that attributes exist.

Now, it’s time to implement. We match two events and delegate to a specific function to make the code a bit more readable:


func (tracer *locationTracer) Trace(event *topdown.Event) {
switch event.Op {
case topdown.UnifyOp:
tracer.traceUnify(event)
case topdown.EvalOp:
tracer.traceEval(event)
}
}

We’ll handle the insertion into our PathTree later in an auxiliary function called used(\*ast.Term). For now, let’s mark both the left- and right-hand sides to the unification as used:


func (tracer *locationTracer) traceUnify(event *topdown.Event) {
if expr, ok := event.Node.(*ast.Expr); ok {
operands := expr.Operands()
if len(operands) == 2 {
// Unification (1)
tracer.used(event.Plug(operands[0]))
tracer.used(event.Plug(operands[1]))
}
}
}

event.Plug is a helper to fill in variables with their actual values.

An EvalOp event covers both (2) and (3) mentioned above. In the case of a built-in function, we will have an array of terms, of which the first element is the function, and the remaining elements are the arguments. We can check that we’re dealing with a built-in function by looking in ast.BuiltinMap.

The case for a standalone expression is easy.


func (tracer *locationTracer) traceEval(event *topdown.Event) {
if expr, ok := event.Node.(*ast.Expr); ok {
switch terms := expr.Terms.(type) {
case []*ast.Term:
if len(terms) < 1 {
// I'm not sure what this is, but it's definitely
// not a built-in function application.
break
}
operator := terms[0]
if _, ok := ast.BuiltinMap[operator.String()]; ok {
// Built-in function call (2)
for _, term := range terms[1:] {
tracer.used(event.Plug(term))
}
}
case *ast.Term:
// Standalone expression (3)
tracer.used(event.Plug(terms))
}
}
}

Annotating terms

When we try to implement used(\*ast.Term), the next question poses itself — given a term, how do we map it to a Path in the input?

One option would be to search the input document for matching terms. But that would produce a lot of false positives, since a given string like 10.0.0.0/24 may appear many times in the input!

Instead, we will annotate all terms with their path. Terms in OPA can contain some metadata, including the location in the Rego source file. We can reuse this field to store an input Path. This is a bit hacky, but with some squinting, we are morally on the right side, since the field is meant to store locations.

The following snippet illustrates how we want to annotate the first few lines of our CloudFormation template:

`
Resources: # ["Resources"]
Vpc: # ["Resources", "Vpc"]
Type: AWS::EC2::VPC # ["Resources", "Vpc", "Type"]
Properties: # ["Resources", "Vpc", "Properties"]
CidrBlock: 10.0.0.0/16 # ["Resources", "Vpc", "Properties", "CidrBlock"]

`

annotate implements a recursive traversal to determine the Path at each node in the value. For conciseness, we only support objects and leave sets and arrays out.


func annotate(path Path, term *ast.Term) {
// Annotate current term by setting location.
if bytes, err := json.Marshal(path); err == nil {
term.Location = &ast.Location{}
term.Location.File = "path:" + string(bytes)
}
// Recursively annotate children.
switch value := term.Value.(type) {
case ast.Object:
for _, key := range value.Keys() {
if str, ok := key.Value.(ast.String); ok {
path = append(path, string(str))
annotate(path, value.Get(key))
path = path[:len(path)-1]
}
}
}
}

With this annotation in place, it’s easy to write used(\*ast.Term). The only thing to keep in mind is that not all values are annotated. We only do that for those coming from the input document, not, for example, literals embedded in the Rego source code.


func (tracer *locationTracer) used(term *ast.Term) {
if term.Location != nil {
val := strings.TrimPrefix(term.Location.File, "path:")
if len(val) != len(term.Location.File) {
// Only when we stripped a "path" suffix.
var path Path
json.Unmarshal([]byte(val), &path)
tracer.tree.Insert(path)
}
}
}

Wrapping up

That’s it, folks! We skipped over a lot of details, such as arrays and how to apply this to a more complex IaC language like HCL.

In addition to that, we’re also marking the Type attributes as used, since we check those in our policy. This isn’t great, and as an alternative, we try to provide a resources-oriented Rego API instead. But that’s beyond the scope of this example.

If you’re interested in learning more about any of these features, we recommend checking out snyk/policy-engine for the core implementation or the updated Snyk IaC, which comes with this and a whole host of other features, including an exhaustive rule bundle.

What follows is a main function to tie everything together and print out some debug information. It’s mostly just wrapping up the primitives we defined so far, and running it on an example. But let’s include it to make this post function as a reproducible standalone example.

`
func infer(policy string, file string) error {
source, err := NewSource(file)
if err != nil {
return err
}

bytes, err := ioutil.ReadFile(file)
if err != nil {
    return err
}

var node yaml.Node
if err := yaml.Unmarshal(bytes, &node); err != nil {
    return err
}

var doc interface{}
if err := yaml.Unmarshal(bytes, &doc); err != nil {
    return err
}

input, err := ast.InterfaceToValue(doc)
if err != nil {
    return err
}

annotate(Path{}, ast.NewTerm(input))
if bytes, err = ioutil.ReadFile(policy); err != nil {
    return err
}

tracer := newLocationTracer()
results, err := rego.New(
    rego.Module(policy, string(bytes)),
    rego.ParsedInput(input),
    rego.Query("data.policy.deny"),
    rego.Tracer(tracer),
).Eval(context.Background())
if err != nil {
    return err
}

fmt.Fprintf(os.Stderr, "Results: %v\n", results)
for _, path := range tracer.tree.List() {
    fmt.Fprintf(os.Stderr, "Location: %s\n", source.Location(path).String())
}
return nil
Enter fullscreen mode Exit fullscreen mode

}
func main() {
if err := infer("policy.rego", "template.yml"); err != nil {
fmt.Fprintf(os.Stderr, "%s\n", err)
os.Exit(1)
}
}
`

The full code for this PoC can be found in this gist.

. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .