From b12fcfef1355a2b3540963566ba748fd6ee29f92 Mon Sep 17 00:00:00 2001
From: Justin Hiemstra <jhiemstra@wisc.edu>
Date: Wed, 2 Oct 2024 16:48:45 +0000
Subject: [PATCH] Handle more errors in DiscoverFederation and test those
 errors

---
 pelican_url/discovery.go      | 16 ++++++++++++++--
 pelican_url/discovery_test.go | 28 ++++++++++++++++++++++++++++
 2 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/pelican_url/discovery.go b/pelican_url/discovery.go
index 90aac1b66..f6c833bd8 100644
--- a/pelican_url/discovery.go
+++ b/pelican_url/discovery.go
@@ -278,8 +278,18 @@ func startMetadataQuery(ctx context.Context, httpClient *http.Client, ua string,
 // This function is for discovering federations as specified by a url during a pelican:// transfer.
 // this does not populate global fields and is more temporary per url
 func DiscoverFederation(ctx context.Context, httpClient *http.Client, ua string, discoveryUrl *url.URL) (metadata FederationDiscovery, err error) {
-	log.Debugln("Performing federation service discovery for against", discoveryUrl.String())
+	if discoveryUrl.Path != "" {
+		return FederationDiscovery{}, errors.Errorf("Discovery URL must not have a path, but you provided '%s'", discoveryUrl.String())
+	}
+	if discoveryUrl.Host == "" {
+		return FederationDiscovery{}, errors.Errorf("Discovery URL must have a host, but you provided '%s'", discoveryUrl.String())
+	}
+	if discoveryUrl.Scheme != "https" && discoveryUrl.Scheme != "http" {
+		return FederationDiscovery{}, errors.Errorf("Discovery URL must use http/s, but you provided '%s'", discoveryUrl.String())
+	}
+
 	discoveryUrl.Path = PelicanDiscoveryPath
+	log.Debugln("Performing federation service discovery for against", discoveryUrl.String())
 
 	var result *http.Response
 	for idx := 1; idx <= 3; idx++ {
@@ -343,7 +353,9 @@ func (p *PelicanURL) PopulateFedInfo(opts ...DiscoveryOption) error {
 		opt(options)
 	}
 
-	discoveryUrl := &url.URL{Scheme: "https", Path: PelicanDiscoveryPath}
+	// We don't set the discovery path here, because that's handled by DiscoverFederation. All we really care to pass to that
+	// is the scheme/host.
+	discoveryUrl := &url.URL{Scheme: "https"}
 	normedScheme := normalizeScheme(p.Scheme)
 	if normedScheme == OsdfScheme || normedScheme == StashScheme {
 		// Prefer OSDF discovery host, but allow someone to overwrite if they really want to
diff --git a/pelican_url/discovery_test.go b/pelican_url/discovery_test.go
index efc74739d..31fe14e44 100644
--- a/pelican_url/discovery_test.go
+++ b/pelican_url/discovery_test.go
@@ -99,6 +99,34 @@ func TestDiscoverFederation(t *testing.T) {
 		assert.Equal(t, "https://broker.com", fedInfo.BrokerEndpoint, "Unexpected BrokerEndpoint")
 	})
 
+	t.Run("TestMalformedDiscUrls", func(t *testing.T) {
+		tr := &http.Transport{
+			TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
+		}
+		client := &http.Client{Transport: tr}
+
+		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+		defer cancel()
+
+		malformedUrl, _ := url.Parse("badurl")
+		_, err = DiscoverFederation(ctx, client, "test-ua", malformedUrl)
+		require.Error(t, err)
+		require.Contains(t, err.Error(), "Discovery URL must not have a path, but you provided 'badurl'")
+
+		// Url with invalid scheme
+		malformedUrl, _ = url.Parse("foobar://badurl")
+		_, err = DiscoverFederation(ctx, client, "test-ua", malformedUrl)
+		require.Error(t, err)
+		require.Contains(t, err.Error(), "Discovery URL must use http/s, but you provided 'foobar://badurl'")
+
+		// Url with no host
+		malformedUrl, _ = url.Parse("https://")
+		_, err = DiscoverFederation(ctx, client, "test-ua", malformedUrl)
+		require.Error(t, err)
+		// Not totally sure why URL parsing results in `https:` instead of `https://` but that's what we get
+		require.Contains(t, err.Error(), "Discovery URL must have a host, but you provided 'https:'")
+	})
+
 	t.Run("TestMetadataDiscoveryTimeout", func(t *testing.T) {
 		tr := &http.Transport{
 			TLSClientConfig: &tls.Config{InsecureSkipVerify: true},