Add examples for encrypt/decrypt; Fix merkle tree

Author: pyropy

example/README.md

@@ -0,0 +1 @@
+## TODO: Add usage examples

example/main.go

@@ -1,63 +1,126 @@
 package main
 
 import (
+	"encoding/hex"
 	"flag"
 	"fmt"
-	bitstream "github.com/pyropy/bitstream-go/pkg"
 	"io"
 	"log"
 	"os"
+
+	"github.com/btcsuite/btcd/btcec/v2"
+	bitstream "github.com/pyropy/bitstream-go/pkg"
 )
 
 var (
-	inPath    string
-	outPath   string
-	preimage  string
-	chunkSize int64
+	action      string
+	inPath      string
+	outPath     string
+	preimage    string
+	paymentHash string
+	chunkSize   int64
+	pk          *btcec.PrivateKey
 )
 
 func init() {
+	flag.StringVar(&action, "action", "encrypt", "encrypt or decrypt")
 	flag.StringVar(&inPath, "in", "", "File to encrypt/decrypt")
 	flag.StringVar(&outPath, "out", "", "Path to output encrypted/decrypted file")
 	flag.StringVar(&preimage, "preimage", "", "Preimage to use for encryption/decryption")
+	flag.StringVar(&paymentHash, "hash", "", "Payment hash")
 	flag.Int64Var(&chunkSize, "size", 32, "Size of chunks to encrypt/decrypt")
+
+	pkBytes, err := hex.DecodeString("80faa7d1c150a903a4028bf87ca8800aff507b24df90e8434bfa1f34d639c053")
+	if err != nil {
+		panic(err)
+	}
+
+	pk, _ = btcec.PrivKeyFromBytes(pkBytes)
 }
 
 // TODO: Implement encrypt and decrypt cli commands
 func main() {
 	flag.Parse()
 
-	err := encrypt(inPath, outPath)
+	var err error
+
+	switch action {
+	case "encrypt":
+		err = encrypt(inPath, outPath, preimage, chunkSize)
+	case "decrypt":
+		err = decrypt(inPath, outPath, preimage, chunkSize)
+	}
+
 	if err != nil {
 		panic(err)
 	}
 }
 
-func encrypt(inPath, outPath string) error {
+func encrypt(inPath, outPath, preimage string, chunkSize int64) error {
 	f, err := os.Open(inPath)
 	if err != nil {
 		return err
 	}
 
-	encryptedFile, tree, err := bitstream.EncryptFile([]byte(preimage), f, chunkSize)
+	defer f.Close()
+
+	out, err := os.OpenFile(outPath, os.O_RDWR|os.O_CREATE, 0644)
+	if err != nil {
+		return err
+	}
+
+	preimageBytes, err := hex.DecodeString(preimage)
+	if err != nil {
+		panic(err)
+	}
+
+	paymentHashBytes, err := hex.DecodeString(paymentHash)
+	if err != nil {
+		return err
+	}
+
+	err = bitstream.Encrypt(pk, paymentHashBytes, preimageBytes, f, out, chunkSize)
 	if err != nil {
 		return err
 	}
 
-	l := fmt.Sprintf("Merkle tree hash 0x%x", tree.GetHash())
-	log.Println(l)
+	return nil
+}
 
-	b, err := io.ReadAll(encryptedFile)
+func decrypt(inPath, outPath, preimage string, chunkSize int64) error {
+	f, err := os.Open(inPath)
 	if err != nil {
 		return err
 	}
 
-	err = os.WriteFile(outPath, b, 0644)
+	defer f.Close()
+
+	preimageBytes, err := hex.DecodeString(preimage)
+	if err != nil {
+		return err
+	}
+
+	decryptedFile, index, err := bitstream.Decrypt(preimageBytes, f, chunkSize)
+	if err != nil {
+		if index != -1 {
+			proof, err := bitstream.GenerateProof(f, chunkSize, 2*index)
+			if err != nil {
+				return err
+			}
+
+			root := fmt.Sprintf("%x", proof.MerkleProof.Root)
+			log.Println("Generated proof", "proof root", root)
+		}
+
+		return err
+	}
+
+	decrypted, err := io.ReadAll(decryptedFile)
 	if err != nil {
 		return err
 	}
 
-	err = f.Close()
+	err = os.WriteFile(outPath, decrypted, 0644)
 	if err != nil {
 		return err
 	}

pkg/cipher.go

@@ -11,7 +11,8 @@ func ChunkCipher(index uint64, preimage []byte, data []byte) []byte {
 
 	// add 1 to index
 	index += 1
-	// convert index to bytes
+
+	// convert index to bytes (original implementation uses btc specific encoding)
 	indexBytes := make([]byte, 8)
 	binary.LittleEndian.PutUint64(indexBytes, index)
 

pkg/decrypt.go

@@ -13,7 +13,7 @@ import (
 //
 // If chunk index is not -1, then error ocurred while decrypting the chunk
 // hence proof for given chunk should be generated.
-func Decrypt(preimage []byte, file io.ReadSeeker, chunkSize int) (io.Reader, int, error) {
+func Decrypt(preimage []byte, file io.ReadSeeker, chunkSize int64) (io.Reader, int, error) {
 	var out bytes.Buffer
 
 	expectedHash := make([]byte, HashSize)
@@ -41,7 +41,7 @@ func Decrypt(preimage []byte, file io.ReadSeeker, chunkSize int) (io.Reader, int
 		}
 
 		// update offset for chunk read
-		offset += int64(chunkSize)
+		offset += chunkSize
 		_, err = file.Seek(offset, 0)
 		if err != nil {
 			return nil, -1, err
@@ -58,7 +58,7 @@ func Decrypt(preimage []byte, file io.ReadSeeker, chunkSize int) (io.Reader, int
 		}
 
 		// update offset for next expected hash read
-		offset += int64(chunkSize)
+		offset += chunkSize
 
 		// decrypt chunk, compute hash and if hashes match
 		decryptedChunk := ChunkCipher(i, preimage, encryptedChunk)

pkg/encrypt.go

@@ -3,28 +3,28 @@ package bitstream
 import (
 	"bytes"
 	"fmt"
+	"io"
+
 	"github.com/btcsuite/btcd/btcec/v2"
 	"github.com/btcsuite/btcd/btcec/v2/schnorr"
-	"io"
 )
 
 type EncryptedFile interface {
 	io.WriterTo
 	io.Reader
 }
 
-func Encrypt(pk *btcec.PrivateKey, paymentHash []byte, preimage []byte, inFile io.ReadSeeker, outFile io.Writer, chunkSize int64) error {
+func Encrypt(pk *btcec.PrivateKey, paymentHash, preimage []byte, inFile io.ReadSeeker, outFile io.Writer, chunkSize int64) error {
 	encrypted, tree, err := EncryptFile(preimage, inFile, chunkSize)
 	if err != nil {
 		return err
 	}
 
 	encryptedRoot := tree.GetHash()
-	message := append(encryptedRoot, paymentHash...)
-
-	sig, err := schnorr.Sign(pk, message)
+	// original implementation signs encrypted root + payment hash
+	sig, err := schnorr.Sign(pk, encryptedRoot)
 	if err != nil {
-		return err
+		return fmt.Errorf("failed to generate schorr sig: %w", err)
 	}
 
 	_, err = outFile.Write(sig.Serialize())

pkg/merkle_tree.go

@@ -23,9 +23,8 @@ type Node struct {
 }
 
 type MerkleProof struct {
-	Root  []byte
-	Bloom []uint8
-	Path  [][]byte
+	Root []byte
+	Path [][]byte
 }
 
 func (n *Node) GetHash() []byte {
@@ -65,14 +64,14 @@ func NewTree(leaves []*Node) *MerkleTree {
 		}
 	}
 
-	if len(leaves)%2 != 0 {
-		leaf := leaves[len(leaves)-1]
-		leaves = append(leaves, leaf)
-	}
-
 	for len(leaves) > 1 {
 		var nextLevel []*Node
 
+		if len(leaves)%2 != 0 {
+			leaf := leaves[len(leaves)-1]
+			leaves = append(leaves, leaf)
+		}
+
 		// go through all leaves on current level in pairs
 		for i := 0; i < len(leaves); i += 2 {
 			left := leaves[i]
@@ -94,14 +93,14 @@ func NewTree(leaves []*Node) *MerkleTree {
 func GenerateMerkleProof(leaves []*Node, index int) *MerkleProof {
 	var path [][]byte
 
-	if len(leaves)%2 != 0 {
-		leaf := leaves[len(leaves)-1]
-		leaves = append(leaves, leaf)
-	}
-
 	for len(leaves) > 1 {
 		var nextLevel []*Node
 
+		if len(leaves)%2 != 0 {
+			leaf := leaves[len(leaves)-1]
+			leaves = append(leaves, leaf)
+		}
+
 		// go through all leaves on current level in pairs
 		for i := 0; i < len(leaves); i += 2 {
 			left := leaves[i]

pkg/merkle_tree_test.go

@@ -21,13 +21,13 @@ func TestTree(t *testing.T) {
 			leaves: []*Node{
 				NewNode([]byte("hello")),
 			},
-			expect: "1d25c19a1a3fb65c78d018561057362916c14bfd36b75aa8cb0f4d696293b183",
+			expect: "2cf24dba5fb0a30e26e83b2ac5b9e29e1b161e5c1fa7425e73043362938b9824",
 		},
 		{
 			leaves: []*Node{
 				NewNode([]byte("world")),
 			},
-			expect: "a9e049081e9eaeef2d30029031314c888955d7181a032c29b11633ac11d6076c",
+			expect: "486ea46224d1bb4fb680f34f7c9ad96a8f24ec88be73ea8e5a6c65260e9cb8a7",
 		},
 		{
 			leaves: []*Node{

pkg/proof.go

@@ -8,12 +8,12 @@ const (
 )
 
 type Proof struct {
-	sig         []byte
-	paymentHash []byte
-	merkleProof *MerkleProof
+	Sig         []byte
+	PaymentHash []byte
+	MerkleProof *MerkleProof
 }
 
-func GenerateProof(file io.ReadSeeker, chunkSize int, chunkIndex int) (*Proof, error) {
+func GenerateProof(file io.ReadSeeker, chunkSize int64, chunkIndex int) (*Proof, error) {
 	sig := make([]byte, SignatureSize)
 	paymentHash := make([]byte, HashSize)
 	hash := make([]byte, chunkSize)
@@ -53,7 +53,7 @@ func GenerateProof(file io.ReadSeeker, chunkSize int, chunkIndex int) (*Proof, e
 		}
 
 		// update offset for next chunk hash read
-		offset += int64(chunkSize * 2)
+		offset += chunkSize * 2
 		node := NewNodeFromHash(hash)
 		leaves = append(leaves, node)
 
@@ -64,8 +64,8 @@ func GenerateProof(file io.ReadSeeker, chunkSize int, chunkIndex int) (*Proof, e
 	merkleProof := GenerateMerkleProof(leaves, chunkIndex)
 
 	return &Proof{
-		sig:         sig,
-		paymentHash: paymentHash,
-		merkleProof: merkleProof,
+		Sig:         sig,
+		PaymentHash: paymentHash,
+		MerkleProof: merkleProof,
 	}, nil
 }