From 7e18e09ad38cc57b855a5234bd25c360a7cba0f2 Mon Sep 17 00:00:00 2001
From: Octogonapus <firey45@gmail.com>
Date: Thu, 15 Jun 2023 10:23:34 -0400
Subject: [PATCH] Fix mlflow name adjustments. The case of the name depends on
 the repository, as specified in the spec
 https://github.com/package-url/purl-spec/blob/master/PURL-TYPES.rst#mlflow
 Also fix the qualifier order in the mlflow test to match the example in the
 spec.

---
 packageurl.go                 | 28 ++++++++++++++++++++++++++--
 testdata/test-suite-data.json |  2 +-
 2 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/packageurl.go b/packageurl.go
index 7b851cb..9a2f5d3 100644
--- a/packageurl.go
+++ b/packageurl.go
@@ -91,6 +91,8 @@ var (
 	TypeSwift = "swift"
 	// TypeHuggingface is pkg:huggingface purl.
 	TypeHuggingface = "huggingface"
+	// TypeMLflow is pkg:mlflow purl.
+	TypeMLFlow = "mlflow"
 )
 
 // Qualifier represents a single key=value qualifier in the package url
@@ -285,7 +287,7 @@ func FromString(purl string) (PackageURL, error) {
 	remainder = nextSplit[1]
 
 	index = strings.LastIndex(remainder, "/")
-	name := typeAdjustName(purlType, remainder[index+1:])
+	name := typeAdjustName(purlType, remainder[index+1:], qualifiers)
 	version := ""
 
 	atIndex := strings.Index(name, "@")
@@ -352,12 +354,15 @@ func typeAdjustNamespace(purlType, ns string) string {
 
 // Make any purl type-specific adjustments to the parsed name.
 // See https://github.com/package-url/purl-spec#known-purl-types
-func typeAdjustName(purlType, name string) string {
+func typeAdjustName(purlType, name string, qualifiers Qualifiers) string {
+	quals := qualifiers.Map()
 	switch purlType {
 	case TypeBitbucket, TypeDebian, TypeGithub, TypeGolang, TypeNPM:
 		return strings.ToLower(name)
 	case TypePyPi:
 		return strings.ToLower(strings.ReplaceAll(name, "_", "-"))
+	case TypeMLFlow:
+		return adjustMlflowName(name, quals)
 	}
 	return name
 }
@@ -372,6 +377,25 @@ func typeAdjustVersion(purlType, version string) string {
 	return version
 }
 
+// https://github.com/package-url/purl-spec/blob/master/PURL-TYPES.rst#mlflow
+func adjustMlflowName(name string, qualifiers map[string]string) string {
+	if repo, ok := qualifiers["repository_url"]; ok {
+		if strings.Contains(repo, "azureml") {
+			// Azure ML is case-sensitive and must be kept as-is
+			return name
+		} else if strings.Contains(repo, "databricks") {
+			// Databricks is case-insensitive and must be lowercased
+			return strings.ToLower(name)
+		} else {
+			// Unknown repository type, keep as-is
+			return name
+		}
+	} else {
+		// No repository qualifier given, keep as-is
+		return name
+	}
+}
+
 // validQualifierKey validates a qualifierKey against our QualifierKeyPattern.
 func validQualifierKey(key string) bool {
 	return QualifierKeyPattern.MatchString(key)
diff --git a/testdata/test-suite-data.json b/testdata/test-suite-data.json
index a819fc8..8b0485f 100644
--- a/testdata/test-suite-data.json
+++ b/testdata/test-suite-data.json
@@ -530,7 +530,7 @@
   {
     "description": "MLflow model with unique identifiers",
     "purl": "pkg:mlflow/trafficsigns@10?model_uuid=36233173b22f4c89b451f1228d700d49&run_id=410a3121-2709-4f88-98dd-dba0ef056b0a&repository_url=https://adb-5245952564735461.0.azuredatabricks.net/api/2.0/mlflow",
-    "canonical_purl": "pkg:mlflow/trafficsigns@10?model_uuid=36233173b22f4c89b451f1228d700d49&repository_url=https://adb-5245952564735461.0.azuredatabricks.net/api/2.0/mlflow&run_id=410a3121-2709-4f88-98dd-dba0ef056b0a",
+    "canonical_purl": "pkg:mlflow/trafficsigns@10?model_uuid=36233173b22f4c89b451f1228d700d49&run_id=410a3121-2709-4f88-98dd-dba0ef056b0a&repository_url=https://adb-5245952564735461.0.azuredatabricks.net/api/2.0/mlflow",
     "type": "mlflow",
     "namespace": null,
     "name": "trafficsigns",