Merge pull request #2 from everFinance/merkle-add

feat(): add merkle

Author: outprog

.gitignore

@@ -1,2 +1,3 @@
 main.go
-test-keyfile.json
+test-keyfile.json
+.DS_Store

Makefile

@@ -0,0 +1,2 @@
+test:
+	sh ./test.sh

README.md

@@ -74,3 +74,11 @@ func main() {
 	fmt.Println(id, stat, err) // {{id}}, Pending, nil
 }
 ```
+
+### Development
+
+#### Test
+
+```
+make test
+```

go.mod

@@ -3,4 +3,8 @@ module github.com/everFinance/goar
 go 1.14
 
 require github.com/stretchr/testify v1.6.1
-require github.com/everFinance/gojwk v1.0.0
+
+require (
+	github.com/everFinance/gojwk v1.0.0
+	github.com/shopspring/decimal v1.2.0
+)

go.sum

@@ -4,6 +4,8 @@ github.com/everFinance/gojwk v1.0.0 h1:le/oI2NgXlrqg3MHU6ka+V30EWcD7TD6+Ilh+go79
 github.com/everFinance/gojwk v1.0.0/go.mod h1:icXSXsIdpAczlpAtSljQlmABkMTRZENr73KHmo0GOGc=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/shopspring/decimal v1.2.0 h1:abSATXmQEYyShuxI4/vyW3tV1MrKAJzCZ/0zLUXYbsQ=
+github.com/shopspring/decimal v1.2.0/go.mod h1:DKyhrW/HYNuLGql+MJL6WCR6knT2jwCFRcu2hWCYk4o=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/testify v1.6.1 h1:hDPOHmpOpP40lSULcqw7IrRb/u7w6RpDC9399XyoNd0=
 github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=

test.sh

@@ -0,0 +1,10 @@
+#! /bin/bash
+
+for i in $(find . -iname "*_test.go" -exec dirname {} \; | uniq)
+do
+    go test -race -cover $i;
+    if [ $? != 0 ]
+    then
+      return 1
+    fi
+done

utils/merkle.go

@@ -2,7 +2,10 @@ package utils
 
 import (
 	"crypto/sha256"
+	"math"
 	"math/big"
+
+	"github.com/shopspring/decimal"
 )
 
 const (
@@ -12,13 +15,15 @@ const (
 	// number of bits in a big.Word
 	wordBits = 32 << (uint64(^big.Word(0)) >> 63)
 	// number of bytes in a big.Word
-	wordBytes = wordBits / 8
+	wordBytes      = wordBits / 8
+	BranchNodeType = "branch"
+	LeafNodeType   = "leaf"
 )
 
 type Chunks struct {
 	DataRoot []byte
 	Chunks   []Chunk
-	Proofs   []Proof
+	Proofs   []*Proof
 }
 
 type Chunk struct {
@@ -27,68 +32,179 @@ type Chunk struct {
 	MaxByteRange int
 }
 
-type LeafNode struct {
+// Node include leaf node and branch node
+type Node struct {
 	ID           []byte
-	DataHash     []byte
-	Type         string
-	MinByteRange int
+	Type         string // "branch" or "leaf"
+	DataHash     []byte // only leaf node
+	MinByteRange int    // only leaf node
 	MaxByteRange int
+	ByteRange    int   // only branch node
+	LeftChild    *Node // only branch node
+	RightChild   *Node // only branch node
 }
 
 type Proof struct {
 	Offest int
 	Proof  []byte
 }
 
-func GenerateChunks(data []byte) (c Chunks) {
+func GenerateChunks(data []byte) Chunks {
 	chunks := chunkData(data)
 	leaves := generateLeaves(chunks)
-	// TODO, calculate root & proofs
+	root := buildLayer(leaves, 0) // leaf node level == 0
+	proofs := generateProofs(root)
+
+	// Discard the last chunk & proof if it's zero length.
+	lastChunk := chunks[len(chunks)-1]
+	if lastChunk.MaxByteRange-lastChunk.MinByteRange == 0 {
+		chunks = chunks[:len(chunks)-1]
+		proofs = proofs[:len(proofs)-1]
+	}
 
-	c.DataRoot = leaves[0].ID
-	return
+	return Chunks{
+		DataRoot: root.ID,
+		Chunks:   chunks,
+		Proofs:   proofs,
+	}
 }
 
 func chunkData(data []byte) (chunks []Chunk) {
 	cursor := 0
-	// TODO
-	// for len(data) >= MAX_CHUNK_SIZE {
-	// }
+	var rest = data
+	// if data length > max size
+	for len(rest) >= MAX_CHUNK_SIZE {
+		chunkSize := MAX_CHUNK_SIZE
+
+		// 查看下一轮的chunkSize 是否小于最小的size，如果是则在这轮中调整chunk size 的大小
+		nextChunkSize := len(rest) - MAX_CHUNK_SIZE
+		if nextChunkSize > 0 && nextChunkSize < MIN_CHUNK_SIZE {
+			dec := decimal.NewFromFloat(math.Ceil(float64(len(rest) / 2)))
+			chunkSize = int(dec.IntPart())
+		}
+
+		chunk := rest[:chunkSize]
+		dataHash := sha256.Sum256(chunk)
+		cursor += len(chunk)
+		chunks = append(chunks, Chunk{
+			DataHash:     dataHash[:],
+			MinByteRange: cursor - len(chunk),
+			MaxByteRange: cursor,
+		})
 
-	hash := sha256.Sum256(data)
+		rest = rest[chunkSize:]
+	}
+
+	hash := sha256.Sum256(rest)
 	chunks = append(chunks, Chunk{
 		DataHash:     hash[:],
 		MinByteRange: cursor,
-		MaxByteRange: cursor + len(data),
+		MaxByteRange: cursor + len(rest),
 	})
 	return
 }
 
-func generateLeaves(chunks []Chunk) (leafs []LeafNode) {
+func generateLeaves(chunks []Chunk) (leafs []*Node) {
 	for _, chunk := range chunks {
 		hDataHash := sha256.Sum256(chunk.DataHash)
 		hMaxByteRange := sha256.Sum256(PaddedBigBytes(big.NewInt(int64(chunk.MaxByteRange)), 32))
 
-		leafs = append(leafs, LeafNode{
+		leafs = append(leafs, &Node{
 			ID: hashArray(
 				[][]byte{hDataHash[:], hMaxByteRange[:]},
 			),
+			Type:         LeafNodeType,
 			DataHash:     chunk.DataHash,
-			Type:         "leaf",
 			MinByteRange: chunk.MinByteRange,
 			MaxByteRange: chunk.MaxByteRange,
+			ByteRange:    0,
+			LeftChild:    nil,
+			RightChild:   nil,
 		})
 	}
+	return
+}
+
+// buildLayer
+func buildLayer(nodes []*Node, level int) (root *Node) {
+	if len(nodes) == 1 {
+		root = nodes[0]
+		return
+	}
+
+	nextLayer := make([]*Node, 0, len(nodes)/2)
+	for i := 0; i < len(nodes); i += 2 {
+		leftNode := nodes[i]
+		var rightNode *Node
+		if i+1 < len(nodes) {
+			rightNode = nodes[i+1]
+		}
+		nextLayer = append(nextLayer, hashBranch(leftNode, rightNode))
+	}
+
+	return buildLayer(nextLayer, level+1)
+}
+
+// hashBranch get branch node by child node
+func hashBranch(leftNode, rightNode *Node) (branchNode *Node) {
+	// 如果只有一个node，则该node 为branch node
+	if rightNode == nil {
+		return leftNode
+	}
+	hLeafNodeId := sha256.Sum256(leftNode.ID)
+	hRightNodeId := sha256.Sum256(rightNode.ID)
+	hLeafNodeMaxRange := sha256.Sum256(PaddedBigBytes(big.NewInt(int64(leftNode.MaxByteRange)), 32))
+	id := hashArray([][]byte{hLeafNodeId[:], hRightNodeId[:], hLeafNodeMaxRange[:]})
+
+	branchNode = &Node{
+		Type:         BranchNodeType,
+		ID:           id,
+		DataHash:     nil,
+		MinByteRange: 0,
+		MaxByteRange: rightNode.MaxByteRange,
+		ByteRange:    leftNode.MaxByteRange,
+		LeftChild:    leftNode,
+		RightChild:   rightNode,
+	}
 
 	return
 }
 
-func buildLayer() {
-	//TODO
+func generateProofs(rootNode *Node) []*Proof {
+	return resolveBranchProofs(rootNode, []byte{}, 0)
 }
 
-func generateProofs() {
-	// TODO
+// resolveBranchProofs 从root node 递归搜索叶子节点并为其生成证明
+func resolveBranchProofs(node *Node, proof []byte, depth int) (proofs []*Proof) {
+
+	if node.Type == LeafNodeType {
+		p := &Proof{
+			Offest: node.MaxByteRange - 1,
+			Proof: ConcatBuffer(
+				proof,
+				node.DataHash,
+				PaddedBigBytes(big.NewInt(int64(node.MaxByteRange)), 32),
+			),
+		}
+		proofs = append(proofs, p)
+		return
+	}
+
+	if node.Type == BranchNodeType {
+		partialProof := ConcatBuffer(
+			proof,
+			node.LeftChild.ID,
+			node.RightChild.ID,
+			PaddedBigBytes(big.NewInt(int64(node.ByteRange)), 32),
+		)
+		leftProofs := resolveBranchProofs(node.LeftChild, partialProof, depth+1)
+		rightProofs := resolveBranchProofs(node.RightChild, partialProof, depth+1)
+		proofs = append(append(proofs, leftProofs...), rightProofs...)
+		return
+	}
+
+	// node type error then return nil
+	return nil
 }
 
 func hashArray(data [][]byte) []byte {
@@ -120,3 +236,18 @@ func ReadBits(bigint *big.Int, buf []byte) {
 		}
 	}
 }
+
+// ConcatBuffer
+func ConcatBuffer(buffers ...[]byte) []byte {
+	totalLength := 0
+	for i := 0; i < len(buffers); i++ {
+		totalLength += len(buffers[i])
+	}
+
+	temp := make([]byte, 0, totalLength)
+
+	for _, val := range buffers {
+		temp = append(temp, val...)
+	}
+	return temp
+}

utils/merkle_test.go

@@ -1,6 +1,7 @@
 package utils
 
 import (
+	"io/ioutil"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -9,6 +10,11 @@ import (
 func TestGenerateChunks(t *testing.T) {
 	data, _ := Base64Decode("NzcyNg")
 	assert.Equal(t, "z3rQGxyiqdQuOh2dxDst176oOKmW3S9MwQNTEh4DK1U", Base64Encode(GenerateChunks(data).DataRoot))
+
+	data, err := ioutil.ReadFile("./testfile/1mb.bin")
+	assert.NoError(t, err)
+	chunks := GenerateChunks(data)
+	assert.Equal(t, "o1tTTjbC7hIZN6KbUUYjlkQoDl2k8VXNuBDcGIs52Hc", Base64Encode(chunks.DataRoot))
 }
 
 func TestChunkData(t *testing.T) {