From ef226cf82a4ef199f9d029433b23eb6291b9e232 Mon Sep 17 00:00:00 2001
From: Leonardo Chaia <leonardochaia@protonmail.com>
Date: Mon, 11 Nov 2024 14:27:49 -0300
Subject: [PATCH] chore: applies suggestions from code review

Signed-off-by: Leonardo Chaia <leonardochaia@protonmail.com>
---
 src/OrasProject.Oras/Registry/IMounter.cs     |  5 ++-
 .../Registry/Remote/BlobStore.cs              | 41 ++++++++++++-------
 .../Registry/Remote/Repository.cs             |  4 +-
 3 files changed, 32 insertions(+), 18 deletions(-)

diff --git a/src/OrasProject.Oras/Registry/IMounter.cs b/src/OrasProject.Oras/Registry/IMounter.cs
index 60f74eb..e9e36be 100644
--- a/src/OrasProject.Oras/Registry/IMounter.cs
+++ b/src/OrasProject.Oras/Registry/IMounter.cs
@@ -33,5 +33,8 @@ public interface IMounter
     /// <param name="getContent"></param>
     /// <param name="cancellationToken"></param>
     /// <returns></returns>
-    Task MountAsync(Descriptor descriptor, string fromRepository, Func<CancellationToken, Task<Stream>>? getContent, CancellationToken cancellationToken);
+    Task MountAsync(Descriptor descriptor,
+        string fromRepository,
+        Func<CancellationToken, Task<Stream>>? getContent = null,
+        CancellationToken cancellationToken = default);
 }
diff --git a/src/OrasProject.Oras/Registry/Remote/BlobStore.cs b/src/OrasProject.Oras/Registry/Remote/BlobStore.cs
index a8dc7e5..92d5324 100644
--- a/src/OrasProject.Oras/Registry/Remote/BlobStore.cs
+++ b/src/OrasProject.Oras/Registry/Remote/BlobStore.cs
@@ -207,6 +207,7 @@ public async Task MountAsync(Descriptor descriptor, string fromRepository,
             {
                 case HttpStatusCode.Created:
                     // 201, layer has been mounted
+                    response.VerifyContentDigest(descriptor.Digest);
                     return;
                 case HttpStatusCode.Accepted:
                 {
@@ -235,21 +236,29 @@ public async Task MountAsync(Descriptor descriptor, string fromRepository,
         //
         // [spec]: https://github.com/opencontainers/distribution-spec/blob/v1.1.0/spec.md#mounting-a-blob-from-another-repository
 
-        Stream contents;
-        if (getContent != null)
+        async Task<Stream> GetContentStream()
         {
-            contents = await getContent(cancellationToken).ConfigureAwait(false);
-        }
-        else
-        {
-            var referenceOptions = repository.Options with
+            Stream stream;
+            if (getContent != null)
+            {
+                stream = await getContent(cancellationToken).ConfigureAwait(false);
+            }
+            else
             {
-                Reference = Reference.Parse(fromRepository),
-            };
-            contents = await new Repository(referenceOptions).FetchAsync(descriptor, cancellationToken);
+                var referenceOptions = repository.Options with
+                {
+                    Reference = Reference.Parse(fromRepository),
+                };
+                stream = await new Repository(referenceOptions).FetchAsync(descriptor, cancellationToken).ConfigureAwait(false);
+            }
+
+            return stream;
         }
 
-        await InternalPushAsync(url, descriptor, contents, cancellationToken).ConfigureAwait(false);
+        await using (var contents = await GetContentStream())
+        {
+            await InternalPushAsync(url, descriptor, contents, cancellationToken).ConfigureAwait(false);
+        }
     }
 
     private async Task InternalPushAsync(Uri url, Descriptor descriptor, Stream content,
@@ -267,11 +276,13 @@ private async Task InternalPushAsync(Uri url, Descriptor descriptor, Stream cont
         // the descriptor media type is ignored as in the API doc.
         req.Content.Headers.ContentType = new MediaTypeHeaderValue(MediaTypeNames.Application.Octet);
 
-        using var response =
-            await Repository.Options.HttpClient.SendAsync(req, cancellationToken).ConfigureAwait(false);
-        if (response.StatusCode != HttpStatusCode.Created)
+        using (var response =
+               await Repository.Options.HttpClient.SendAsync(req, cancellationToken).ConfigureAwait(false))
         {
-            throw await response.ParseErrorResponseAsync(cancellationToken).ConfigureAwait(false);
+            if (response.StatusCode != HttpStatusCode.Created)
+            {
+                throw await response.ParseErrorResponseAsync(cancellationToken).ConfigureAwait(false);
+            }
         }
     }
 }
diff --git a/src/OrasProject.Oras/Registry/Remote/Repository.cs b/src/OrasProject.Oras/Registry/Remote/Repository.cs
index ca58a99..764f788 100644
--- a/src/OrasProject.Oras/Registry/Remote/Repository.cs
+++ b/src/OrasProject.Oras/Registry/Remote/Repository.cs
@@ -347,6 +347,6 @@ internal Reference ParseReferenceFromContentReference(string reference)
     /// <param name="getContent"></param>
     /// <param name="cancellationToken"></param>
     /// <returns></returns>
-    public Task MountAsync(Descriptor descriptor, string fromRepository, Func<CancellationToken, Task<Stream>>? getContent, CancellationToken cancellationToken) 
-        => ((IMounter)Blobs).MountAsync(descriptor, fromRepository, getContent, cancellationToken);
+    public async Task MountAsync(Descriptor descriptor, string fromRepository, Func<CancellationToken, Task<Stream>>? getContent, CancellationToken cancellationToken = default) 
+        => await ((IMounter)Blobs).MountAsync(descriptor, fromRepository, getContent, cancellationToken).ConfigureAwait(false);
 }