Add support for scoped data sources (#67)

Author: wata727

integration/data_sources/main.tf

@@ -5,3 +5,13 @@ data "aws_ami" "valid" {
 data "aws_ami" "invalid" {
   owners = ["amazon"]
 }
+
+check "scoped" {
+  data "aws_ami" "scoped_valid" {
+    owners = ["self"]
+  }
+
+  data "aws_ami" "scoped_invalid" {
+    owners = ["amazon"]
+  }
+}

integration/data_sources/policies/main_test.rego

@@ -4,12 +4,18 @@ import future.keywords
 mock_data_sources(type, schema, options) := terraform.mock_data_sources(type, schema, options, {"main.tf": `
 data "aws_ami" "main" {
   owners = ["amazon"]
+}
+
+check "scope" {
+  data "aws_ami" "scoped" {
+    owners = ["amazon"]
+  }
 }`})
 
 test_deny_other_ami_owners_passed if {
   issues := deny_other_ami_owners with terraform.data_sources as mock_data_sources
 
-  count(issues) == 1
+  count(issues) == 2
   issue := issues[_]
   issue.msg == "third-party AMI is not allowed"
 }

integration/data_sources/result.json

@@ -19,6 +19,26 @@
         }
       },
       "callers": []
+    },
+    {
+      "rule": {
+        "name": "opa_deny_other_ami_owners",
+        "severity": "error",
+        "link": "policies/main.rego:3"
+      },
+      "message": "third-party AMI is not allowed",
+      "range": {
+        "filename": "main.tf",
+        "start": {
+          "line": 15,
+          "column": 14
+        },
+        "end": {
+          "line": 15,
+          "column": 24
+        }
+      },
+      "callers": []
     }
   ],
   "errors": []

integration/data_sources/result_test.json

@@ -4,7 +4,7 @@
       "rule": {
         "name": "opa_test_deny_other_ami_owners_failed",
         "severity": "error",
-        "link": "policies/main_test.rego:17"
+        "link": "policies/main_test.rego:23"
       },
       "message": "test failed",
       "range": {

opa/functions.go

@@ -156,7 +156,80 @@ func dataSourcesFunc(runner tflint.Runner) *function3 {
 			},
 		},
 		Func: func(_ rego.BuiltinContext, dataType *ast.Term, schema *ast.Term, options *ast.Term) (*ast.Term, error) {
-			return typedBlockFunc(dataType, schema, options, "data", runner)
+			var typeName string
+			if err := ast.As(dataType.Value, &typeName); err != nil {
+				return nil, err
+			}
+			var schemaJSON map[string]any
+			if err := ast.As(schema.Value, &schemaJSON); err != nil {
+				return nil, err
+			}
+			innerSchema, tyMap, err := jsonToSchema(schemaJSON, map[string]cty.Type{}, "schema")
+			if err != nil {
+				return nil, err
+			}
+			var optionJSON map[string]string
+			if err := ast.As(options.Value, &optionJSON); err != nil {
+				return nil, err
+			}
+			option, err := jsonToOption(optionJSON)
+			if err != nil {
+				return nil, err
+			}
+
+			content, err := runner.GetModuleContent(&hclext.BodySchema{
+				Blocks: []hclext.BlockSchema{
+					{
+						Type:       "data",
+						LabelNames: []string{"type", "name"},
+						Body:       innerSchema,
+					},
+					{
+						Type:       "check",
+						LabelNames: []string{"name"},
+						Body: &hclext.BodySchema{
+							Blocks: []hclext.BlockSchema{
+								{
+									Type:       "data",
+									LabelNames: []string{"type", "name"},
+									Body:       innerSchema,
+								},
+							},
+						},
+					},
+				},
+			}, option.AsGetModuleContentOptions())
+			if err != nil {
+				return nil, err
+			}
+
+			blocks := []*hclext.Block{}
+			for _, block := range content.Blocks {
+				switch block.Type {
+				case "data":
+					// "*" is a special character that returns all blocks
+					if typeName == block.Labels[0] || typeName == "*" {
+						blocks = append(blocks, block)
+					}
+				case "check":
+					for _, inner := range block.Body.Blocks {
+						if typeName == inner.Labels[0] || typeName == "*" {
+							blocks = append(blocks, inner)
+						}
+					}
+				}
+			}
+
+			out, err := typedBlocksToJSON(blocks, tyMap, "schema", runner)
+			if err != nil {
+				return nil, err
+			}
+			v, err := ast.InterfaceToValue(out)
+			if err != nil {
+				return nil, err
+			}
+
+			return ast.NewTerm(v), nil
 		},
 	}
 }

opa/functions_test.go

@@ -350,6 +350,56 @@ data "aws_ami" "main" {
 				},
 			},
 		},
+		{
+			name: "scoped data source",
+			config: `
+check "scoped" {
+	data "aws_ami" "main" {
+		owners = ["self"]
+	}
+}`,
+			dataType: "aws_ami",
+			schema:   map[string]any{"owners": "list(string)"},
+			want: []map[string]any{
+				{
+					"type": "aws_ami",
+					"name": "main",
+					"config": map[string]any{
+						"owners": map[string]any{
+							"value":     []string{"self"},
+							"unknown":   false,
+							"sensitive": false,
+							"range": map[string]any{
+								"filename": "main.tf",
+								"start": map[string]int{
+									"line":   4,
+									"column": 12,
+									"byte":   54,
+								},
+								"end": map[string]int{
+									"line":   4,
+									"column": 20,
+									"byte":   62,
+								},
+							},
+						},
+					},
+					"decl_range": map[string]any{
+						"filename": "main.tf",
+						"start": map[string]int{
+							"line":   3,
+							"column": 2,
+							"byte":   19,
+						},
+						"end": map[string]int{
+							"line":   3,
+							"column": 23,
+							"byte":   40,
+						},
+					},
+				},
+			},
+		},
 	}
 
 	for _, test := range tests {