From d4556d4b85a5d37ec857b42ecb0fb8ceb4b9ddd0 Mon Sep 17 00:00:00 2001
From: Evans Mungai <evans@replicated.com>
Date: Wed, 29 Nov 2023 16:58:28 +0000
Subject: [PATCH] fix(support-bundle): correct user-agent bevahiour

Only kots.io requires a user agent when downloading troubleshoot specs
---
 internal/specs/specs.go      | 30 +++++++++++++++++-------------
 internal/specs/specs_test.go | 23 +++++++++++++++++++++++
 2 files changed, 40 insertions(+), 13 deletions(-)

diff --git a/internal/specs/specs.go b/internal/specs/specs.go
index 1fd50c55a..9231cd1f7 100644
--- a/internal/specs/specs.go
+++ b/internal/specs/specs.go
@@ -170,22 +170,26 @@ 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" {
+					// To download specs from kots.io, we need to set the User-Agent header
+					rawSpec, err := downloadFromHttpURL(ctx, v, map[string]string{
+						"User-Agent": "Replicated_Troubleshoot/v1beta1",
+					})
+					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)
 			}
 		}
 	}
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)
+}