From 520e10cf8b14737577890d638e04a5444bd2ec78 Mon Sep 17 00:00:00 2001
From: Ian Duffy <iduffy@cloudsmith.io>
Date: Mon, 6 Jan 2025 20:18:12 +0000
Subject: [PATCH] Fix SAML import, allow for arbitrary headers

---
 cloudsmith/provider.go              |  9 ++++++++-
 cloudsmith/provider_config.go       | 21 +++++++++++++++++++--
 cloudsmith/resource_saml_auth.go    |  2 +-
 cloudsmith/resource_service_test.go |  2 +-
 go.sum                              |  2 --
 5 files changed, 29 insertions(+), 7 deletions(-)

diff --git a/cloudsmith/provider.go b/cloudsmith/provider.go
index c7487f5..c8fae78 100644
--- a/cloudsmith/provider.go
+++ b/cloudsmith/provider.go
@@ -27,6 +27,12 @@ func Provider() *schema.Provider {
 				Optional:    true,
 				DefaultFunc: schema.EnvDefaultFunc("CLOUDSMITH_API_HOST", "https://api.cloudsmith.io/v1"),
 			},
+			"headers": {
+				Type:        schema.TypeMap,
+				Elem:        &schema.Schema{Type: schema.TypeString},
+				Description: "Additional HTTP headers to include in API requests",
+				Optional:    true,
+			},
 		},
 		DataSourcesMap: map[string]*schema.Resource{
 			"cloudsmith_namespace":             dataSourceNamespace(),
@@ -72,8 +78,9 @@ func Provider() *schema.Provider {
 		apiHost := requiredString(d, "api_host")
 		apiKey := requiredString(d, "api_key")
 		userAgent := fmt.Sprintf("(%s %s) Terraform/%s", runtime.GOOS, runtime.GOARCH, terraformVersion)
+		headers := d.Get("headers").(map[string]interface{})
 
-		return newProviderConfig(apiHost, apiKey, userAgent)
+		return newProviderConfig(apiHost, apiKey, headers, userAgent)
 	}
 
 	return p
diff --git a/cloudsmith/provider_config.go b/cloudsmith/provider_config.go
index 92f6860..c029eb9 100644
--- a/cloudsmith/provider_config.go
+++ b/cloudsmith/provider_config.go
@@ -3,6 +3,7 @@ package cloudsmith
 import (
 	"context"
 	"errors"
+	"fmt"
 	"net/http"
 
 	"github.com/cloudsmith-io/cloudsmith-api-go"
@@ -20,17 +21,21 @@ type providerConfig struct {
 	APIClient *cloudsmith.APIClient
 }
 
-func newProviderConfig(apiHost, apiKey, userAgent string) (*providerConfig, diag.Diagnostics) {
+func newProviderConfig(apiHost string, apiKey string, headers map[string]interface{}, userAgent string) (*providerConfig, diag.Diagnostics) {
 	if apiKey == "" {
 		return nil, diag.FromErr(errMissingCredentials)
 	}
 
 	httpClient := http.DefaultClient
-	httpClient.Transport = logging.NewSubsystemLoggingHTTPTransport("Cloudsmith", http.DefaultTransport)
+	httpClient.Transport = logging.NewSubsystemLoggingHTTPTransport("Cloudsmith", &headerTransport{
+		headers: headers,
+		rt:      http.DefaultTransport,
+	})
 
 	config := cloudsmith.NewConfiguration()
 	config.Debug = logging.IsDebugOrHigher()
 	config.HTTPClient = httpClient
+
 	config.Servers = cloudsmith.ServerConfigurations{
 		{URL: apiHost},
 	}
@@ -58,3 +63,15 @@ func (pc *providerConfig) GetAPIKey() string {
 	apiKeys, _ := pc.Auth.Value(cloudsmith.ContextAPIKeys).(map[string]cloudsmith.APIKey)
 	return apiKeys["apikey"].Key
 }
+
+type headerTransport struct {
+	headers map[string]interface{}
+	rt      http.RoundTripper
+}
+
+func (t *headerTransport) RoundTrip(req *http.Request) (*http.Response, error) {
+	for k, v := range t.headers {
+		req.Header.Add(k, fmt.Sprint(v))
+	}
+	return t.rt.RoundTrip(req)
+}
diff --git a/cloudsmith/resource_saml_auth.go b/cloudsmith/resource_saml_auth.go
index b9244aa..f2b4e72 100644
--- a/cloudsmith/resource_saml_auth.go
+++ b/cloudsmith/resource_saml_auth.go
@@ -185,7 +185,7 @@ func generateSAMLAuthID(organization string, samlAuth *cloudsmith.OrganizationSA
 		data += fmt.Sprintf("-%t", samlAuth.GetSamlAuthEnabled())
 		data += fmt.Sprintf("-%t", samlAuth.GetSamlAuthEnforced())
 
-		if url, hasURL := samlAuth.GetSamlMetadataUrlOk(); hasURL {
+		if url, hasURL := samlAuth.GetSamlMetadataUrlOk(); url != nil && hasURL {
 			data += fmt.Sprintf("-%s", *url)
 		}
 
diff --git a/cloudsmith/resource_service_test.go b/cloudsmith/resource_service_test.go
index b263466..feb2a62 100644
--- a/cloudsmith/resource_service_test.go
+++ b/cloudsmith/resource_service_test.go
@@ -50,7 +50,7 @@ func TestAccService_basic(t *testing.T) {
 				Check: resource.ComposeTestCheckFunc(
 					testAccServiceCheckExists("cloudsmith_service.test"),
 					resource.TestCheckResourceAttrSet("cloudsmith_service.test", "team.#"),
-          resource.TestMatchTypeSetElemNestedAttrs("cloudsmith_service.test", "team.*", map[string]*regexp.Regexp{
+					resource.TestMatchTypeSetElemNestedAttrs("cloudsmith_service.test", "team.*", map[string]*regexp.Regexp{
 						"slug": regexp.MustCompile("^tf-test-team-svc(-[^2].*)?$"),
 						"role": regexp.MustCompile("^Member$"),
 					}),
diff --git a/go.sum b/go.sum
index 35a0aaf..0265a5e 100644
--- a/go.sum
+++ b/go.sum
@@ -19,8 +19,6 @@ github.com/apparentlymart/go-textseg/v13 v13.0.0/go.mod h1:ZK2fH7c4NqDTLtiYLvIkE
 github.com/armon/go-socks5 v0.0.0-20160902184237-e75332964ef5/go.mod h1:wHh0iHkYZB8zMSxRWpUBQtwG5a7fFgvEO+odwuTv2gs=
 github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU=
 github.com/client9/misspell v0.3.4/go.mod h1:qj6jICC3Q7zFZvVWo7KLAzC3yx5G7kyvSDkc90ppPyw=
-github.com/cloudsmith-io/cloudsmith-api-go v0.0.41 h1:7bPbyD3FIN8PmmZQg43O/y8OTkFFZ2ZxcSbrMcdrjHE=
-github.com/cloudsmith-io/cloudsmith-api-go v0.0.41/go.mod h1:xHTDgU8TyMXiWUZdn2OQR7AaZdz9M5SrNenyQU44zGQ=
 github.com/cloudsmith-io/cloudsmith-api-go v0.0.42-0.20241129202450-bd5381591ce5 h1:vU6M9CeLmTU9xJj5G7pCQ9F8o+zl4pSfKnmGxB+miLE=
 github.com/cloudsmith-io/cloudsmith-api-go v0.0.42-0.20241129202450-bd5381591ce5/go.mod h1:xHTDgU8TyMXiWUZdn2OQR7AaZdz9M5SrNenyQU44zGQ=
 github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=