fix: fix nextKey values when using multiple prefixes (#43486)

This PR makes `generic.Service` correctly implementing `List*` functions when multiple key prefixes are defined

Signed-off-by: Tiago Silva <[email protected]>

lib/services/local/generic/generic.go

@@ -78,7 +78,6 @@ func (c *ServiceConfig[T]) CheckAndSetDefaults() error {
 	if c.UnmarshalFunc == nil {
 		return trace.BadParameter("unmarshal func is missing")
 	}
-
 	return nil
 }
 
@@ -167,17 +166,17 @@ func (s *Service[T]) GetResources(ctx context.Context) ([]T, error) {
 
 // ListResources returns a paginated list of resources.
 func (s *Service[T]) ListResources(ctx context.Context, pageSize int, pageToken string) ([]T, string, error) {
-	resources, next, err := s.ListResourcesReturnNextResource(ctx, pageSize, pageToken)
-	var nextKey string
-	if next != nil {
-		nextKey = backend.GetPaginationKey(*next)
-	}
+	resources, _, nextKey, err := s.listResourcesReturnNextResourceWithKey(ctx, pageSize, pageToken)
 	return resources, nextKey, trace.Wrap(err)
 }
 
 // ListResourcesReturnNextResource returns a paginated list of resources. The next resource is returned, which allows consumers to construct
 // the next pagination key as appropriate.
 func (s *Service[T]) ListResourcesReturnNextResource(ctx context.Context, pageSize int, pageToken string) ([]T, *T, error) {
+	resources, next, _, err := s.listResourcesReturnNextResourceWithKey(ctx, pageSize, pageToken)
+	return resources, next, trace.Wrap(err)
+}
+func (s *Service[T]) listResourcesReturnNextResourceWithKey(ctx context.Context, pageSize int, pageToken string) ([]T, *T, string, error) {
 	rangeStart := backend.Key(s.backendPrefix, pageToken)
 	rangeEnd := backend.RangeEnd(backend.ExactKey(s.backendPrefix))
 
@@ -191,26 +190,32 @@ func (s *Service[T]) ListResourcesReturnNextResource(ctx context.Context, pageSi
 	// no filter provided get the range directly
 	result, err := s.backend.GetRange(ctx, rangeStart, rangeEnd, limit)
 	if err != nil {
-		return nil, nil, trace.Wrap(err)
+		return nil, nil, "", trace.Wrap(err)
 	}
 
 	out := make([]T, 0, len(result.Items))
+	var lastKey []byte
 	for _, item := range result.Items {
 		resource, err := s.unmarshalFunc(item.Value, services.WithRevision(item.Revision), services.WithResourceID(item.ID))
 		if err != nil {
-			return nil, nil, trace.Wrap(err)
+			return nil, nil, "", trace.Wrap(err)
 		}
 		out = append(out, resource)
+		lastKey = item.Key
 	}
 
-	var next *T
+	var (
+		next    *T
+		nextKey string
+	)
 	if len(out) > pageSize {
 		next = &out[pageSize]
 		// Truncate the last item that was used to determine next row existence.
 		out = out[:pageSize]
+		nextKey = trimLastKey(string(lastKey), s.backendPrefix)
 	}
 
-	return out, next, nil
+	return out, next, nextKey, nil
 }
 
 // ListResourcesWithFilter returns a paginated list of resources that match the given filter.
@@ -226,6 +231,7 @@ func (s *Service[T]) ListResourcesWithFilter(ctx context.Context, pageSize int,
 	limit := pageSize + 1
 
 	var resources []T
+	var lastKey []byte
 	if err := backend.IterateRange(
 		ctx,
 		s.backend,
@@ -239,6 +245,7 @@ func (s *Service[T]) ListResourcesWithFilter(ctx context.Context, pageSize int,
 					return false, trace.Wrap(err)
 				}
 				if matcher(resource) {
+					lastKey = item.Key
 					resources = append(resources, resource)
 				}
 				if len(resources) == pageSize {
@@ -252,7 +259,7 @@ func (s *Service[T]) ListResourcesWithFilter(ctx context.Context, pageSize int,
 
 	var nextKey string
 	if len(resources) > pageSize {
-		nextKey = backend.GetPaginationKey(resources[pageSize])
+		nextKey = trimLastKey(string(lastKey), s.backendPrefix)
 		// Truncate the last item that was used to determine next row existence.
 		resources = resources[:pageSize]
 	}
@@ -460,3 +467,13 @@ func (s *Service[T]) RunWhileLocked(ctx context.Context, lockName string, ttl ti
 			return fn(ctx, s.backend)
 		}))
 }
+
+func trimLastKey(lastKey, prefix string) string {
+	if !strings.HasSuffix(prefix, string(backend.Separator)) {
+		prefix += string(backend.Separator)
+	}
+	if !strings.HasPrefix(prefix, string(backend.Separator)) {
+		prefix = string(backend.Separator) + prefix
+	}
+	return strings.TrimPrefix(lastKey, prefix)
+}

lib/services/local/generic/generic_test.go

@@ -333,6 +333,65 @@ func TestGenericListResourcesReturnNextResource(t *testing.T) {
 	require.Nil(t, next)
 }
 
+func TestGenericListResourcesWithMultiplePrefixes(t *testing.T) {
+	ctx := context.Background()
+
+	memBackend, err := memory.New(memory.Config{
+		Context: ctx,
+		Clock:   clockwork.NewFakeClock(),
+	})
+	require.NoError(t, err)
+
+	service, err := NewService(&ServiceConfig[*testResource]{
+		Backend:       memBackend,
+		ResourceKind:  "generic resource",
+		PageLimit:     200,
+		BackendPrefix: "generic_prefix",
+		UnmarshalFunc: unmarshalResource,
+		MarshalFunc:   marshalResource,
+	})
+	require.NoError(t, err)
+
+	// Create a couple test resources.
+	r1 := newTestResource("r1")
+	r2 := newTestResource("r2")
+
+	_, err = service.WithPrefix("a-unique-prefix").UpsertResource(ctx, r1)
+	require.NoError(t, err)
+	_, err = service.WithPrefix("another-unique-prefix").UpsertResource(ctx, r2)
+	require.NoError(t, err)
+
+	page, next, err := service.ListResources(ctx, 1, "")
+	require.NoError(t, err)
+	require.Empty(t, cmp.Diff([]*testResource{r1}, page,
+		cmpopts.IgnoreFields(types.Metadata{}, "Revision"),
+	))
+	require.Equal(t, next, "another-unique-prefix"+string(backend.Separator)+r2.GetName())
+
+	page, next, err = service.ListResources(ctx, 1, next)
+	require.NoError(t, err)
+	require.Empty(t, cmp.Diff([]*testResource{r2}, page,
+		cmpopts.IgnoreFields(types.Metadata{}, "Revision"),
+	))
+	require.Empty(t, next)
+
+	_, err = service.WithPrefix("a-unique-prefix").UpsertResource(ctx, r2)
+	require.NoError(t, err)
+	page, next, err = service.WithPrefix("a-unique-prefix").ListResources(ctx, 1, "")
+	require.NoError(t, err)
+	require.Empty(t, cmp.Diff([]*testResource{r1}, page,
+		cmpopts.IgnoreFields(types.Metadata{}, "Revision"),
+	))
+	require.Equal(t, next, r2.GetName())
+
+	page, next, err = service.WithPrefix("a-unique-prefix").ListResources(ctx, 1, next)
+	require.NoError(t, err)
+	require.Empty(t, cmp.Diff([]*testResource{r2}, page,
+		cmpopts.IgnoreFields(types.Metadata{}, "Revision"),
+	))
+	require.Empty(t, next)
+}
+
 func TestGenericListResourcesWithFilter(t *testing.T) {
 	ctx := context.Background()
 

lib/services/local/generic/generic_wrapper.go

@@ -33,8 +33,8 @@ func NewServiceWrapper[T types.ResourceMetadata](
 	resourceKind string,
 	backendPrefix string,
 	marshalFunc MarshalFunc[T],
-	unmarshalFunc UnmarshalFunc[T]) (*ServiceWrapper[T], error) {
-
+	unmarshalFunc UnmarshalFunc[T],
+) (*ServiceWrapper[T], error) {
 	cfg := &ServiceConfig[resourceMetadataAdapter[T]]{
 		Backend:       backend,
 		ResourceKind:  resourceKind,