feat: validate the sha256 after pulling the blob

Signed-off-by: chlins <[email protected]>

Author: chlins

pkg/backend/fetch_test.go

@@ -19,13 +19,15 @@ package backend
 import (
 	"context"
 	"encoding/json"
+	"fmt"
 	"net/http"
 	"net/http/httptest"
 	"os"
 	"strings"
 	"testing"
 
 	modelspec "github.com/CloudNativeAI/model-spec/specs-go/v1"
+	godigest "github.com/opencontainers/go-digest"
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -39,6 +41,15 @@ func TestFetch(t *testing.T) {
 	require.NoError(t, err)
 	defer os.RemoveAll(tempDir)
 
+	// Setup mock file
+	const (
+		file1Content = "file1 content..."
+		file2Content = "file2 content..."
+	)
+
+	file1Digest := godigest.FromString(file1Content)
+	file2Digest := godigest.FromString(file2Content)
+
 	// Setup mock server
 	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		switch r.URL.Path {
@@ -51,16 +62,16 @@ func TestFetch(t *testing.T) {
 				Layers: []ocispec.Descriptor{
 					{
 						MediaType: "application/octet-stream.raw",
-						Digest:    "sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",
-						Size:      0,
+						Digest:    file1Digest,
+						Size:      int64(len(file1Content)),
 						Annotations: map[string]string{
 							modelspec.AnnotationFilepath: "file1.txt",
 						},
 					},
 					{
 						MediaType: "application/octet-stream.raw",
-						Digest:    "sha256:a3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",
-						Size:      0,
+						Digest:    file2Digest,
+						Size:      int64(len(file2Content)),
 						Annotations: map[string]string{
 							modelspec.AnnotationFilepath: "file2.txt",
 						},
@@ -69,10 +80,13 @@ func TestFetch(t *testing.T) {
 			}
 			w.Header().Set("Content-Type", "application/json")
 			require.NoError(t, json.NewEncoder(w).Encode(manifest))
-		case "/v2/test/model/blobs/sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",
-			"/v2/test/model/blobs/sha256:a3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855":
-			// Return empty content for blobs
-			w.WriteHeader(http.StatusOK)
+		case fmt.Sprintf("/v2/test/model/blobs/%s", file1Digest):
+			_, err := w.Write([]byte(file1Content))
+			require.NoError(t, err)
+
+		case fmt.Sprintf("/v2/test/model/blobs/%s", file2Digest):
+			_, err := w.Write([]byte(file2Content))
+			require.NoError(t, err)
 		default:
 			t.Logf("Unexpected request to %s", r.URL.Path)
 			w.WriteHeader(http.StatusNotFound)

pkg/backend/pull.go

@@ -22,14 +22,15 @@ import (
 	"fmt"
 	"io"
 
+	retry "github.com/avast/retry-go/v4"
+	sha256 "github.com/minio/sha256-simd"
+	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
+	"golang.org/x/sync/errgroup"
+
 	internalpb "github.com/CloudNativeAI/modctl/internal/pb"
 	"github.com/CloudNativeAI/modctl/pkg/backend/remote"
 	"github.com/CloudNativeAI/modctl/pkg/config"
 	"github.com/CloudNativeAI/modctl/pkg/storage"
-
-	retry "github.com/avast/retry-go/v4"
-	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
-	"golang.org/x/sync/errgroup"
 )
 
 // Pull pulls an artifact from a registry.
@@ -136,6 +137,8 @@ func pullIfNotExist(ctx context.Context, pb *internalpb.ProgressBar, prompt stri
 	defer content.Close()
 
 	reader := pb.Add(prompt, desc.Digest.String(), desc.Size, content)
+	hash := sha256.New()
+	reader = io.TeeReader(reader, hash)
 
 	// push the content to the destination, and wrap the content reader for progress bar,
 	// manifest should use dst.Manifests().Push, others should use dst.Blobs().Push.
@@ -185,6 +188,13 @@ func pullIfNotExist(ctx context.Context, pb *internalpb.ProgressBar, prompt stri
 		}
 	}
 
+	// validate the digest of the blob.
+	if err := validateDigest(desc.Digest.String(), hash.Sum(nil)); err != nil {
+		err = fmt.Errorf("failed to validate the digest of the blob %s, err: %w", desc.Digest.String(), err)
+		pb.Abort(desc.Digest.String(), err)
+		return err
+	}
+
 	return nil
 }
 
@@ -199,11 +209,38 @@ func pullAndExtractFromRemote(ctx context.Context, pb *internalpb.ProgressBar, p
 	defer content.Close()
 
 	reader := pb.Add(prompt, desc.Digest.String(), desc.Size, content)
+	hash := sha256.New()
+	reader = io.TeeReader(reader, hash)
+
 	if err := extractLayer(desc, outputDir, reader); err != nil {
 		err = fmt.Errorf("failed to extract the blob %s to output directory: %w", desc.Digest.String(), err)
 		pb.Abort(desc.Digest.String(), err)
 		return err
 	}
 
+	// validate the digest of the blob.
+	if err := validateDigest(desc.Digest.String(), hash.Sum(nil)); err != nil {
+		err = fmt.Errorf("failed to validate the digest of the blob %s, err: %w", desc.Digest.String(), err)
+		pb.Abort(desc.Digest.String(), err)
+		return err
+	}
+
+	return nil
+}
+
+// validateDigest validates the hash digest whether matches the expected digest.
+func validateDigest(digest string, hash []byte) error {
+	if digest == "" {
+		return fmt.Errorf("digest is empty")
+	}
+
+	if len(hash) != sha256.Size {
+		return fmt.Errorf("invalid hash length")
+	}
+
+	if digest != fmt.Sprintf("sha256:%x", hash) {
+		return fmt.Errorf("actual digest %s does not match the expected digest %s", fmt.Sprintf("sha256:%x", hash), digest)
+	}
+
 	return nil
 }