Regoによるソースコード位置の自動特定

1Resources:
2  Vpc:
3    Type: AWS::EC2::VPC
4    Properties:
5      CidrBlock: 10.0.0.0/16
6  PublicSubnet:
7    Type: AWS::EC2::Subnet
8    Properties:
9      VpcId:
10        Ref: Vpc
11      CidrBlock: 10.0.0.0/24
12  PrivateSubnet:
13    Type: AWS::EC2::Subnet
14    Properties:
15      VpcId:
16        Ref: Vpc
17      CidrBlock: 10.0.128.0/20

1package policy
2
3deny[resourceId] {
4    resource := input.Resources[resourceId]
5    resource.Type == "AWS::EC2::Subnet"
6    [_, mask] = split(resource.Properties.CidrBlock, "/")
7    to_number(mask) < 24
8}

1type Location struct {
2    File   string
3    Line   int
4    Column int
5}
6func (loc Location) String() string {
7    return fmt.Sprintf("%s:%d:%d", loc.File, loc.Line, loc.Column)
8}

1some_array:
2- hello
3- world
4some_object:
5  foo: bar

1type Source struct {
2    file string
3    root *yaml.Node
4}
5func NewSource(file string) (*Source, error) {
6    bytes, err := ioutil.ReadFile(file)
7    if err != nil {
8        return nil, err
9    }
10
11    var root yaml.Node
12    if err := yaml.Unmarshal(bytes, &root); err != nil {
13        return nil, err
14    }
15
16    return &Source{file: file, root: &root}, nil
17}

1func (source *Source) Location(path Path) *Location {
2    cursor := source.root
3    for len(path) > 0 {
4        switch cursor.Kind {
5        // Ignore multiple docs in our PoC.
6        case yaml.DocumentNode:
7            cursor = cursor.Content[0]
8        case yaml.MappingNode:
9            // Objects are stored as an array.
10            // Content[2 * n] holds to the key and
11            // Content[2 * n + 1] to the value.
12            for i := 0; i < len(cursor.Content); i += 2 {
13                if cursor.Content[i].Value == path[0] {
14                    cursor = cursor.Content[i+1]
15                    path = path[1:]
16                }
17            }
18        }
19    }
20    return &Location{
21        File:   source.file,
22        Line:   cursor.Line,
23        Column: cursor.Column,
24    }
25}

1func (tree PathTree) Insert(path Path) {
2    if len(path) > 0 {
3        if _, ok := tree[path[0]]; !ok {
4            tree[path[0]] = map[string]PathTree{}
5        }
6        tree[path[0]].Insert(path[1:])
7    }
8}

1func (tree PathTree) List() []Path {
2    if len(tree) == 0 {
3        // Return the empty path.
4        return []Path{{}}
5    } else {
6        out := []Path{}
7        for k, child := range tree {
8            // Prepend `k` to every child path.
9            for _, childPath := range child.List() {
10                path := Path{k}
11                path = append(path, childPath...)
12                out = append(out, path)
13            }
14        }
15        return out
16    }
17}

1has_bad_subnet(props) {
2    [_, mask] = split(props.CidrBlock, "/")
3    to_number(mask) < 24
4}
5
6deny[resourceId] {
7    resource := input.Resources[resourceId]
8    resource.Type == "AWS::EC2::Subnet"
9    has_bad_subnet(object.get(resource, "Properties", {}))
10}

1type locationTracer struct {
2    tree PathTree
3}
4func newLocationTracer() *locationTracer {
5    return &locationTracer{tree: PathTree{}}
6}
7func (tracer *locationTracer) Enabled() bool {
8    return true
9}

1func (tracer *locationTracer) Trace(event *topdown.Event) {
2    switch event.Op {
3    case topdown.UnifyOp:
4        tracer.traceUnify(event)
5    case topdown.EvalOp:
6        tracer.traceEval(event)
7    }
8}

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

1func (tracer *locationTracer) traceEval(event *topdown.Event) {
2    if expr, ok := event.Node.(*ast.Expr); ok {
3        switch terms := expr.Terms.(type) {
4        case []*ast.Term:
5            if len(terms) < 1 {
6                // I'm not sure what this is, but it's definitely
7                // not a built-in function application.
8                break
9            }
10            operator := terms[0]
11            if _, ok := ast.BuiltinMap[operator.String()]; ok {
12                // Built-in function call (2)
13                for _, term := range terms[1:] {
14                    tracer.used(event.Plug(term))
15                }
16            }
17        case *ast.Term:
18            // Standalone expression (3)
19            tracer.used(event.Plug(terms))
20        }
21    }
22}

1func annotate(path Path, term *ast.Term) {
2    // Annotate current term by setting location.
3    if bytes, err := json.Marshal(path); err == nil {
4        term.Location = &ast.Location{}
5        term.Location.File = "path:" + string(bytes)
6    }
7    // Recursively annotate children.
8    switch value := term.Value.(type) {
9    case ast.Object:
10        for _, key := range value.Keys() {
11            if str, ok := key.Value.(ast.String); ok {
12                path = append(path, string(str))
13                annotate(path, value.Get(key))
14                path = path[:len(path)-1]
15            }
16        }
17    }
18}

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

1func infer(policy string, file string) error {
2    source, err := NewSource(file)
3    if err != nil {
4        return err
5    }
6
7    bytes, err := ioutil.ReadFile(file)
8    if err != nil {
9        return err
10    }
11
12    var node yaml.Node
13    if err := yaml.Unmarshal(bytes, &node); err != nil {
14        return err
15    }
16
17    var doc interface{}
18    if err := yaml.Unmarshal(bytes, &doc); err != nil {
19        return err
20    }
21
22    input, err := ast.InterfaceToValue(doc)
23    if err != nil {
24        return err
25    }
26
27    annotate(Path{}, ast.NewTerm(input))
28    if bytes, err = ioutil.ReadFile(policy); err != nil {
29        return err
30    }
31
32    tracer := newLocationTracer()
33    results, err := rego.New(
34        rego.Module(policy, string(bytes)),
35        rego.ParsedInput(input),
36        rego.Query("data.policy.deny"),
37        rego.Tracer(tracer),
38    ).Eval(context.Background())
39    if err != nil {
40        return err
41    }
42
43    fmt.Fprintf(os.Stderr, "Results: %v\n", results)
44    for _, path := range tracer.tree.List() {
45        fmt.Fprintf(os.Stderr, "Location: %s\n", source.Location(path).String())
46    }
47    return nil
48}
49func main() {
50    if err := infer("policy.rego", "template.yml"); err != nil {
51        fmt.Fprintf(os.Stderr, "%s\n", err)
52        os.Exit(1)
53    }
54}