src/mth.go

   1 /*
   2 NNCP -- Node to Node copy, utilities for store-and-forward data exchange
   3 Copyright (C) 2016-2021 Sergey Matveev <stargrave@stargrave.org>
   4
   5 This program is free software: you can redistribute it and/or modify
   6 it under the terms of the GNU General Public License as published by
   7 the Free Software Foundation, version 3 of the License.
   8
   9 This program is distributed in the hope that it will be useful,
  10 but WITHOUT ANY WARRANTY; without even the implied warranty of
  11 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  12 GNU General Public License for more details.
  13
  14 You should have received a copy of the GNU General Public License
  15 along with this program.  If not, see <http://www.gnu.org/licenses/>.
  16 */
  17
  18 package nncp
  19
  20 import (
  21         "bytes"
  22         "errors"
  23         "io"
  24
  25         "lukechampine.com/blake3"
  26 )
  27
  28 const (
  29         MTHBlockSize = 128 * 1024
  30         MTHSize      = 32
  31 )
  32
  33 var (
  34         MTHLeafKey = blake3.Sum256([]byte("NNCP MTH LEAF"))
  35         MTHNodeKey = blake3.Sum256([]byte("NNCP MTH NODE"))
  36 )
  37
  38 type MTHEventType uint8
  39
  40 const (
  41         MTHEventAppend  MTHEventType = iota
  42         MTHEventPrepend MTHEventType = iota
  43         MTHEventFold    MTHEventType = iota
  44 )
  45
  46 type MTHEvent struct {
  47         Type  MTHEventType
  48         Level int
  49         Ctr   int
  50         Hsh   []byte
  51 }
  52
  53 type MTH struct {
  54         size        int64
  55         PrependSize int64
  56         skip        int64
  57         skipped     bool
  58         hasher      *blake3.Hasher
  59         hashes      [][MTHSize]byte
  60         buf         *bytes.Buffer
  61         finished    bool
  62         Events      chan MTHEvent
  63         PktName     string
  64 }
  65
  66 func MTHNew(size, offset int64) *MTH {
  67         mth := MTH{
  68                 hasher: blake3.New(MTHSize, MTHLeafKey[:]),
  69                 buf:    bytes.NewBuffer(make([]byte, 0, 2*MTHBlockSize)),
  70         }
  71         if size == 0 {
  72                 return &mth
  73         }
  74         prepends := int(offset / MTHBlockSize)
  75         skip := MTHBlockSize - (offset - int64(prepends)*MTHBlockSize)
  76         if skip == MTHBlockSize {
  77                 skip = 0
  78         } else if skip > 0 {
  79                 prepends++
  80         }
  81         prependSize := int64(prepends * MTHBlockSize)
  82         if prependSize > size {
  83                 prependSize = size
  84         }
  85         if offset+skip > size {
  86                 skip = size - offset
  87         }
  88         mth.size = size
  89         mth.PrependSize = prependSize
  90         mth.skip = skip
  91         mth.hashes = make([][MTHSize]byte, prepends, 1+size/MTHBlockSize)
  92         return &mth
  93 }
  94
  95 func (mth *MTH) Reset() { panic("not implemented") }
  96
  97 func (mth *MTH) Size() int { return MTHSize }
  98
  99 func (mth *MTH) BlockSize() int { return MTHBlockSize }
 100
 101 func (mth *MTH) Write(data []byte) (int, error) {
 102         if mth.finished {
 103                 return 0, errors.New("already Sum()ed")
 104         }
 105         n, err := mth.buf.Write(data)
 106         if err != nil {
 107                 return n, err
 108         }
 109         if mth.skip > 0 && int64(mth.buf.Len()) >= mth.skip {
 110                 mth.buf.Next(int(mth.skip))
 111                 mth.skip = 0
 112         }
 113         for mth.buf.Len() >= MTHBlockSize {
 114                 if _, err = mth.hasher.Write(mth.buf.Next(MTHBlockSize)); err != nil {
 115                         return n, err
 116                 }
 117                 h := new([MTHSize]byte)
 118                 mth.hasher.Sum(h[:0])
 119                 mth.hasher.Reset()
 120                 mth.hashes = append(mth.hashes, *h)
 121                 if mth.Events != nil {
 122                         mth.Events <- MTHEvent{
 123                                 MTHEventAppend,
 124                                 0, len(mth.hashes) - 1,
 125                                 mth.hashes[len(mth.hashes)-1][:],
 126                         }
 127                 }
 128         }
 129         return n, err
 130 }
 131
 132 func (mth *MTH) PrependFrom(r io.Reader) (int, error) {
 133         if mth.finished {
 134                 return 0, errors.New("already Sum()ed")
 135         }
 136         var err error
 137         buf := make([]byte, MTHBlockSize)
 138         var i, n, read int
 139         fullsize := mth.PrependSize
 140         les := LEs{{"Pkt", mth.PktName}, {"FullSize", fullsize}, {"Size", 0}}
 141         for mth.PrependSize >= MTHBlockSize {
 142                 n, err = io.ReadFull(r, buf)
 143                 read += n
 144                 mth.PrependSize -= MTHBlockSize
 145                 if err != nil {
 146                         return read, err
 147                 }
 148                 if _, err = mth.hasher.Write(buf); err != nil {
 149                         panic(err)
 150                 }
 151                 mth.hasher.Sum(mth.hashes[i][:0])
 152                 mth.hasher.Reset()
 153                 if mth.Events != nil {
 154                         mth.Events <- MTHEvent{MTHEventPrepend, 0, i, mth.hashes[i][:]}
 155                 }
 156                 if mth.PktName != "" {
 157                         les[len(les)-1].V = int64(read)
 158                         Progress("check", les)
 159                 }
 160                 i++
 161         }
 162         if mth.PrependSize > 0 {
 163                 n, err = io.ReadFull(r, buf[:mth.PrependSize])
 164                 read += n
 165                 if err != nil {
 166                         return read, err
 167                 }
 168                 if _, err = mth.hasher.Write(buf[:mth.PrependSize]); err != nil {
 169                         panic(err)
 170                 }
 171                 mth.hasher.Sum(mth.hashes[i][:0])
 172                 mth.hasher.Reset()
 173                 if mth.Events != nil {
 174                         mth.Events <- MTHEvent{MTHEventPrepend, 0, i, mth.hashes[i][:]}
 175                 }
 176                 if mth.PktName != "" {
 177                         les[len(les)-1].V = fullsize
 178                         Progress("check", les)
 179                 }
 180         }
 181         return read, nil
 182 }
 183
 184 func (mth *MTH) Sum(b []byte) []byte {
 185         if mth.finished {
 186                 return append(b, mth.hashes[0][:]...)
 187         }
 188         if mth.buf.Len() > 0 {
 189                 b := mth.buf.Next(MTHBlockSize)
 190                 if _, err := mth.hasher.Write(b); err != nil {
 191                         panic(err)
 192                 }
 193                 h := new([MTHSize]byte)
 194                 mth.hasher.Sum(h[:0])
 195                 mth.hasher.Reset()
 196                 mth.hashes = append(mth.hashes, *h)
 197                 if mth.Events != nil {
 198                         mth.Events <- MTHEvent{
 199                                 MTHEventAppend,
 200                                 0, len(mth.hashes) - 1,
 201                                 mth.hashes[len(mth.hashes)-1][:],
 202                         }
 203                 }
 204         }
 205         switch len(mth.hashes) {
 206         case 0:
 207                 h := new([MTHSize]byte)
 208                 if _, err := mth.hasher.Write(nil); err != nil {
 209                         panic(err)
 210                 }
 211                 mth.hasher.Sum(h[:0])
 212                 mth.hasher.Reset()
 213                 mth.hashes = append(mth.hashes, *h)
 214                 if mth.Events != nil {
 215                         mth.Events <- MTHEvent{MTHEventAppend, 0, 0, mth.hashes[0][:]}
 216                 }
 217                 fallthrough
 218         case 1:
 219                 mth.hashes = append(mth.hashes, mth.hashes[0])
 220                 if mth.Events != nil {
 221                         mth.Events <- MTHEvent{MTHEventAppend, 0, 1, mth.hashes[1][:]}
 222                 }
 223         }
 224         mth.hasher = blake3.New(MTHSize, MTHNodeKey[:])
 225         level := 1
 226         for len(mth.hashes) != 1 {
 227                 hashesUp := make([][MTHSize]byte, 0, 1+len(mth.hashes)/2)
 228                 pairs := (len(mth.hashes) / 2) * 2
 229                 for i := 0; i < pairs; i += 2 {
 230                         if _, err := mth.hasher.Write(mth.hashes[i][:]); err != nil {
 231                                 panic(err)
 232                         }
 233                         if _, err := mth.hasher.Write(mth.hashes[i+1][:]); err != nil {
 234                                 panic(err)
 235                         }
 236                         h := new([MTHSize]byte)
 237                         mth.hasher.Sum(h[:0])
 238                         mth.hasher.Reset()
 239                         hashesUp = append(hashesUp, *h)
 240                         if mth.Events != nil {
 241                                 mth.Events <- MTHEvent{
 242                                         MTHEventFold,
 243                                         level, len(hashesUp) - 1,
 244                                         hashesUp[len(hashesUp)-1][:],
 245                                 }
 246                         }
 247                 }
 248                 if len(mth.hashes)%2 == 1 {
 249                         hashesUp = append(hashesUp, mth.hashes[len(mth.hashes)-1])
 250                         if mth.Events != nil {
 251                                 mth.Events <- MTHEvent{
 252                                         MTHEventAppend,
 253                                         level, len(hashesUp) - 1,
 254                                         hashesUp[len(hashesUp)-1][:],
 255                                 }
 256                         }
 257                 }
 258                 mth.hashes = hashesUp
 259                 level++
 260         }
 261         mth.finished = true
 262         if mth.Events != nil {
 263                 close(mth.Events)
 264         }
 265         return append(b, mth.hashes[0][:]...)
 266 }