github.com/anacrolix/torrent@v1.61.0/merkle/hash.go (about)

     1  package merkle
     2  
     3  import (
     4  	"crypto/sha256"
     5  	"hash"
     6  	"unsafe"
     7  )
     8  
     9  func NewHash() *Hash {
    10  	h := &Hash{
    11  		nextBlock: sha256.New(),
    12  	}
    13  	return h
    14  }
    15  
    16  type Hash struct {
    17  	blocks    [][32]byte
    18  	nextBlock hash.Hash
    19  	// How many bytes have been written to nextBlock so far.
    20  	nextBlockWritten int
    21  }
    22  
    23  func (h *Hash) remaining() int {
    24  	return BlockSize - h.nextBlockWritten
    25  }
    26  
    27  func (h *Hash) Write(p []byte) (n int, err error) {
    28  	for len(p) > 0 {
    29  		var n1 int
    30  		n1, err = h.nextBlock.Write(p[:min(len(p), h.remaining())])
    31  		n += n1
    32  		h.nextBlockWritten += n1
    33  		p = p[n1:]
    34  		if h.remaining() == 0 {
    35  			h.blocks = append(h.blocks, h.nextBlockSum())
    36  			h.nextBlock.Reset()
    37  			h.nextBlockWritten = 0
    38  		}
    39  		if err != nil {
    40  			break
    41  		}
    42  	}
    43  	return
    44  }
    45  
    46  func (h *Hash) nextBlockSum() (sum [32]byte) {
    47  	if unsafe.SliceData(h.nextBlock.Sum(sum[:0])) != unsafe.SliceData(sum[:]) {
    48  		panic("go sux")
    49  	}
    50  	return
    51  }
    52  
    53  func (h *Hash) curBlocks() [][32]byte {
    54  	blocks := h.blocks
    55  	if h.nextBlockWritten != 0 {
    56  		blocks = append(blocks, h.nextBlockSum())
    57  	}
    58  	return blocks
    59  }
    60  
    61  func (h *Hash) Sum(b []byte) []byte {
    62  	sum := RootWithPadHash(h.curBlocks(), [32]byte{})
    63  	return append(b, sum[:]...)
    64  }
    65  
    66  // Sums by extending with zero hashes for blocks missing to meet the given length. Necessary for
    67  // piece layers hashes for file tail blocks that don't pad to the piece length.
    68  func (h *Hash) SumMinLength(b []byte, length int) []byte {
    69  	blocks := h.curBlocks()
    70  	minBlocks := (length + BlockSize - 1) / BlockSize
    71  	blocks = append(blocks, make([][32]byte, minBlocks-len(blocks))...)
    72  	sum := RootWithPadHash(blocks, [32]byte{})
    73  	return append(b, sum[:]...)
    74  }
    75  
    76  func (h *Hash) Reset() {
    77  	h.blocks = h.blocks[:0]
    78  	h.nextBlock.Reset()
    79  	h.nextBlockWritten = 0
    80  }
    81  
    82  func (h *Hash) Size() int {
    83  	return 32
    84  }
    85  
    86  func (h *Hash) BlockSize() int {
    87  	return h.nextBlock.BlockSize()
    88  }
    89  
    90  var _ hash.Hash = (*Hash)(nil)