fixing issues with search depth

Signed-off-by: chaosinthecrd <[email protected]>
ChaosInTheCRD · Sep 4, 2024 · 3c3f4af · 3c3f4af
1 parent 9dbe8c5
commit 3c3f4af
Show file tree

Hide file tree

Showing 8 changed files with 13 additions and 14 deletions.
diff --git a/policy/policy.go b/policy/policy.go
@@ -244,7 +244,7 @@ func (p Policy) Verify(ctx context.Context, opts ...VerifyOption) (bool, map[str
 				resultsByStep[stepName] = StepResult{Step: stepName}
 			}
 
-			collections, err := vo.verifiedSource.Search(ctx, stepName, vo.subjectDigests, attestationsByStep[stepName])
+			collections, err := vo.verifiedSource.Search(ctx, depth, stepName, vo.subjectDigests, attestationsByStep[stepName])
 			if err != nil {
 				return false, nil, err
 			}

diff --git a/policy/policy_test.go b/policy/policy_test.go
@@ -417,7 +417,7 @@ func newDummyVerifiedSourcer(verifiedCollections []source.CollectionVerification
 	return &dummyVerifiedSourcer{verifiedCollections}
 }
 
-func (s *dummyVerifiedSourcer) Search(ctx context.Context, collectionName string, subjectDigests, attestations []string) ([]source.CollectionVerificationResult, error) {
+func (s *dummyVerifiedSourcer) Search(ctx context.Context, depth int, collectionName string, subjectDigests, attestations []string) ([]source.CollectionVerificationResult, error) {
 	return s.verifiedCollections, nil
 }
 

diff --git a/source/archivista.go b/source/archivista.go
@@ -32,7 +32,7 @@ func NewArchvistSource(client *archivista.Client) *ArchivistaSource {
 	}
 }
 
-func (s *ArchivistaSource) Search(ctx context.Context, collectionName string, subjectDigests, attestations []string) ([]CollectionEnvelope, error) {
+func (s *ArchivistaSource) Search(ctx context.Context, depth int, collectionName string, subjectDigests, attestations []string) ([]CollectionEnvelope, error) {
 	gitoids, err := s.client.SearchGitoids(ctx, archivista.SearchGitoidVariables{
 		CollectionName: collectionName,
 		SubjectDigests: subjectDigests,

diff --git a/source/memory.go b/source/memory.go
@@ -106,12 +106,11 @@ func (s *MemorySource) LoadEnvelope(reference string, env dsse.Envelope) error {
 	return nil
 }
 
-func (s *MemorySource) Search(ctx context.Context, collectionName string, subjectDigests, attestations []string) ([]CollectionEnvelope, error) {
-	if s.searched {
+func (s *MemorySource) Search(ctx context.Context, depth int, collectionName string, subjectDigests, attestations []string) ([]CollectionEnvelope, error) {
+	if depth > 0 {
+		fmt.Println("skipping memory source search")
 		log.Debug("skipping memory source search: already performed")
 		return []CollectionEnvelope{}, nil
-	} else {
-		s.searched = true
 	}
 
 	matches := make([]CollectionEnvelope, 0)

diff --git a/source/memory_test.go b/source/memory_test.go
@@ -304,7 +304,7 @@ func TestSearch(t *testing.T) {
 			}
 
 			// Run the search query on the MemorySource
-			got, err := s.Search(tt.searchQuery.ctx, tt.searchQuery.collectionName, tt.searchQuery.subDigest, tt.searchQuery.attestations)
+			got, err := s.Search(tt.searchQuery.ctx, 0, tt.searchQuery.collectionName, tt.searchQuery.subDigest, tt.searchQuery.attestations)
 			if (err != nil) != tt.wantErr {
 				t.Fatalf("MemorySource.Search() error = %v, wantErr %v", err, tt.wantErr)
 			}

diff --git a/source/multi.go b/source/multi.go
@@ -28,7 +28,7 @@ func NewMultiSource(sources ...Sourcer) *MultiSource {
 }
 
 // Search concurrently queries all sources and returns the combined results.
-func (s *MultiSource) Search(ctx context.Context, collectionName string, subjectDigests, attestations []string) ([]CollectionEnvelope, error) {
+func (s *MultiSource) Search(ctx context.Context, depth int, collectionName string, subjectDigests, attestations []string) ([]CollectionEnvelope, error) {
 	results := []CollectionEnvelope{}
 	errors := []error{}
 
@@ -61,7 +61,7 @@ func (s *MultiSource) Search(ctx context.Context, collectionName string, subject
 		// Goroutine for querying a source and collecting the results or error
 		go func(src Sourcer) {
 			defer wg.Done()
-			res, err := src.Search(ctx, collectionName, subjectDigests, attestations)
+			res, err := src.Search(ctx, depth, collectionName, subjectDigests, attestations)
 			if err != nil {
 				errs <- err
 			} else {

diff --git a/source/source.go b/source/source.go
@@ -31,7 +31,7 @@ type CollectionEnvelope struct {
 }
 
 type Sourcer interface {
-	Search(ctx context.Context, collectionName string, subjectDigests, attestations []string) ([]CollectionEnvelope, error)
+	Search(ctx context.Context, depth int, collectionName string, subjectDigests, attestations []string) ([]CollectionEnvelope, error)
 }
 
 func envelopeToCollectionEnvelope(reference string, env dsse.Envelope) (CollectionEnvelope, error) {

diff --git a/source/verified.go b/source/verified.go
@@ -31,7 +31,7 @@ type CollectionVerificationResult struct {
 }
 
 type VerifiedSourcer interface {
-	Search(ctx context.Context, collectionName string, subjectDigests, attestations []string) ([]CollectionVerificationResult, error)
+	Search(ctx context.Context, depth int, collectionName string, subjectDigests, attestations []string) ([]CollectionVerificationResult, error)
 }
 
 type VerifiedSource struct {
@@ -43,8 +43,8 @@ func NewVerifiedSource(source Sourcer, verifyOpts ...dsse.VerificationOption) *V
 	return &VerifiedSource{source, verifyOpts}
 }
 
-func (s *VerifiedSource) Search(ctx context.Context, collectionName string, subjectDigests, attestations []string) ([]CollectionVerificationResult, error) {
-	unverified, err := s.source.Search(ctx, collectionName, subjectDigests, attestations)
+func (s *VerifiedSource) Search(ctx context.Context, depth int, collectionName string, subjectDigests, attestations []string) ([]CollectionVerificationResult, error) {
+	unverified, err := s.source.Search(ctx, depth, collectionName, subjectDigests, attestations)
 	if err != nil {
 		return nil, err
 	}