feat(policies): read all pages in the chainloop.discover builtin

pkg/policies/engine/builtins/discover.go

@@ -1,5 +1,5 @@
 //
-// Copyright 2025 The Chainloop Authors.
+// Copyright 2025-2026 The Chainloop Authors.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -22,25 +22,87 @@ import (
 	"google.golang.org/grpc"
 )
 
+const (
+	// discoverPageSize is the number of references requested per page, the maximum that the API accepts
+	discoverPageSize = 100
+	// discoverMaxPages caps the number of requests of one discover call
+	discoverMaxPages = 10
+)
+
+// DiscoverResult is the referrer with the references of all the pages read.
+// It serializes to the same shape as a ReferrerItem, plus the truncated field.
+type DiscoverResult struct {
+	*v1.ReferrerItem
+	// Truncated is true when the read stopped before the last page of references: either the
+	// page cap was reached, or the control plane repeated its next cursor (it ignores the cursor)
+	Truncated bool `json:"truncated"`
+}
+
 // DiscoverService wraps the gRPC discover functionality to be shared across engines
 type DiscoverService struct {
-	conn *grpc.ClientConn
+	client v1.ReferrerServiceClient
 }
 
 // NewDiscoverService creates a new discover service
 func NewDiscoverService(conn *grpc.ClientConn) *DiscoverService {
-	return &DiscoverService{conn: conn}
+	if conn == nil {
+		return &DiscoverService{}
+	}
+
+	return &DiscoverService{client: v1.NewReferrerServiceClient(conn)}
 }
 
-// Discover calls the DiscoverPrivate gRPC endpoint to get artifact graph data
-func (s *DiscoverService) Discover(ctx context.Context, digest, kind string) (*v1.ReferrerServiceDiscoverPrivateResponse, error) {
-	if s.conn == nil {
+// Discover calls the DiscoverPrivate gRPC endpoint to get artifact graph data.
+// It follows the pagination cursor and merges the references of all pages,
+// up to discoverMaxPages requests.
+//
+// A control plane without pagination support ignores the pagination request and
+// returns all the references in one response with no next cursor, so the read
+// ends after the first page. A next cursor that does not advance also ends the
+// read, so the same page is never merged twice, and marks the result as truncated.
+func (s *DiscoverService) Discover(ctx context.Context, digest, kind string) (*DiscoverResult, error) {
+	if s.client == nil {
 		return nil, nil
 	}
 
-	client := v1.NewReferrerServiceClient(s.conn)
-	return client.DiscoverPrivate(ctx, &v1.ReferrerServiceDiscoverPrivateRequest{
-		Digest: digest,
-		Kind:   kind,
-	})
+	readPage := func(cursor string) (*v1.ReferrerServiceDiscoverPrivateResponse, error) {
+		return s.client.DiscoverPrivate(ctx, &v1.ReferrerServiceDiscoverPrivateRequest{
+			Digest:     digest,
+			Kind:       kind,
+			Pagination: &v1.CursorPaginationRequest{Cursor: cursor, Limit: discoverPageSize},
+		})
+	}
+
+	resp, err := readPage("")
+	if err != nil {
+		return nil, err
+	}
+
+	result := &DiscoverResult{ReferrerItem: resp.GetResult()}
+	for pages := 1; result.ReferrerItem != nil; pages++ {
+		cursor := resp.GetPagination().GetNextCursor()
+		if cursor == "" {
+			break
+		}
+
+		if pages == discoverMaxPages {
+			result.Truncated = true
+			break
+		}
+
+		if resp, err = readPage(cursor); err != nil {
+			return nil, err
+		}
+
+		// a server that ignores the cursor returns the same page and the same next cursor again,
+		// so the rest of the references cannot be read
+		if resp.GetPagination().GetNextCursor() == cursor {
+			result.Truncated = true
+			break
+		}
+
+		result.References = append(result.References, resp.GetResult().GetReferences()...)
+	}
+
+	return result, nil
 }

pkg/policies/engine/builtins/discover_test.go

@@ -0,0 +1,217 @@
+//
+// Copyright 2026 The Chainloop Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package builtins
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"strconv"
+	"strings"
+	"testing"
+
+	v1 "github.com/chainloop-dev/chainloop/app/controlplane/api/controlplane/v1"
+	"github.com/open-policy-agent/opa/v1/ast"
+	"github.com/open-policy-agent/opa/v1/topdown"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc"
+)
+
+// testKind is the kind of the referrer that the tests discover
+const testKind = "CONTAINER_IMAGE"
+
+// fakeReferrerClient serves total references of one digest in pages of the requested limit.
+// The cursor is the index of the first reference of the page.
+type fakeReferrerClient struct {
+	v1.ReferrerServiceClient
+	total  int
+	failAt int // fail the request with this page number (1-based), 0 means never
+	// noPagination acts as a control plane without pagination: it ignores the
+	// pagination request and returns all the references with no pagination response
+	noPagination bool
+	// ignoreCursor acts as a control plane that pages but ignores the cursor
+	ignoreCursor bool
+	requests     []*v1.ReferrerServiceDiscoverPrivateRequest
+}
+
+func (f *fakeReferrerClient) DiscoverPrivate(_ context.Context, req *v1.ReferrerServiceDiscoverPrivateRequest, _ ...grpc.CallOption) (*v1.ReferrerServiceDiscoverPrivateResponse, error) {
+	f.requests = append(f.requests, req)
+	if f.failAt == len(f.requests) {
+		return nil, errors.New("boom")
+	}
+
+	if f.noPagination {
+		return &v1.ReferrerServiceDiscoverPrivateResponse{
+			Result: &v1.ReferrerItem{Digest: req.GetDigest(), Kind: testKind, References: f.refs(0, f.total)},
+		}, nil
+	}
+
+	start := 0
+	if c := req.GetPagination().GetCursor(); c != "" && !f.ignoreCursor {
+		var err error
+		if start, err = strconv.Atoi(c); err != nil {
+			return nil, err
+		}
+	}
+	end := min(start+int(req.GetPagination().GetLimit()), f.total)
+
+	var next string
+	if end < f.total {
+		next = strconv.Itoa(end)
+	}
+
+	return &v1.ReferrerServiceDiscoverPrivateResponse{
+		Result:     &v1.ReferrerItem{Digest: req.GetDigest(), Kind: testKind, References: f.refs(start, end)},
+		Pagination: &v1.CursorPaginationResponse{NextCursor: next},
+	}, nil
+}
+
+func (f *fakeReferrerClient) refs(start, end int) []*v1.ReferrerItem {
+	refs := make([]*v1.ReferrerItem, 0, end-start)
+	for i := start; i < end; i++ {
+		refs = append(refs, &v1.ReferrerItem{Digest: fmt.Sprintf("sha256:%d", i), Kind: "ATTESTATION"})
+	}
+	return refs
+}
+
+func TestDiscoverReadsAllPages(t *testing.T) {
+	const maxRefs = discoverMaxPages * discoverPageSize
+	testCases := []struct {
+		name          string
+		total         int
+		failAt        int
+		noPagination  bool
+		ignoreCursor  bool
+		wantRequests  int
+		wantRefs      int
+		wantTruncated bool
+		wantErr       bool
+	}{
+		{name: "no references", total: 0, wantRequests: 1},
+		{name: "single page", total: 20, wantRequests: 1, wantRefs: 20},
+		{name: "exactly one full page", total: discoverPageSize, wantRequests: 1, wantRefs: discoverPageSize},
+		{name: "multiple pages", total: 2*discoverPageSize + 50, wantRequests: 3, wantRefs: 2*discoverPageSize + 50},
+		{name: "last allowed page ends the read", total: maxRefs, wantRequests: discoverMaxPages, wantRefs: maxRefs},
+		{name: "page cap stops the read", total: maxRefs + 1, wantRequests: discoverMaxPages, wantRefs: maxRefs, wantTruncated: true},
+		{name: "error on a later page", total: 2*discoverPageSize + 50, failAt: 2, wantRequests: 2, wantErr: true},
+		// a control plane without pagination returns everything in one response
+		{name: "control plane without pagination", total: maxRefs + 500, noPagination: true, wantRequests: 1, wantRefs: maxRefs + 500},
+		// a control plane that ignores the cursor repeats the first page: keep it once, stop and report truncated
+		{name: "control plane ignores the cursor", total: 2*discoverPageSize + 50, ignoreCursor: true, wantRequests: 2, wantRefs: discoverPageSize, wantTruncated: true},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			client := &fakeReferrerClient{total: tc.total, failAt: tc.failAt, noPagination: tc.noPagination, ignoreCursor: tc.ignoreCursor}
+			svc := &DiscoverService{client: client}
+
+			got, err := svc.Discover(context.Background(), "sha256:deadbeef", testKind)
+			assert.Len(t, client.requests, tc.wantRequests)
+			if tc.wantErr {
+				require.Error(t, err)
+				return
+			}
+			require.NoError(t, err)
+
+			// every request asks for the largest page and keeps the digest and kind
+			for i, req := range client.requests {
+				assert.Equal(t, "sha256:deadbeef", req.GetDigest())
+				assert.Equal(t, testKind, req.GetKind())
+				assert.Equal(t, int32(discoverPageSize), req.GetPagination().GetLimit())
+				if i == 0 {
+					assert.Empty(t, req.GetPagination().GetCursor())
+				}
+			}
+
+			assert.Equal(t, "sha256:deadbeef", got.GetDigest())
+			assert.Equal(t, tc.wantTruncated, got.Truncated)
+			require.Len(t, got.GetReferences(), tc.wantRefs)
+			// the references keep the order of the pages
+			for i, ref := range got.GetReferences() {
+				assert.Equal(t, fmt.Sprintf("sha256:%d", i), ref.GetDigest())
+			}
+		})
+	}
+}
+
+func TestDiscoverRegoImpl(t *testing.T) {
+	testCases := []struct {
+		name string
+		svc  *DiscoverService
+		want string
+	}{
+		{
+			name: "merged references and truncated field",
+			svc:  &DiscoverService{client: &fakeReferrerClient{total: discoverPageSize + 1}},
+			want: fmt.Sprintf(`{"digest": "sha256:deadbeef", "kind": "CONTAINER_IMAGE", "truncated": false, "references": [%s]}`, refsJSON(discoverPageSize+1)),
+		},
+		{
+			name: "no connection",
+			svc:  NewDiscoverService(nil),
+			want: "null",
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			var got *ast.Term
+			err := getDiscoverRegoImpl(tc.svc)(
+				topdown.BuiltinContext{Context: context.Background()},
+				[]*ast.Term{ast.StringTerm("sha256:deadbeef"), ast.StringTerm("")},
+				func(term *ast.Term) error { got = term; return nil },
+			)
+			require.NoError(t, err)
+			require.NotNil(t, got)
+			assert.JSONEq(t, tc.want, got.String())
+		})
+	}
+}
+
+// refsJSON returns the JSON of the first n references that fakeReferrerClient serves
+func refsJSON(n int) string {
+	refs := make([]string, n)
+	for i := range refs {
+		refs[i] = fmt.Sprintf(`{"digest": "sha256:%d", "kind": "ATTESTATION"}`, i)
+	}
+	return strings.Join(refs, ",")
+}
+
+func TestDiscoverNoConnection(t *testing.T) {
+	got, err := NewDiscoverService(nil).Discover(context.Background(), "sha256:deadbeef", "")
+	require.NoError(t, err)
+	assert.Nil(t, got)
+}
+
+func TestDiscoverResultJSONShape(t *testing.T) {
+	res := &DiscoverResult{
+		ReferrerItem: &v1.ReferrerItem{
+			Digest:     "sha256:deadbeef",
+			Kind:       testKind,
+			References: []*v1.ReferrerItem{{Digest: "sha256:1", Kind: "ATTESTATION"}},
+		},
+	}
+
+	raw, err := json.Marshal(res)
+	require.NoError(t, err)
+	assert.JSONEq(t, `{
+		"digest": "sha256:deadbeef",
+		"kind": "CONTAINER_IMAGE",
+		"references": [{"digest": "sha256:1", "kind": "ATTESTATION"}],
+		"truncated": false
+	}`, string(raw))
+}

pkg/policies/engine/builtins/rego.go

@@ -1,5 +1,5 @@
 //
-// Copyright 2025 The Chainloop Authors.
+// Copyright 2025-2026 The Chainloop Authors.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -31,6 +31,10 @@ const discoverBuiltinName = "chainloop.discover"
 //
 //	chainloop.discover(digest, kind)
 //
+// The result holds the references of all the pages of the API, up to a cap.
+// When the read stops before the last page, the result has "truncated": true. This happens
+// when the cap is reached, or when the control plane repeats its next cursor.
+//
 // For instance, to get the references for an CONTAINER_IMAGE material, and fail if any of them is an attestation with policy violations:
 // ```
 //
@@ -55,15 +59,13 @@ func RegisterDiscoverBuiltin(conn *grpc.ClientConn) error {
 				types.Named("digest", types.S).Description("digest of the artifact to discover"),
 				types.Named("kind", types.S).Description("optional filter by kind to disambiguate"),
 			),
-			types.Named("response", types.A).Description("response object as in the `chainloop discover` CLI output"),
+			types.Named("response", types.A).Description("response object as in the `chainloop discover` CLI output, with the references of all pages and a truncated field"),
 		),
 		Nondeterministic: true,
-	}, getDiscoverRegoImpl(conn))
+	}, getDiscoverRegoImpl(NewDiscoverService(conn)))
 }
 
-func getDiscoverRegoImpl(conn *grpc.ClientConn) topdown.BuiltinFunc {
-	discoverSvc := NewDiscoverService(conn)
-
+func getDiscoverRegoImpl(discoverSvc *DiscoverService) topdown.BuiltinFunc {
 	return func(bctx topdown.BuiltinContext, operands []*ast.Term, iter func(*ast.Term) error) error {
 		if len(operands) < 1 {
 			return errors.New("need at least one operand")
@@ -92,8 +94,13 @@ func getDiscoverRegoImpl(conn *grpc.ClientConn) topdown.BuiltinFunc {
 			return fmt.Errorf("failed to call discover endpoint: %w", err)
 		}
 
+		// no connection to the control plane
+		if resp == nil {
+			return iter(ast.NullTerm())
+		}
+
 		// call the iterator with the output value
-		return iter(ast.NewTerm(ast.MustInterfaceToValue(resp.Result)))
+		return iter(ast.NewTerm(ast.MustInterfaceToValue(resp)))
 	}
 }
 

pkg/policies/engine/builtins/wasm.go

@@ -1,5 +1,5 @@
 //
-// Copyright 2024-2025 The Chainloop Authors.
+// Copyright 2024-2026 The Chainloop Authors.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -57,7 +57,7 @@ func CreateDiscoverHostFunctions(conn *grpc.ClientConn) []extism.HostFunction {
 		}
 
 		// Serialize response to JSON
-		jsonData, err := json.Marshal(resp.Result)
+		jsonData, err := json.Marshal(resp)
 		if err != nil {
 			// Return 0 to signal error
 			stack[0] = 0