From ac7452009bf7ca0fa8ee1de8807c792eabad405a Mon Sep 17 00:00:00 2001
From: Nikita Pivkin <nikita.pivkin@smartforce.io>
Date: Fri, 24 May 2024 09:07:04 +0600
Subject: [PATCH] feat(misconf): resolve tf module from OpenTofu compatible
 registry (#6743)

---
 .../terraform/parser/module_retrieval.go      | 14 ++++-----
 .../terraform/parser/resolvers/registry.go    | 22 +++++++++++--
 .../resolvers/registry_integration_test.go    | 31 +++++++++++++++++++
 3 files changed, 58 insertions(+), 9 deletions(-)
 create mode 100644 pkg/iac/scanners/terraform/parser/resolvers/registry_integration_test.go

diff --git a/pkg/iac/scanners/terraform/parser/module_retrieval.go b/pkg/iac/scanners/terraform/parser/module_retrieval.go
index cae84359498d..2ae6221afc73 100644
--- a/pkg/iac/scanners/terraform/parser/module_retrieval.go
+++ b/pkg/iac/scanners/terraform/parser/module_retrieval.go
@@ -5,21 +5,21 @@ import (
 	"fmt"
 	"io/fs"
 
-	resolvers2 "github.com/aquasecurity/trivy/pkg/iac/scanners/terraform/parser/resolvers"
+	"github.com/aquasecurity/trivy/pkg/iac/scanners/terraform/parser/resolvers"
 )
 
 type ModuleResolver interface {
-	Resolve(context.Context, fs.FS, resolvers2.Options) (filesystem fs.FS, prefix string, downloadPath string, applies bool, err error)
+	Resolve(context.Context, fs.FS, resolvers.Options) (filesystem fs.FS, prefix string, downloadPath string, applies bool, err error)
 }
 
 var defaultResolvers = []ModuleResolver{
-	resolvers2.Cache,
-	resolvers2.Local,
-	resolvers2.Remote,
-	resolvers2.Registry,
+	resolvers.Cache,
+	resolvers.Local,
+	resolvers.Remote,
+	resolvers.Registry,
 }
 
-func resolveModule(ctx context.Context, current fs.FS, opt resolvers2.Options) (filesystem fs.FS, sourcePrefix, downloadPath string, err error) {
+func resolveModule(ctx context.Context, current fs.FS, opt resolvers.Options) (filesystem fs.FS, sourcePrefix, downloadPath string, err error) {
 	opt.Debug("Resolving module '%s' with source: '%s'...", opt.Name, opt.Source)
 	for _, resolver := range defaultResolvers {
 		if filesystem, prefix, path, applies, err := resolver.Resolve(ctx, current, opt); err != nil {
diff --git a/pkg/iac/scanners/terraform/parser/resolvers/registry.go b/pkg/iac/scanners/terraform/parser/resolvers/registry.go
index 22778cc10230..93d80fa0af86 100644
--- a/pkg/iac/scanners/terraform/parser/resolvers/registry.go
+++ b/pkg/iac/scanners/terraform/parser/resolvers/registry.go
@@ -122,11 +122,29 @@ func (r *registryResolver) Resolve(ctx context.Context, target fs.FS, opt Option
 		return nil, "", "", true, err
 	}
 	defer func() { _ = resp.Body.Close() }()
-	if resp.StatusCode != http.StatusNoContent {
+
+	// OpenTofu may return 200 with body
+	switch resp.StatusCode {
+	case http.StatusOK:
+		// https://opentofu.org/docs/internals/module-registry-protocol/#sample-response-1
+		var downloadResponse struct {
+			Location string `json:"location"`
+		}
+		if err := json.NewDecoder(resp.Body).Decode(&downloadResponse); err != nil {
+			return nil, "", "", true, fmt.Errorf("failed to decode download response: %w", err)
+		}
+
+		opt.Source = downloadResponse.Location
+	case http.StatusNoContent:
+		opt.Source = resp.Header.Get("X-Terraform-Get")
+	default:
 		return nil, "", "", true, fmt.Errorf("unexpected status code: %d", resp.StatusCode)
 	}
 
-	opt.Source = resp.Header.Get("X-Terraform-Get")
+	if opt.Source == "" {
+		return nil, "", "", true, fmt.Errorf("no source was found for the registry at %s", hostname)
+	}
+
 	opt.Debug("Module '%s' resolved via registry to new source: '%s'", opt.Name, opt.Source)
 	opt.RelativePath = relativePath
 	filesystem, prefix, downloadPath, _, err = Remote.Resolve(ctx, target, opt)
diff --git a/pkg/iac/scanners/terraform/parser/resolvers/registry_integration_test.go b/pkg/iac/scanners/terraform/parser/resolvers/registry_integration_test.go
new file mode 100644
index 000000000000..e2d87104da2d
--- /dev/null
+++ b/pkg/iac/scanners/terraform/parser/resolvers/registry_integration_test.go
@@ -0,0 +1,31 @@
+package resolvers_test
+
+import (
+	"context"
+	"io/fs"
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/aquasecurity/trivy/pkg/iac/scanners/terraform/parser/resolvers"
+)
+
+func TestResolveModuleFromOpenTofuRegistry(t *testing.T) {
+	if testing.Short() {
+		t.Skip("skipping integration test in short mode")
+	}
+
+	fsys, _, path, _, err := resolvers.Registry.Resolve(context.Background(), nil, resolvers.Options{
+		Source:         "registry.opentofu.org/terraform-aws-modules/s3-bucket/aws",
+		RelativePath:   "test",
+		Name:           "bucket",
+		Version:        "4.1.2",
+		AllowDownloads: true,
+		SkipCache:      true,
+	})
+	require.NoError(t, err)
+
+	_, err = fs.Stat(fsys, filepath.Join(path, "main.tf"))
+	require.NoError(t, err)
+}