From 9090fa81f817cd0b10cfa6af26e80385c9a7bb69 Mon Sep 17 00:00:00 2001
From: Eric Bao <baojn@zju.edu.cn>
Date: Wed, 19 May 2021 14:32:14 +0800
Subject: [PATCH] verify layer digest when push the layer (#118)

---
 registry/layer.go          | 25 ++++++++++++++++++-------
 utils/compress/compress.go | 23 ++++++++++-------------
 2 files changed, 28 insertions(+), 20 deletions(-)

diff --git a/registry/layer.go b/registry/layer.go
index 1ef44e0ff40..bde1e7d5ef5 100644
--- a/registry/layer.go
+++ b/registry/layer.go
@@ -2,6 +2,7 @@ package registry
 
 import (
 	"context"
+	"encoding/hex"
 	"errors"
 	"io"
 	"net/http"
@@ -33,25 +34,35 @@ func (registry *Registry) DownloadLayer(ctx context.Context, repository string,
 }
 
 // UploadLayer uploads a specific layer by digest for a repository.
-func (registry *Registry) UploadLayer(ctx context.Context, repository string, digest reference.Reference, content io.Reader) error {
+func (registry *Registry) UploadLayer(ctx context.Context, repository string, digs reference.Reference, content io.Reader) error {
 	uploadURL, token, err := registry.initiateUpload(ctx, repository)
 	if err != nil {
 		return err
 	}
+
 	q := uploadURL.Query()
-	q.Set("digest", digest.String())
+	q.Set("digest", digs.String())
 	uploadURL.RawQuery = q.Encode()
-
-	registry.Logf("registry.layer.upload url=%s repository=%s digest=%s", uploadURL, repository, digest)
-
-	upload, err := http.NewRequest("PUT", uploadURL.String(), content)
+	dig, ok := digs.(digest.Digest)
+	if !ok {
+		return errors.New("layer digest is necessary when push a layer")
+	}
+	// used to verify compressed layer hash
+	hash := dig.Algorithm().Hash()
+	uploadReader := io.TeeReader(content, hash)
+	registry.Logf("registry.layer.upload url=%s repository=%s digest=%s", uploadURL, repository, digs)
+	upload, err := http.NewRequest("PUT", uploadURL.String(), uploadReader)
 	if err != nil {
 		return err
 	}
+
 	upload.Header.Set("Content-Type", "application/octet-stream")
 	upload.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token))
-
 	_, err = registry.Client.Do(upload.WithContext(ctx))
+	currentDigestHex := hex.EncodeToString(hash.Sum(nil))
+	if currentDigestHex != dig.Hex() {
+		return fmt.Errorf("failed to push layer, err: layer digest changed, original: %s, current: %s", dig.Hex(), currentDigestHex)
+	}
 	return err
 }
 
diff --git a/utils/compress/compress.go b/utils/compress/compress.go
index a098c376802..6c7da965632 100644
--- a/utils/compress/compress.go
+++ b/utils/compress/compress.go
@@ -186,7 +186,7 @@ func Decompress(src io.Reader, dst string) error {
 
 		switch header.Typeflag {
 		case tar.TypeDir:
-			if _, err := os.Stat(target); err != nil {
+			if _, err = os.Stat(target); err != nil {
 				if err = os.MkdirAll(target, os.FileMode(header.Mode)); err != nil {
 					return err
 				}
@@ -207,25 +207,22 @@ func Decompress(src io.Reader, dst string) error {
 		case tar.TypeReg:
 			err = func() error {
 				// regularly won't mkdir, unless add newFolder on compressing
-				err := utils.MkDirIfNotExists(filepath.Dir(target))
-				if err != nil {
-					return err
+				inErr := utils.MkDirIfNotExists(filepath.Dir(target))
+				if inErr != nil {
+					return inErr
 				}
 
-				fileToWrite, err := os.OpenFile(target, os.O_CREATE|os.O_TRUNC|os.O_RDWR, os.FileMode(header.Mode))
-				if err != nil {
-					return err
+				fileToWrite, inErr := os.OpenFile(target, os.O_CREATE|os.O_TRUNC|os.O_RDWR, os.FileMode(header.Mode))
+				if inErr != nil {
+					return inErr
 				}
 
 				defer fileToWrite.Close()
-				if _, err := io.Copy(fileToWrite, tr); err != nil {
-					return err
+				if _, inErr = io.Copy(fileToWrite, tr); inErr != nil {
+					return inErr
 				}
 				// for not changing
-				if err = os.Chtimes(target, header.AccessTime, header.ModTime); err != nil {
-					return err
-				}
-				return nil
+				return os.Chtimes(target, header.AccessTime, header.ModTime)
 			}()
 
 			if err != nil {