From 4b1d01612d89bb6ab4e613f7fb0ed8ccbef11d99 Mon Sep 17 00:00:00 2001
From: Billy Zha <jinzha1@microsoft.com>
Date: Mon, 12 Dec 2022 15:25:44 +0800
Subject: [PATCH] fix: use proxy when copying node with successors (#384)

Resolves #383

Signed-off-by: Billy Zha <jinzha1@microsoft.com>
---
 copy.go      | 46 ++++++++++++++++++++++------------------------
 copy_test.go | 15 ++++++++++++++-
 2 files changed, 36 insertions(+), 25 deletions(-)

diff --git a/copy.go b/copy.go
index 2a35f692..92222e54 100644
--- a/copy.go
+++ b/copy.go
@@ -219,38 +219,36 @@ func copyGraph(ctx context.Context, src content.ReadOnlyStorage, dst content.Sto
 		}
 		successors = removeForeignLayers(successors)
 
-		// handle leaf nodes
-		if len(successors) == 0 {
-			exists, err = proxy.Cache.Exists(ctx, desc)
-			if err != nil {
+		if len(successors) != 0 {
+			// for non-leaf nodes, process successors and wait for them to complete
+			region.End()
+			if err := syncutil.Go(ctx, limiter, fn, successors...); err != nil {
 				return err
 			}
-			if exists {
-				return copyNode(ctx, proxy.Cache, dst, desc, opts)
+			for _, node := range successors {
+				done, committed := tracker.TryCommit(node)
+				if committed {
+					return fmt.Errorf("%s: %s: successor not committed", desc.Digest, node.Digest)
+				}
+				select {
+				case <-done:
+				case <-ctx.Done():
+					return ctx.Err()
+				}
+			}
+			if err := region.Start(); err != nil {
+				return err
 			}
-			return copyNode(ctx, src, dst, desc, opts)
 		}
 
-		// for non-leaf nodes, process successors and wait for them to complete
-		region.End()
-		if err := syncutil.Go(ctx, limiter, fn, successors...); err != nil {
+		exists, err = proxy.Cache.Exists(ctx, desc)
+		if err != nil {
 			return err
 		}
-		for _, node := range successors {
-			done, committed := tracker.TryCommit(node)
-			if committed {
-				return fmt.Errorf("%s: %s: successor not committed", desc.Digest, node.Digest)
-			}
-			select {
-			case <-done:
-			case <-ctx.Done():
-				return ctx.Err()
-			}
-		}
-		if err := region.Start(); err != nil {
-			return err
+		if exists {
+			return copyNode(ctx, proxy.Cache, dst, desc, opts)
 		}
-		return copyNode(ctx, proxy.Cache, dst, desc, opts)
+		return copyNode(ctx, src, dst, desc, opts)
 	}
 
 	return syncutil.Go(ctx, limiter, fn, root)
diff --git a/copy_test.go b/copy_test.go
index 73df3317..119357b8 100644
--- a/copy_test.go
+++ b/copy_test.go
@@ -303,7 +303,7 @@ func TestCopyGraph_FullCopy(t *testing.T) {
 		}
 	}
 
-	// test copy
+	// test copy graph
 	srcTracker := &storageTracker{Storage: src}
 	dstTracker := &storageTracker{Storage: dst}
 	root := descs[len(descs)-1]
@@ -1353,6 +1353,19 @@ func TestCopyGraph_WithOptions(t *testing.T) {
 		}
 	}
 
+	// test successor descriptors not obtained from src
+	root = descs[3]
+	opts = oras.DefaultCopyGraphOptions
+	opts.FindSuccessors = func(ctx context.Context, fetcher content.Fetcher, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
+		if content.Equal(desc, root) {
+			return descs[1:3], nil
+		}
+		return content.Successors(ctx, fetcher, desc)
+	}
+	if err := oras.CopyGraph(ctx, src, cas.NewMemory(), root, opts); err != nil {
+		t.Fatalf("CopyGraph() error = %v, wantErr %v", err, false)
+	}
+
 	// test partial copy
 	var preCopyCount int64
 	var postCopyCount int64