From b640a6011a1f2761f857d024d700d4363a0dc927 Mon Sep 17 00:00:00 2001
From: Daniel Pettersen <38071012+daniepett@users.noreply.github.com>
Date: Mon, 14 Mar 2022 23:11:34 +0000
Subject: [PATCH] fix: Allow legacy version of GrantIDs to be used with new
 grant functionality (#923)

---
 pkg/resources/grant_helpers.go               | 20 ++++++++--
 pkg/resources/grant_helpers_internal_test.go | 41 ++++++++++++++++++--
 2 files changed, 55 insertions(+), 6 deletions(-)

diff --git a/pkg/resources/grant_helpers.go b/pkg/resources/grant_helpers.go
index cf6701981f..da4478aefa 100644
--- a/pkg/resources/grant_helpers.go
+++ b/pkg/resources/grant_helpers.go
@@ -115,13 +115,27 @@ func grantIDFromString(stringID string) (*grantID, error) {
 	if len(lines) != 1 {
 		return nil, fmt.Errorf("1 line per grant")
 	}
-	if len(lines[0]) != 5 && len(lines[0]) != 6 {
-		return nil, fmt.Errorf("5 or 6 fields allowed")
+
+	// Len 4 is allowing for legacy IDs where role names are not included
+	if len(lines[0]) < 4 || len(lines[0]) > 6 {
+		return nil, fmt.Errorf("4 to 6 fields allowed in ID")
+	}
+
+	// Splitting string list if new ID structure, will cause issues if roles names passed are "true" or "false".
+	// Checking for true/false to eliminate scenarios where it would pick up the grant option.
+	// Roles will be empty list if legacy IDs are used, roles from grants are not
+	// used in Read functions, just for uniqueness in IDs of resources
+	roles := []string{}
+	if len(lines[0]) > 4 && lines[0][4] != "true" && lines[0][4] != "false" {
+		roles = strings.Split(lines[0][4], ",")
 	}
 
+	// Allowing legacy IDs to check grant option
 	grantOption := false
 	if len(lines[0]) == 6 && lines[0][5] == "true" {
 		grantOption = true
+	} else if len(lines[0]) == 5 && lines[0][4] == "true" {
+		grantOption = true
 	}
 
 	grantResult := &grantID{
@@ -129,7 +143,7 @@ func grantIDFromString(stringID string) (*grantID, error) {
 		SchemaName:   lines[0][1],
 		ObjectName:   lines[0][2],
 		Privilege:    lines[0][3],
-		Roles:        strings.Split(lines[0][4], ","),
+		Roles:        roles,
 		GrantOption:  grantOption,
 	}
 	return grantResult, nil
diff --git a/pkg/resources/grant_helpers_internal_test.go b/pkg/resources/grant_helpers_internal_test.go
index f7c1b91493..ec60aed7b2 100644
--- a/pkg/resources/grant_helpers_internal_test.go
+++ b/pkg/resources/grant_helpers_internal_test.go
@@ -44,17 +44,17 @@ func TestGrantIDFromString(t *testing.T) {
 	// Bad ID -- not enough fields
 	id = "database|name-privilege"
 	_, err = grantIDFromString(id)
-	r.Equal(fmt.Errorf("5 or 6 fields allowed"), err)
+	r.Equal(fmt.Errorf("4 to 6 fields allowed in ID"), err)
 
 	// Bad ID -- privilege in wrong area
 	id = "database||name-privilege"
 	_, err = grantIDFromString(id)
-	r.Equal(fmt.Errorf("5 or 6 fields allowed"), err)
+	r.Equal(fmt.Errorf("4 to 6 fields allowed in ID"), err)
 
 	// too many fields
 	id = "database_name|schema|view_name|privilege|false|2|too-many"
 	_, err = grantIDFromString(id)
-	r.Equal(fmt.Errorf("5 or 6 fields allowed"), err)
+	r.Equal(fmt.Errorf("4 to 6 fields allowed in ID"), err)
 
 	// 0 lines
 	id = ""
@@ -109,3 +109,38 @@ func TestGrantStruct(t *testing.T) {
 	r.Equal([]string{"test3", "test4"}, newGrant.Roles)
 	r.Equal(false, newGrant.GrantOption)
 }
+
+func TestGrantLegacyID(t *testing.T) {
+	// Testing that grants with legacy ID structure resolves to expected output
+	r := require.New(t)
+	gID := "database_name|schema|view_name|priv|true"
+	grant, err := grantIDFromString(gID)
+	r.NoError(err)
+	r.Equal("database_name", grant.ResourceName)
+	r.Equal("schema", grant.SchemaName)
+	r.Equal("view_name", grant.ObjectName)
+	r.Equal("priv", grant.Privilege)
+	r.Equal([]string{}, grant.Roles)
+	r.Equal(true, grant.GrantOption)
+
+	gID = "database_name|schema|view_name|priv|false"
+	grant, err = grantIDFromString(gID)
+	r.NoError(err)
+	r.Equal("database_name", grant.ResourceName)
+	r.Equal("schema", grant.SchemaName)
+	r.Equal("view_name", grant.ObjectName)
+	r.Equal("priv", grant.Privilege)
+	r.Equal([]string{}, grant.Roles)
+	r.Equal(false, grant.GrantOption)
+
+	gID = "database_name|schema|view_name|priv"
+	grant, err = grantIDFromString(gID)
+	r.NoError(err)
+	r.Equal("database_name", grant.ResourceName)
+	r.Equal("schema", grant.SchemaName)
+	r.Equal("view_name", grant.ObjectName)
+	r.Equal("priv", grant.Privilege)
+	r.Equal([]string{}, grant.Roles)
+	r.Equal(false, grant.GrantOption)
+
+}