Modify default hash function (#15)

* Modify default hash functions.

Signed-off-by: txaty <[email protected]>

* Add unit tests.

Signed-off-by: txaty <[email protected]>

Author: Tommy TIAN

.gitignore

@@ -10,6 +10,9 @@
 # Test binary, built with `go test -c`
 *.test
 
+# coverage files
+coverage.*
+
 # Output of the go coverage tool, specifically when used with LiteIDE
 *.out
 *.html

README.md

@@ -56,6 +56,8 @@ func NewHashFunc(data []byte) ([]byte, error) {
 }
 ```
 
+*Important Notice*: please make sure the hash function used by paralleled algorithms is concurrent-safe.
+
 ## Example
 
 ### Proof generation and verification of all blocks

default_hash.go

@@ -0,0 +1,53 @@
+// MIT License
+//
+// Copyright (c) 2022 Tommy TIAN
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in all
+// copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+// SOFTWARE.
+
+package merkletree
+
+import (
+	"crypto/sha256"
+)
+
+var sha256Digest = sha256.New()
+
+// defaultHashFunc is used when no user hash function is specified.
+// It implements SHA256 hash function.
+func defaultHashFunc(data []byte) ([]byte, error) {
+	defer sha256Digest.Reset()
+	_, err := sha256Digest.Write(data)
+	if err != nil {
+		return nil, err
+	}
+	return sha256Digest.Sum(nil), nil
+
+}
+
+// defaultHashFuncParal is used by parallel algorithms when no user hash function is specified.
+// It implements SHA256 hash function.
+// When implementing hash functions for paralleled algorithms, please make sure it is concurrent safe.
+func defaultHashFuncParal(data []byte) ([]byte, error) {
+	digest := sha256.New()
+	_, err := digest.Write(data)
+	if err != nil {
+		return nil, err
+	}
+	return digest.Sum(nil), nil
+}

default_hash_test.go

@@ -0,0 +1,72 @@
+package merkletree
+
+import (
+	"crypto/sha256"
+	"errors"
+	"hash"
+	"reflect"
+	"testing"
+)
+
+type mockHash struct {
+}
+
+func NewMock() hash.Hash {
+	return &mockHash{}
+}
+
+func (m *mockHash) Reset() {
+}
+
+func (m *mockHash) BlockSize() int {
+	return 0
+}
+
+func (m *mockHash) Size() int {
+	return 0
+}
+
+func (m *mockHash) Sum(in []byte) []byte {
+	return in
+}
+
+func (m *mockHash) Write([]byte) (nn int, err error) {
+	return 0, errors.New("mockHash.Write error")
+}
+
+func Test_defaultHashFunc(t *testing.T) {
+	sha256Digest = NewMock()
+	defer func() {
+		sha256Digest = sha256.New()
+	}()
+	type args struct {
+		data []byte
+	}
+	tests := []struct {
+		name    string
+		args    args
+		want    []byte
+		wantErr bool
+	}{
+		{
+			name: "test_write_err",
+			args: args{
+				data: []byte{},
+			},
+			want:    nil,
+			wantErr: true,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := defaultHashFunc(tt.args.data)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("defaultHashFunc() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+			if !reflect.DeepEqual(got, tt.want) {
+				t.Errorf("defaultHashFunc() got = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}

merkle_tree.go

@@ -25,7 +25,6 @@ package merkletree
 import (
 	"bytes"
 	"crypto/rand"
-	"crypto/sha256"
 	"errors"
 	"math"
 	"runtime"
@@ -109,7 +108,11 @@ func New(config *Config, blocks []DataBlock) (m *MerkleTree, err error) {
 		config = new(Config)
 	}
 	if config.HashFunc == nil {
-		config.HashFunc = defaultHashFunc
+		if config.RunInParallel {
+			config.HashFunc = defaultHashFuncParal
+		} else {
+			config.HashFunc = defaultHashFunc
+		}
 	}
 	// If the configuration mode is not set, then set it to ModeProofGen by default.
 	if config.Mode == 0 {
@@ -389,14 +392,6 @@ func getDummyHash() ([]byte, error) {
 	return dummyBytes, nil
 }
 
-// defaultHashFunc is used when no user hash function is specified.
-// It implements SHA256 hash function.
-func defaultHashFunc(data []byte) ([]byte, error) {
-	sha256Func := sha256.New()
-	sha256Func.Write(data)
-	return sha256Func.Sum(nil), nil
-}
-
 func (m *MerkleTree) leafGen(blocks []DataBlock) ([][]byte, error) {
 	var (
 		lenLeaves = len(blocks)

merkle_tree_test.go

@@ -34,7 +34,7 @@ import (
 	"github.com/agiledragon/gomonkey/v2"
 )
 
-const benchSize = 100000
+const benchSize = 10000
 
 type mockDataBlock struct {
 	data []byte