From bf198c3918d980ec842ba3151313efb28582fb79 Mon Sep 17 00:00:00 2001
From: Michael Yang <mxyng@pm.me>
Date: Thu, 20 Jul 2023 11:44:05 -0700
Subject: [PATCH] verify blob digest

---
 server/images.go | 27 +++++++++++++++++++++++++++
 1 file changed, 27 insertions(+)

diff --git a/server/images.go b/server/images.go
index 4258b201..579ff713 100644
--- a/server/images.go
+++ b/server/images.go
@@ -623,6 +623,13 @@ func PullModel(name, username, password string, fn func(api.ProgressResponse)) e
 		completed += layer.Size
 	}
 
+	fn(api.ProgressResponse{Status: "verifying sha256 digest"})
+	for _, layer := range layers {
+		if err := verifyBlob(layer.Digest); err != nil {
+			return err
+		}
+	}
+
 	fn(api.ProgressResponse{Status: "writing manifest"})
 
 	manifestJSON, err := json.Marshal(manifest)
@@ -917,3 +924,23 @@ func makeRequest(method, url string, headers map[string]string, body io.Reader,
 
 	return resp, nil
 }
+
+func verifyBlob(digest string) error {
+	fp, err := GetBlobsPath(digest)
+	if err != nil {
+		return err
+	}
+
+	f, err := os.Open(fp)
+	if err != nil {
+		return err
+	}
+	defer f.Close()
+
+	fileDigest, _ := GetSHA256Digest(f)
+	if digest != fileDigest {
+		return fmt.Errorf("digest mismatch: want %s, got %s", digest, fileDigest)
+	}
+
+	return nil
+}