refactor: break up LoadFileOrURL by scheme

Signed-off-by: Zachary Newman <z@znewman.net>
sigstore · Apr 24, 2022 · 89dd9a9 · 89dd9a9
1 parent c3e0952
commit 89dd9a9
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 10 deletions.
diff --git a/pkg/blob/load.go b/pkg/blob/load.go
@@ -15,6 +15,7 @@
 package blob
 
 import (
+	"fmt"
 	"io"
 	"net/http"
 	"os"
@@ -25,16 +26,25 @@ import (
 func LoadFileOrURL(fileRef string) ([]byte, error) {
 	var raw []byte
 	var err error
-	if strings.HasPrefix(fileRef, "http://") || strings.HasPrefix(fileRef, "https://") {
-		// #nosec G107
-		resp, err := http.Get(fileRef)
-		if err != nil {
-			return nil, err
-		}
-		defer resp.Body.Close()
-		raw, err = io.ReadAll(resp.Body)
-		if err != nil {
-			return nil, err
+	parts := strings.SplitAfterN(fileRef, "://", 2)
+	if len(parts) == 2 {
+		scheme := parts[0]
+		switch scheme {
+		case "http://":
+			fallthrough
+		case "https://":
+			// #nosec G107
+			resp, err := http.Get(fileRef)
+			if err != nil {
+				return nil, err
+			}
+			defer resp.Body.Close()
+			raw, err = io.ReadAll(resp.Body)
+			if err != nil {
+				return nil, err
+			}
+		default:
+			return nil, fmt.Errorf("loading URL: unrecognized scheme: %s", scheme)
 		}
 	} else {
 		raw, err = os.ReadFile(filepath.Clean(fileRef))

diff --git a/pkg/blob/load_test.go b/pkg/blob/load_test.go
@@ -61,4 +61,9 @@ func TestLoadFileOrURL(t *testing.T) {
 	} else if !bytes.Equal(actual, data) {
 		t.Errorf("LoadFileOrURL(HTTP) = '%s'; want '%s'", actual, data)
 	}
+
+	_, err = LoadFileOrURL("invalid://url")
+	if err == nil {
+		t.Error("LoadFileOrURL(): expected error")
+	}
 }