소스 검색

remove file on digest mismatch

Bruce MacDonald 1 년 전
부모
커밋
5a5ca8e7ff
1개의 변경된 파일15개의 추가작업 그리고 2개의 파일을 삭제
  1. 15 2
      server/images.go

+ 15 - 2
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