From 0a944d8f0355bac712028eae289c815ed3e180c8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Adolfo=20Garc=C3=ADa=20Veytia=20=28Puerco=29?=
 <puerco@stacklok.com>
Date: Tue, 23 Apr 2024 15:19:15 -0600
Subject: [PATCH] Add mediatype string builder
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This commit adds a mediatype build string to generate media type strings

Signed-off-by: Adolfo García Veytia (Puerco) <puerco@stacklok.com>
---
 pkg/bundle/bundle.go      | 27 ++++++++++++++++++++++-----
 pkg/bundle/bundle_test.go | 30 ++++++++++++++++++++++++++++++
 2 files changed, 52 insertions(+), 5 deletions(-)

diff --git a/pkg/bundle/bundle.go b/pkg/bundle/bundle.go
index 376895c7..b8e4f812 100644
--- a/pkg/bundle/bundle.go
+++ b/pkg/bundle/bundle.go
@@ -42,6 +42,8 @@ var ErrMissingEnvelope = fmt.Errorf("%w: missing envelope", ErrInvalidAttestatio
 var ErrDecodingJSON = fmt.Errorf("%w: decoding json", ErrInvalidAttestation)
 var ErrDecodingB64 = fmt.Errorf("%w: decoding base64", ErrInvalidAttestation)
 
+const mediaTypeBase = "application/vnd.dev.sigstore.bundle"
+
 func ErrValidationError(err error) error {
 	return fmt.Errorf("%w: %w", ErrValidation, err)
 }
@@ -113,17 +115,32 @@ func (b *ProtobufBundle) validate() error {
 	return nil
 }
 
+// MediaTypeString returns a mediatype string for the specified bundle version.
+// The function returns an error if the resulting string does validate.
+func MediaTypeString(version string) (string, error) {
+	if version == "" {
+		return "", fmt.Errorf("unable to build media type string, no version defined")
+	}
+
+	mtString := fmt.Sprintf("%s+json;version=%s", mediaTypeBase, strings.TrimPrefix(version, "v"))
+	if _, err := getBundleVersion(mtString); err != nil {
+		return "", fmt.Errorf("unable to build mediatype: %w", err)
+	}
+
+	return mtString, nil
+}
+
 func getBundleVersion(mediaType string) (string, error) {
 	switch mediaType {
-	case "application/vnd.dev.sigstore.bundle+json;version=0.1":
+	case mediaTypeBase + "+json;version=0.1":
 		return "v0.1", nil
-	case "application/vnd.dev.sigstore.bundle+json;version=0.2":
+	case mediaTypeBase + "+json;version=0.2":
 		return "v0.2", nil
-	case "application/vnd.dev.sigstore.bundle+json;version=0.3":
+	case mediaTypeBase + "+json;version=0.3":
 		return "v0.3", nil
 	}
-	if strings.HasPrefix(mediaType, "application/vnd.dev.sigstore.bundle.v") && strings.HasSuffix(mediaType, "+json") {
-		version := strings.TrimPrefix(mediaType, "application/vnd.dev.sigstore.bundle.")
+	if strings.HasPrefix(mediaType, mediaTypeBase+".v") && strings.HasSuffix(mediaType, "+json") {
+		version := strings.TrimPrefix(mediaType, mediaTypeBase+".")
 		version = strings.TrimSuffix(version, "+json")
 		if semver.IsValid(version) {
 			return version, nil
diff --git a/pkg/bundle/bundle_test.go b/pkg/bundle/bundle_test.go
index f1d94196..ba0ca6a8 100644
--- a/pkg/bundle/bundle_test.go
+++ b/pkg/bundle/bundle_test.go
@@ -17,6 +17,8 @@ package bundle
 import (
 	"fmt"
 	"testing"
+
+	"github.com/stretchr/testify/require"
 )
 
 func Test_getBundleVersion(t *testing.T) {
@@ -94,3 +96,31 @@ func Test_getBundleVersion(t *testing.T) {
 		})
 	}
 }
+
+func TestMediaTypeString(t *testing.T) {
+	t.Parallel()
+	for _, tc := range []struct {
+		name     string
+		ver      string
+		expected string
+		mustErr  bool
+	}{
+		// Test for input, and valid an invalid versions. All validation
+		// is in getBundleVersion.
+		{"normal-semver", "v0.2", "application/vnd.dev.sigstore.bundle+json;version=0.2", false},
+		{"blank", "", "", true},
+		{"invalid", "garbage", "", true},
+	} {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			res, err := MediaTypeString(tc.ver)
+			if tc.mustErr {
+				require.Error(t, err)
+				return
+			}
+			require.NoError(t, err)
+			require.Equal(t, tc.expected, res)
+		})
+	}
+}