From 3673518b3089d51f4dfc97067dd00ee2a1bf8225 Mon Sep 17 00:00:00 2001
From: Andrew Ayer <andrew@sslmate.com>
Date: Thu, 5 Jan 2017 14:48:35 -0800
Subject: [PATCH] Add constructor functions for MerkleTreeBuilder

---
 auditing.go | 19 +++++++++++++++++++
 scanner.go  |  7 +++++--
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/auditing.go b/auditing.go
index fd2fc1d..69158bc 100644
--- a/auditing.go
+++ b/auditing.go
@@ -12,6 +12,7 @@ package certspotter
 import (
 	"bytes"
 	"crypto/sha256"
+	"errors"
 	"software.sslmate.com/src/certspotter/ct"
 )
 
@@ -141,6 +142,24 @@ type MerkleTreeBuilder struct {
 	numLeaves  uint64 // number of hashes added so far
 }
 
+func calculateStackSize (numLeaves uint64) int {
+	stackSize := 0
+	for numLeaves > 0 {
+		stackSize += int(numLeaves & 1)
+		numLeaves >>= 1
+	}
+	return stackSize
+}
+func EmptyMerkleTreeBuilder () *MerkleTreeBuilder {
+	return &MerkleTreeBuilder{}
+}
+func NewMerkleTreeBuilder (stack []ct.MerkleTreeNode, numLeaves uint64) (*MerkleTreeBuilder, error) {
+	if len(stack) != calculateStackSize(numLeaves) {
+		return nil, errors.New("NewMerkleTreeBuilder: incorrect stack size")
+	}
+	return &MerkleTreeBuilder{stack: stack, numLeaves: numLeaves}, nil
+}
+
 func (builder *MerkleTreeBuilder) Add(hash ct.MerkleTreeNode) {
 	builder.stack = append(builder.stack, hash)
 	builder.numLeaves++
diff --git a/scanner.go b/scanner.go
index 1331cec..f3e88a1 100644
--- a/scanner.go
+++ b/scanner.go
@@ -254,9 +254,12 @@ func (s *Scanner) MakeMerkleTreeBuilder(sth *ct.SignedTreeHead) (*MerkleTreeBuil
 			return nil, err
 		}
 		reverseHashes(auditPath)
-		builder = &MerkleTreeBuilder{numLeaves: sth.TreeSize - 1, stack: auditPath}
+		builder, err = NewMerkleTreeBuilder(auditPath, sth.TreeSize - 1)
+		if err != nil {
+			return nil, fmt.Errorf("Error returned bad audit proof for %x to %d", leafHash, sth.TreeSize)
+		}
 	} else {
-		builder = &MerkleTreeBuilder{numLeaves: 0}
+		builder = EmptyMerkleTreeBuilder()
 	}
 
 	builder.Add(leafHash)