From ffb48cf2c03a6729a2782230a12e15184639df8e Mon Sep 17 00:00:00 2001
From: RebeccaMahany <rebeccamahany@gmail.com>
Date: Wed, 11 Dec 2024 15:58:56 -0500
Subject: [PATCH] Return identifier type from SplitKey too

---
 ee/agent/startupsettings/writer.go |  2 +-
 ee/agent/storage/keys.go           |  8 +++---
 ee/agent/storage/keys_test.go      | 39 +++++++++++++++++-------------
 pkg/osquery/table/table.go         |  2 +-
 4 files changed, 28 insertions(+), 23 deletions(-)

diff --git a/ee/agent/startupsettings/writer.go b/ee/agent/startupsettings/writer.go
index 2e9cc7f71..d489228ad 100644
--- a/ee/agent/startupsettings/writer.go
+++ b/ee/agent/startupsettings/writer.go
@@ -148,7 +148,7 @@ func (s *startupSettingsWriter) extractAutoTableConstructionConfig(registrationI
 func (s *startupSettingsWriter) extractKATCConstructionConfig(registrationId string) (string, error) {
 	kolideCfg := make(map[string]string)
 	if err := s.knapsack.KatcConfigStore().ForEach(func(k []byte, v []byte) error {
-		key, identifier := storage.SplitKey(k)
+		key, _, identifier := storage.SplitKey(k)
 		if string(identifier) == registrationId {
 			kolideCfg[string(key)] = string(v)
 		}
diff --git a/ee/agent/storage/keys.go b/ee/agent/storage/keys.go
index a780a28db..84ea26b42 100644
--- a/ee/agent/storage/keys.go
+++ b/ee/agent/storage/keys.go
@@ -36,17 +36,17 @@ func KeyByIdentifier(key []byte, identifierType []byte, identifier []byte) []byt
 	return newKey
 }
 
-func SplitKey(key []byte) ([]byte, []byte) {
+func SplitKey(key []byte) ([]byte, []byte, []byte) {
 	if !bytes.Contains(key, []byte{keyDelimiter}) {
-		return key, defaultIdentifier
+		return key, nil, defaultIdentifier
 	}
 
 	// Key takes the form `<key>:<identifierType>:<identifier>` -- split
 	// on the keyDelimiter.
 	parts := bytes.SplitN(key, []byte{keyDelimiter}, 3)
 	if len(parts) != 3 {
-		return key, defaultIdentifier
+		return key, nil, defaultIdentifier
 	}
 
-	return parts[0], parts[2]
+	return parts[0], parts[1], parts[2]
 }
diff --git a/ee/agent/storage/keys_test.go b/ee/agent/storage/keys_test.go
index 0209bb37a..4e4b84fb7 100644
--- a/ee/agent/storage/keys_test.go
+++ b/ee/agent/storage/keys_test.go
@@ -51,36 +51,41 @@ func TestSplitKey(t *testing.T) {
 	t.Parallel()
 
 	for _, tt := range []struct {
-		testCaseName       string
-		key                []byte
-		expectedKey        []byte
-		expectedIdentifier []byte
+		testCaseName           string
+		key                    []byte
+		expectedKey            []byte
+		expectedIdentifierType []byte
+		expectedIdentifier     []byte
 	}{
 		{
-			testCaseName:       "default node key",
-			key:                []byte("nodeKey"),
-			expectedKey:        []byte("nodeKey"),
-			expectedIdentifier: []byte("default"),
+			testCaseName:           "default node key",
+			key:                    []byte("nodeKey"),
+			expectedKey:            []byte("nodeKey"),
+			expectedIdentifierType: nil,
+			expectedIdentifier:     []byte("default"),
 		},
 		{
-			testCaseName:       "uuid by registration",
-			key:                []byte("uuid:registration:some-test-registration-id"),
-			expectedKey:        []byte("uuid"),
-			expectedIdentifier: []byte("some-test-registration-id"),
+			testCaseName:           "uuid by registration",
+			key:                    []byte("uuid:registration:some-test-registration-id"),
+			expectedKey:            []byte("uuid"),
+			expectedIdentifierType: IdentifierTypeRegistration,
+			expectedIdentifier:     []byte("some-test-registration-id"),
 		},
 		{
-			testCaseName:       "katc table by registration",
-			key:                []byte("katc_some_test_table:registration:another-test-registration-id"),
-			expectedKey:        []byte("katc_some_test_table"),
-			expectedIdentifier: []byte("another-test-registration-id"),
+			testCaseName:           "katc table by registration",
+			key:                    []byte("katc_some_test_table:registration:another-test-registration-id"),
+			expectedKey:            []byte("katc_some_test_table"),
+			expectedIdentifierType: IdentifierTypeRegistration,
+			expectedIdentifier:     []byte("another-test-registration-id"),
 		},
 	} {
 		tt := tt
 		t.Run(tt.testCaseName, func(t *testing.T) {
 			t.Parallel()
 
-			splitKey, identifier := SplitKey(tt.key)
+			splitKey, identifierType, identifier := SplitKey(tt.key)
 			require.Equal(t, tt.expectedKey, splitKey)
+			require.Equal(t, tt.expectedIdentifierType, identifierType)
 			require.Equal(t, tt.expectedIdentifier, identifier)
 		})
 	}
diff --git a/pkg/osquery/table/table.go b/pkg/osquery/table/table.go
index 2bd17b609..e02ede881 100644
--- a/pkg/osquery/table/table.go
+++ b/pkg/osquery/table/table.go
@@ -108,7 +108,7 @@ func katcFromDb(k types.Knapsack, registrationId string) (map[string]string, err
 	}
 	katcCfg := make(map[string]string)
 	if err := k.KatcConfigStore().ForEach(func(k []byte, v []byte) error {
-		key, identifier := storage.SplitKey(k)
+		key, _, identifier := storage.SplitKey(k)
 		if string(identifier) == registrationId {
 			katcCfg[string(key)] = string(v)
 		}