Repository: dragonflyoss/Dragonfly2
Pull request #1304 (2022-05-13 03:10:59)

feat: support more digest methods like sha256 in client
Signed-off-by: Jim Ma <majinjing3@gmail.com>

<!--- Provide a general summary of your changes in the Title above -->

## Description

<!--- Describe your changes in detail -->

## Related Issue

<!--- This project only accepts pull requests related to open issues -->
<!--- If suggesting a new feature or change, please discuss it in an issue first -->
<!--- If fixing a bug, there should be an issue describing it with steps to reproduce -->
<!--- Please link to the issue here: -->

## Motivation and Context

<!--- Why is this change required? What problem does it solve? -->

## Screenshots (if appropriate)

## Types of changes

<!--- What types of changes does your code introduce? Put an `x` in all the boxes that apply: -->

- [ ] Bug fix (non-breaking change which fixes an issue)
- [x] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing functionality to change)
- [ ] Documentation Update (if none of the other choices apply)

## Checklist

<!--- Go over all the following points, and put an `x` in all the boxes that apply. -->
<!--- If you're unsure about any of these, don't hesitate to ask. We're here to help! -->

- [ ] My change requires a change to the documentation.
- [ ] I have updated the documentation accordingly.
- [ ] I have read the **CONTRIBUTING** document.
- [ ] I have added tests to cover my changes.


codecov-commenter: # [Codecov](https://codecov.io/gh/dragonflyoss/Dragonfly2/pull/1304?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=dragonflyoss) Report
> Merging [#1304](https://codecov.io/gh/dragonflyoss/Dragonfly2/pull/1304?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=dragonflyoss) (44978a0) into [main](https://codecov.io/gh/dragonflyoss/Dragonfly2/commit/b5de837cbe67e9b5164b177d0613f12dd53c778c?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=dragonflyoss) (b5de837) will **increase** coverage by `0.02%`.
> The diff coverage is `38.23%`.

[![Impacted file tree graph](https://codecov.io/gh/dragonflyoss/Dragonfly2/pull/1304/graphs/tree.svg?width=650&height=150&src=pr&token=PUICJYSAT7&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=dragonflyoss)](https://codecov.io/gh/dragonflyoss/Dragonfly2/pull/1304?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=dragonflyoss)

```diff
@@            Coverage Diff             @@
##             main    #1304      +/-   ##
==========================================
+ Coverage   46.33%   46.35%   +0.02%     
==========================================
  Files         139      139              
  Lines       11689    11716      +27     
==========================================
+ Hits         5416     5431      +15     
- Misses       5684     5694      +10     
- Partials      589      591       +2     
```

| Flag | Coverage Δ | |
|---|---|---|
| Object-compatibility-e2etests | `?` | |
| e2etests | `?` | |
| unittests | `46.35% <38.23%> (+0.02%)` | :arrow_up: |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=dragonflyoss#carryforward-flags-in-the-pull-request-comment) to find out more.

| [Impacted Files](https://codecov.io/gh/dragonflyoss/Dragonfly2/pull/1304?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=dragonflyoss) | Coverage Δ | |
|---|---|---|
| [client/daemon/peer/piece\_downloader.go](https://codecov.io/gh/dragonflyoss/Dragonfly2/pull/1304/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=dragonflyoss#diff-Y2xpZW50L2RhZW1vbi9wZWVyL3BpZWNlX2Rvd25sb2FkZXIuZ28=) | `46.37% <20.00%> (-2.86%)` | :arrow_down: |
| [client/daemon/peer/piece\_manager.go](https://codecov.io/gh/dragonflyoss/Dragonfly2/pull/1304/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=dragonflyoss#diff-Y2xpZW50L2RhZW1vbi9wZWVyL3BpZWNlX21hbmFnZXIuZ28=) | `52.52% <33.33%> (-0.45%)` | :arrow_down: |
| [pkg/util/digestutils/digest\_reader.go](https://codecov.io/gh/dragonflyoss/Dragonfly2/pull/1304/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=dragonflyoss#diff-cGtnL3V0aWwvZGlnZXN0dXRpbHMvZGlnZXN0X3JlYWRlci5nbw==) | `56.09% <43.47%> (-20.10%)` | :arrow_down: |
| [scheduler/service/service.go](https://codecov.io/gh/dragonflyoss/Dragonfly2/pull/1304/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=dragonflyoss#diff-c2NoZWR1bGVyL3NlcnZpY2Uvc2VydmljZS5nbw==) | `82.77% <0.00%> (-0.64%)` | :arrow_down: |
| [client/daemon/peer/peertask\_conductor.go](https://codecov.io/gh/dragonflyoss/Dragonfly2/pull/1304/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=dragonflyoss#diff-Y2xpZW50L2RhZW1vbi9wZWVyL3BlZXJ0YXNrX2NvbmR1Y3Rvci5nbw==) | `58.87% <0.00%> (+1.03%)` | :arrow_up: |
| [client/daemon/peer/piece\_broker.go](https://codecov.io/gh/dragonflyoss/Dragonfly2/pull/1304/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=dragonflyoss#diff-Y2xpZW50L2RhZW1vbi9wZWVyL3BpZWNlX2Jyb2tlci5nbw==) | `92.30% <0.00%> (+2.56%)` | :arrow_up: |


Patch:
diff --git a/client/daemon/peer/piece_downloader.go b/client/daemon/peer/piece_downloader.go
index 035a9fba..1d0c3d83 100644
--- a/client/daemon/peer/piece_downloader.go
+++ b/client/daemon/peer/piece_downloader.go
@@ -179,7 +179,12 @@ func (p *pieceDownloader) DownloadPiece(ctx context.Context, req *DownloadPieceR
 	reader, closer := resp.Body.(io.Reader), resp.Body.(io.Closer)
 	if req.CalcDigest {
 		req.log.Debugf("calculate digest for piece %d, digest: %s", req.piece.PieceNum, req.piece.PieceMd5)
-		reader = digestutils.NewDigestReader(req.log, io.LimitReader(resp.Body, int64(req.piece.RangeSize)), req.piece.PieceMd5)
+		reader, err = digestutils.NewDigestReader(req.log, io.LimitReader(resp.Body, int64(req.piece.RangeSize)), req.piece.PieceMd5)
+		if err != nil {
+			_ = closer.Close()
+			req.log.Errorf("init digest reader error: %s", err.Error())
+			return nil, nil, err
+		}
 	}
 	return reader, closer, nil
 }
diff --git a/client/daemon/peer/piece_manager.go b/client/daemon/peer/piece_manager.go
index 51ebfcce..7a958aa4 100644
--- a/client/daemon/peer/piece_manager.go
+++ b/client/daemon/peer/piece_manager.go
@@ -204,7 +204,7 @@ func (pm *pieceManager) processPieceFromSource(pt Task,
 	}
 	if pm.calculateDigest {
 		pt.Log().Debugf("calculate digest")
-		reader = digestutils.NewDigestReader(pt.Log(), reader)
+		reader, _ = digestutils.NewDigestReader(pt.Log(), reader)
 	}
 	var n int64
 	result.Size, err = pt.GetStorage().WritePiece(
@@ -290,7 +290,11 @@ func (pm *pieceManager) DownloadSource(ctx context.Context, pt Task, request *sc
 
 	// calc total
 	if pm.calculateDigest {
-		reader = digestutils.NewDigestReader(pt.Log(), response.Body, request.UrlMeta.Digest)
+		reader, err = digestutils.NewDigestReader(pt.Log(), response.Body, request.UrlMeta.Digest)
+		if err != nil {
+			log.Errorf("init digest reader error: %s", err.Error())
+			return err
+		}
 	}
 	// we must calculate piece size
 	pieceSize := pm.computePieceSize(contentLength)
@@ -470,7 +474,7 @@ func (pm *pieceManager) processPieceFromFile(ctx context.Context, ptm storage.Pe
 
 	if pm.calculateDigest {
 		log.Debugf("calculate digest in processPieceFromFile")
-		reader = digestutils.NewDigestReader(log, r)
+		reader, _ = digestutils.NewDigestReader(log, r)
 	}
 	n, err := tsd.WritePiece(ctx,
 		&storage.WritePieceRequest{
diff --git a/pkg/util/digestutils/digest_reader.go b/pkg/util/digestutils/digest_reader.go
index 00645e90..f5682df7 100644
--- a/pkg/util/digestutils/digest_reader.go
+++ b/pkg/util/digestutils/digest_reader.go
@@ -18,9 +18,14 @@ package digestutils
 
 import (
 	"crypto/md5"
+	"crypto/sha1"
+	"crypto/sha256"
+	"crypto/sha512"
 	"encoding/hex"
+	"fmt"
 	"hash"
 	"io"
+	"strings"
 
 	"github.com/pkg/errors"
 
@@ -46,18 +51,38 @@ type DigestReader interface {
 
 // TODO add AF_ALG digest https://github.com/golang/sys/commit/e24f485414aeafb646f6fca458b0bf869c0880a1
 
-func NewDigestReader(log *logger.SugaredLoggerOnWith, reader io.Reader, digest ...string) io.Reader {
-	var d string
+func NewDigestReader(log *logger.SugaredLoggerOnWith, reader io.Reader, digest ...string) (io.Reader, error) {
+	var (
+		d          string
+		hashMethod hash.Hash
+	)
 	if len(digest) > 0 {
 		d = digest[0]
 	}
+	ds := strings.Split(d, ":")
+	if len(ds) == 2 {
+		d = ds[1]
+		switch ds[0] {
+		case "sha1":
+			hashMethod = sha1.New()
+		case "sha256":
+			hashMethod = sha256.New()
+		case "sha512":
+			hashMethod = sha512.New()
+		case "md5":
+			hashMethod = md5.New()
+		default:
+			return nil, fmt.Errorf("unsupport digest method: %s", ds[0])
+		}
+	} else {
+		hashMethod = md5.New()
+	}
 	return &digestReader{
 		SugaredLoggerOnWith: log,
 		digest:              d,
-		// TODO support more digest method like sha1, sha256
-		hash: md5.New(),
-		r:    reader,
-	}
+		hash:                hashMethod,
+		r:                   reader,
+	}, nil
 }
 
 func (dr *digestReader) Read(p []byte) (int, error) {


Test patch:
diff --git a/pkg/util/digestutils/digest_reader_test.go b/pkg/util/digestutils/digest_reader_test.go
index 687bb574..12d9b83a 100644
--- a/pkg/util/digestutils/digest_reader_test.go
+++ b/pkg/util/digestutils/digest_reader_test.go
@@ -42,7 +42,8 @@ func TestNewDigestReader(t *testing.T) {
 	digest := hex.EncodeToString(hash.Sum(nil)[:16])
 
 	buf := bytes.NewBuffer(testBytes)
-	reader := NewDigestReader(logger.With("test", "test"), buf, digest)
+	reader, err := NewDigestReader(logger.With("test", "test"), buf, digest)
+	assert.Nil(err)
 	data, err := io.ReadAll(reader)
 
 	assert.Nil(err)