From 234292e6986e7ea6d24a335746ff00bf60651a25 Mon Sep 17 00:00:00 2001
From: Evans Mungai <evans@replicated.com>
Date: Thu, 23 Nov 2023 19:38:49 +0000
Subject: [PATCH] fix(support-bundle): do not download specs from urls twice

---
 internal/specs/specs.go      | 39 ++++++++++++++++++++++++------------
 internal/specs/specs_test.go | 23 +++++++++++++++++++++
 2 files changed, 49 insertions(+), 13 deletions(-)

diff --git a/internal/specs/specs.go b/internal/specs/specs.go
index 1fd50c55a..b66a82007 100644
--- a/internal/specs/specs.go
+++ b/internal/specs/specs.go
@@ -170,22 +170,35 @@ func LoadFromCLIArgs(ctx context.Context, client kubernetes.Interface, args []st
 					return nil, types.NewExitCodeError(constants.EXIT_CODE_SPEC_ISSUES, fmt.Errorf("%s is not a URL and was not found (err %s)", v, err))
 				}
 
-				// Download preflight specs
-				rawSpec, err := downloadFromHttpURL(ctx, v, map[string]string{"User-Agent": "Replicated_Preflight/v1beta2"})
+				parsedURL, err := url.ParseRequestURI(v)
 				if err != nil {
-					return nil, err
+					return nil, types.NewExitCodeError(constants.EXIT_CODE_SPEC_ISSUES, err)
 				}
-				rawSpecs = append(rawSpecs, rawSpec)
-
-				// Download support bundle specs
-				rawSpec, err = downloadFromHttpURL(ctx, v, map[string]string{
-					"User-Agent":         "Replicated_Troubleshoot/v1beta1",
-					"Bundle-Upload-Host": fmt.Sprintf("%s://%s", u.Scheme, u.Host),
-				})
-				if err != nil {
-					return nil, err
+				if parsedURL.Host == "kots.io" {
+					// Download support bundle specs from kots.io which needs this header
+					rawSpec, err := downloadFromHttpURL(ctx, v, map[string]string{
+						"User-Agent":         "Replicated_Troubleshoot/v1beta1",
+						"Bundle-Upload-Host": fmt.Sprintf("%s://%s", u.Scheme, u.Host), // Not sure why this is needed
+					})
+					if err != nil {
+						return nil, err
+					}
+					rawSpecs = append(rawSpecs, rawSpec)
+				} else {
+					rawSpec, err := downloadFromHttpURL(ctx, v, nil)
+					if err != nil {
+						return nil, err
+					}
+					rawSpecs = append(rawSpecs, rawSpec)
 				}
-				rawSpecs = append(rawSpecs, rawSpec)
+
+				// TODO: I could not find which host uses this header
+				// Download preflight specs
+				// rawSpec, err := downloadFromHttpURL(ctx, v, map[string]string{"User-Agent": "Replicated_Preflight/v1beta2"})
+				// if err != nil {
+				// 	return nil, err
+				// }
+				// rawSpecs = append(rawSpecs, rawSpec)
 			}
 		}
 	}
diff --git a/internal/specs/specs_test.go b/internal/specs/specs_test.go
index 026289fda..b32df9b29 100644
--- a/internal/specs/specs_test.go
+++ b/internal/specs/specs_test.go
@@ -3,6 +3,8 @@ package specs
 import (
 	"context"
 	"fmt"
+	"net/http"
+	"net/http/httptest"
 	"reflect"
 	"testing"
 
@@ -10,6 +12,7 @@ import (
 	"github.com/replicatedhq/troubleshoot/internal/testutils"
 	troubleshootv1beta2 "github.com/replicatedhq/troubleshoot/pkg/apis/troubleshoot/v1beta2"
 	"github.com/replicatedhq/troubleshoot/pkg/loader"
+	"github.com/spf13/viper"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	corev1 "k8s.io/api/core/v1"
@@ -201,3 +204,23 @@ spec:
 		},
 	}
 }
+
+func TestLoadFromURI(t *testing.T) {
+	m := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Write([]byte(`apiVersion: troubleshoot.sh/v1beta2
+apiVersion: troubleshoot.sh/v1beta2
+kind: HostCollector
+metadata:
+  name: cpu
+spec:
+  collectors:
+    - cpu: {}
+`))
+	}))
+	defer m.Close()
+
+	client := testclient.NewSimpleClientset()
+	specs, err := LoadFromCLIArgs(context.Background(), client, []string{m.URL}, viper.New())
+	require.NoError(t, err)
+	require.Len(t, specs.HostCollectorsV1Beta2, 1)
+}