From b5c20aaec2a867793e1d73880cc28473b2dc3ed1 Mon Sep 17 00:00:00 2001
From: rosstimothy <39066650+rosstimothy@users.noreply.github.com>
Date: Mon, 7 Oct 2024 17:55:07 +0000
Subject: [PATCH] Fix backends with prefixes (#47238)

#46701 incorrectly changed how the prefixes for the etcd and dynamo
backends were applied and removed. In addition, the backend lock
key was not being constructed correctly and caused acquisition of
locks to fail on a cluster that had been upgraded.

Tests have been added to ensure that keys are now being constructed
correctly in all cases. Additionally, backend/helpers.go was moved
to backed/lock.go to better described what functionality the file
contains.

Backends with prefixes (etcd, dynamo) were incorrectly constructing
the key for locks. In addition the prefixes were not being trimmed
correctly which broke CAS operations.
---
 lib/backend/dynamo/dynamodbbk.go              |  2 +-
 lib/backend/dynamo/dynamodbbk_test.go         | 19 +++++++++++++
 lib/backend/etcdbk/etcd.go                    |  6 ++---
 lib/backend/etcdbk/etcd_test.go               | 27 +++++++++++++++++++
 lib/backend/{helpers.go => lock.go}           |  6 +++--
 lib/backend/{helpers_test.go => lock_test.go} | 15 +++++++++++
 6 files changed, 69 insertions(+), 6 deletions(-)
 rename lib/backend/{helpers.go => lock.go} (97%)
 rename lib/backend/{helpers_test.go => lock_test.go} (89%)

diff --git a/lib/backend/dynamo/dynamodbbk.go b/lib/backend/dynamo/dynamodbbk.go
index cd174ccea4b58..0a8c2bf897fbc 100644
--- a/lib/backend/dynamo/dynamodbbk.go
+++ b/lib/backend/dynamo/dynamodbbk.go
@@ -919,7 +919,7 @@ const (
 // prependPrefix adds leading 'teleport/' to the key for backwards compatibility
 // with previous implementation of DynamoDB backend
 func prependPrefix(key backend.Key) string {
-	return keyPrefix + string(key)
+	return keyPrefix + key.String()
 }
 
 // trimPrefix removes leading 'teleport' from the key
diff --git a/lib/backend/dynamo/dynamodbbk_test.go b/lib/backend/dynamo/dynamodbbk_test.go
index 0b5c05897bfcb..b467ec6599511 100644
--- a/lib/backend/dynamo/dynamodbbk_test.go
+++ b/lib/backend/dynamo/dynamodbbk_test.go
@@ -31,6 +31,7 @@ import (
 	"github.com/gravitational/trace"
 	"github.com/jonboulle/clockwork"
 	log "github.com/sirupsen/logrus"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	"github.com/gravitational/teleport/lib/backend"
@@ -222,3 +223,21 @@ func TestCreateTable(t *testing.T) {
 		})
 	}
 }
+
+func TestKeyPrefix(t *testing.T) {
+	t.Run("leading separator in key", func(t *testing.T) {
+		prefixed := prependPrefix(backend.NewKey("test", "llama"))
+		assert.Equal(t, "teleport/test/llama", prefixed)
+
+		key := trimPrefix(prefixed)
+		assert.Equal(t, "/test/llama", key.String())
+	})
+
+	t.Run("no leading separator in key", func(t *testing.T) {
+		prefixed := prependPrefix(backend.Key(".locks/test/llama"))
+		assert.Equal(t, "teleport.locks/test/llama", prefixed)
+
+		key := trimPrefix(prefixed)
+		assert.Equal(t, ".locks/test/llama", key.String())
+	})
+}
diff --git a/lib/backend/etcdbk/etcd.go b/lib/backend/etcdbk/etcd.go
index a268be31c78bc..48f99d0871077 100644
--- a/lib/backend/etcdbk/etcd.go
+++ b/lib/backend/etcdbk/etcd.go
@@ -1110,10 +1110,10 @@ func fromType(eventType mvccpb.Event_EventType) types.OpType {
 	}
 }
 
-func (b *EtcdBackend) trimPrefix(in backend.Key) backend.Key {
-	return in.TrimPrefix(backend.Key(b.cfg.Key))
+func (b *EtcdBackend) trimPrefix(in []byte) backend.Key {
+	return backend.Key(in).TrimPrefix(backend.Key(b.cfg.Key))
 }
 
 func (b *EtcdBackend) prependPrefix(in backend.Key) string {
-	return b.cfg.Key + string(in)
+	return b.cfg.Key + in.String()
 }
diff --git a/lib/backend/etcdbk/etcd_test.go b/lib/backend/etcdbk/etcd_test.go
index 2367e313ab8a7..a148fe9ac800e 100644
--- a/lib/backend/etcdbk/etcd_test.go
+++ b/lib/backend/etcdbk/etcd_test.go
@@ -29,6 +29,7 @@ import (
 
 	"github.com/gravitational/trace"
 	"github.com/jonboulle/clockwork"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	"github.com/gravitational/teleport/lib/backend"
@@ -253,3 +254,29 @@ func etcdTestEndpoint() string {
 	}
 	return "https://127.0.0.1:2379"
 }
+
+func TestKeyPrefix(t *testing.T) {
+	prefixes := []string{"teleport", "/teleport", "/teleport/"}
+
+	for _, prefix := range prefixes {
+		t.Run("prefix="+prefix, func(t *testing.T) {
+			bk := EtcdBackend{cfg: &Config{Key: prefix}}
+
+			t.Run("leading separator in key", func(t *testing.T) {
+				prefixed := bk.prependPrefix(backend.NewKey("test", "llama"))
+				assert.Equal(t, prefix+"/test/llama", prefixed)
+
+				key := bk.trimPrefix([]byte(prefixed))
+				assert.Equal(t, "/test/llama", key.String())
+			})
+
+			t.Run("no leading separator in key", func(t *testing.T) {
+				prefixed := bk.prependPrefix(backend.Key(".locks/test/llama"))
+				assert.Equal(t, prefix+".locks/test/llama", prefixed)
+
+				key := bk.trimPrefix([]byte(prefixed))
+				assert.Equal(t, ".locks/test/llama", key.String())
+			})
+		})
+	}
+}
diff --git a/lib/backend/helpers.go b/lib/backend/lock.go
similarity index 97%
rename from lib/backend/helpers.go
rename to lib/backend/lock.go
index 85cb4bcc68236..54082968de38d 100644
--- a/lib/backend/helpers.go
+++ b/lib/backend/lock.go
@@ -21,11 +21,13 @@ package backend
 import (
 	"bytes"
 	"context"
+	"log/slog"
 	"time"
 
 	"github.com/google/uuid"
 	"github.com/gravitational/trace"
-	log "github.com/sirupsen/logrus"
+
+	logutils "github.com/gravitational/teleport/lib/utils/log"
 )
 
 const (
@@ -205,7 +207,7 @@ func RunWhileLocked(ctx context.Context, cfg RunWhileLockedConfig, fn func(conte
 			case <-cfg.Backend.Clock().After(refreshAfter):
 				if err := lock.resetTTL(ctx, cfg.Backend); err != nil {
 					cancelFunction()
-					log.Errorf("%v", err)
+					slog.ErrorContext(ctx, "failed to reset lock ttl", "error", err, "lock", logutils.StringerAttr(lock.key))
 					return
 				}
 			case <-stopRefresh:
diff --git a/lib/backend/helpers_test.go b/lib/backend/lock_test.go
similarity index 89%
rename from lib/backend/helpers_test.go
rename to lib/backend/lock_test.go
index b652ee157e83a..466a9cd4cdf60 100644
--- a/lib/backend/helpers_test.go
+++ b/lib/backend/lock_test.go
@@ -23,9 +23,24 @@ import (
 	"time"
 
 	"github.com/google/go-cmp/cmp"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
+func TestLockKey(t *testing.T) {
+	t.Run("empty parts", func(t *testing.T) {
+		key := LockKey()
+		assert.Equal(t, ".locks", key.String())
+		assert.Equal(t, [][]byte{[]byte(".locks")}, key.Components())
+	})
+
+	t.Run("with parts", func(t *testing.T) {
+		key := LockKey("test", "llama")
+		assert.Equal(t, ".locks/test/llama", key.String())
+		assert.Equal(t, [][]byte{[]byte(".locks"), []byte("test"), []byte("llama")}, key.Components())
+	})
+}
+
 func TestLockConfiguration_CheckAndSetDefaults(t *testing.T) {
 	type mockBackend struct {
 		Backend