From 5a5ca8e7ff042842013cf7ad1462a5d6c705d39b Mon Sep 17 00:00:00 2001
From: Bruce MacDonald <brucewmacdonald@gmail.com>
Date: Mon, 24 Jul 2023 14:53:01 -0400
Subject: [PATCH] remove file on digest mismatch

---
 server/images.go | 17 +++++++++++++++--
 1 file changed, 15 insertions(+), 2 deletions(-)

diff --git a/server/images.go b/server/images.go
index babfca54..9fe16307 100644
--- a/server/images.go
+++ b/server/images.go
@@ -510,7 +510,7 @@ func CopyModel(src, dest string) error {
 		return err
 	}
 
-	err = ioutil.WriteFile(destPath, input, 0644)
+	err = ioutil.WriteFile(destPath, input, 0o644)
 	if err != nil {
 		fmt.Println("Error reading file:", err)
 		return err
@@ -700,6 +700,17 @@ func PullModel(name string, regOpts *RegistryOptions, fn func(api.ProgressRespon
 	fn(api.ProgressResponse{Status: "verifying sha256 digest"})
 	for _, layer := range layers {
 		if err := verifyBlob(layer.Digest); err != nil {
+			if errors.Is(err, errDigestMismatch) {
+				// something went wrong, delete the blob
+				fp, err := GetBlobsPath(layer.Digest)
+				if err != nil {
+					return err
+				}
+				if err := os.Remove(fp); err != nil {
+					// log this, but return the original error
+					log.Printf("couldn't remove file with digest mismatch '%s': %v", fp, err)
+				}
+			}
 			return err
 		}
 	}
@@ -1059,6 +1070,8 @@ func makeRequest(method, url string, headers map[string]string, body io.Reader,
 	return resp, nil
 }
 
+var errDigestMismatch = fmt.Errorf("digest mismatch, file must be downloaded again")
+
 func verifyBlob(digest string) error {
 	fp, err := GetBlobsPath(digest)
 	if err != nil {
@@ -1073,7 +1086,7 @@ func verifyBlob(digest string) error {
 
 	fileDigest, _ := GetSHA256Digest(f)
 	if digest != fileDigest {
-		return fmt.Errorf("digest mismatch: want %s, got %s", digest, fileDigest)
+		return fmt.Errorf("%w: want %s, got %s", errDigestMismatch, digest, fileDigest)
 	}
 
 	return nil