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)