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         "hash"
  24         "io"
  25
  26         "lukechampine.com/blake3"
  27 )
  28
  29 const (
  30         MTHBlockSize = 128 * 1024
  31         MTHSize      = 32
  32 )
  33
  34 var (
  35         MTHLeafKey = blake3.Sum256([]byte("NNCP MTH LEAF"))
  36         MTHNodeKey = blake3.Sum256([]byte("NNCP MTH NODE"))
  37 )
  38
  39 type MTHEventType uint8
  40
  41 const (
  42         MTHEventAppend  MTHEventType = iota
  43         MTHEventPrepend MTHEventType = iota
  44         MTHEventFold    MTHEventType = iota
  45 )
  46
  47 type MTHEvent struct {
  48         Type  MTHEventType
  49         Level int
  50         Ctr   int
  51         Hsh   []byte
  52 }
  53
  54 type MTH interface {
  55         hash.Hash
  56         PrependFrom(r io.Reader) (int, error)
  57         SetPktName(n string)
  58         PrependSize() int64
  59         Events() chan MTHEvent
  60 }
  61
  62 type MTHFat struct {
  63         size        int64
  64         prependSize int64
  65         skip        int64
  66         skipped     bool
  67         hasher      *blake3.Hasher
  68         hashes      [][MTHSize]byte
  69         buf         *bytes.Buffer
  70         finished    bool
  71         events      chan MTHEvent
  72         pktName     string
  73 }
  74
  75 func MTHFatNew(size, offset int64) MTH {
  76         mth := MTHFat{
  77                 hasher: blake3.New(MTHSize, MTHLeafKey[:]),
  78                 buf:    bytes.NewBuffer(make([]byte, 0, 2*MTHBlockSize)),
  79         }
  80         if size == 0 {
  81                 return &mth
  82         }
  83         prepends := int(offset / MTHBlockSize)
  84         skip := MTHBlockSize - (offset - int64(prepends)*MTHBlockSize)
  85         if skip == MTHBlockSize {
  86                 skip = 0
  87         } else if skip > 0 {
  88                 prepends++
  89         }
  90         prependSize := int64(prepends * MTHBlockSize)
  91         if prependSize > size {
  92                 prependSize = size
  93         }
  94         if offset+skip > size {
  95                 skip = size - offset
  96         }
  97         mth.size = size
  98         mth.prependSize = prependSize
  99         mth.skip = skip
 100         mth.hashes = make([][MTHSize]byte, prepends, 1+size/MTHBlockSize)
 101         return &mth
 102 }
 103
 104 func (mth *MTHFat) Events() chan MTHEvent {
 105         mth.events = make(chan MTHEvent)
 106         return mth.events
 107 }
 108
 109 func (mth *MTHFat) SetPktName(pktName string) { mth.pktName = pktName }
 110
 111 func (mth *MTHFat) PrependSize() int64 { return mth.prependSize }
 112
 113 func (mth *MTHFat) Reset() { panic("not implemented") }
 114
 115 func (mth *MTHFat) Size() int { return MTHSize }
 116
 117 func (mth *MTHFat) BlockSize() int { return MTHBlockSize }
 118
 119 func (mth *MTHFat) Write(data []byte) (int, error) {
 120         if mth.finished {
 121                 return 0, errors.New("already Sum()ed")
 122         }
 123         n, err := mth.buf.Write(data)
 124         if err != nil {
 125                 return n, err
 126         }
 127         if mth.skip > 0 && int64(mth.buf.Len()) >= mth.skip {
 128                 mth.buf.Next(int(mth.skip))
 129                 mth.skip = 0
 130         }
 131         for mth.buf.Len() >= MTHBlockSize {
 132                 if _, err = mth.hasher.Write(mth.buf.Next(MTHBlockSize)); err != nil {
 133                         return n, err
 134                 }
 135                 h := new([MTHSize]byte)
 136                 mth.hasher.Sum(h[:0])
 137                 mth.hasher.Reset()
 138                 mth.hashes = append(mth.hashes, *h)
 139                 if mth.events != nil {
 140                         mth.events <- MTHEvent{
 141                                 MTHEventAppend,
 142                                 0, len(mth.hashes) - 1,
 143                                 mth.hashes[len(mth.hashes)-1][:],
 144                         }
 145                 }
 146         }
 147         return n, err
 148 }
 149
 150 func (mth *MTHFat) PrependFrom(r io.Reader) (int, error) {
 151         if mth.finished {
 152                 return 0, errors.New("already Sum()ed")
 153         }
 154         var err error
 155         buf := make([]byte, MTHBlockSize)
 156         var i, n, read int
 157         fullsize := mth.prependSize
 158         les := LEs{{"Pkt", mth.pktName}, {"FullSize", fullsize}, {"Size", 0}}
 159         for mth.prependSize >= MTHBlockSize {
 160                 n, err = io.ReadFull(r, buf)
 161                 read += n
 162                 mth.prependSize -= MTHBlockSize
 163                 if err != nil {
 164                         return read, err
 165                 }
 166                 if _, err = mth.hasher.Write(buf); err != nil {
 167                         panic(err)
 168                 }
 169                 mth.hasher.Sum(mth.hashes[i][:0])
 170                 mth.hasher.Reset()
 171                 if mth.events != nil {
 172                         mth.events <- MTHEvent{MTHEventPrepend, 0, i, mth.hashes[i][:]}
 173                 }
 174                 if mth.pktName != "" {
 175                         les[len(les)-1].V = int64(read)
 176                         Progress("check", les)
 177                 }
 178                 i++
 179         }
 180         if mth.prependSize > 0 {
 181                 n, err = io.ReadFull(r, buf[:mth.prependSize])
 182                 read += n
 183                 if err != nil {
 184                         return read, err
 185                 }
 186                 if _, err = mth.hasher.Write(buf[:mth.prependSize]); err != nil {
 187                         panic(err)
 188                 }
 189                 mth.hasher.Sum(mth.hashes[i][:0])
 190                 mth.hasher.Reset()
 191                 if mth.events != nil {
 192                         mth.events <- MTHEvent{MTHEventPrepend, 0, i, mth.hashes[i][:]}
 193                 }
 194                 if mth.pktName != "" {
 195                         les[len(les)-1].V = fullsize
 196                         Progress("check", les)
 197                 }
 198         }
 199         return read, nil
 200 }
 201
 202 func (mth *MTHFat) Sum(b []byte) []byte {
 203         if mth.finished {
 204                 return append(b, mth.hashes[0][:]...)
 205         }
 206         if mth.buf.Len() > 0 {
 207                 b := mth.buf.Next(MTHBlockSize)
 208                 if _, err := mth.hasher.Write(b); err != nil {
 209                         panic(err)
 210                 }
 211                 h := new([MTHSize]byte)
 212                 mth.hasher.Sum(h[:0])
 213                 mth.hasher.Reset()
 214                 mth.hashes = append(mth.hashes, *h)
 215                 if mth.events != nil {
 216                         mth.events <- MTHEvent{
 217                                 MTHEventAppend,
 218                                 0, len(mth.hashes) - 1,
 219                                 mth.hashes[len(mth.hashes)-1][:],
 220                         }
 221                 }
 222         }
 223         switch len(mth.hashes) {
 224         case 0:
 225                 h := new([MTHSize]byte)
 226                 if _, err := mth.hasher.Write(nil); err != nil {
 227                         panic(err)
 228                 }
 229                 mth.hasher.Sum(h[:0])
 230                 mth.hasher.Reset()
 231                 mth.hashes = append(mth.hashes, *h)
 232                 if mth.events != nil {
 233                         mth.events <- MTHEvent{MTHEventAppend, 0, 0, mth.hashes[0][:]}
 234                 }
 235                 fallthrough
 236         case 1:
 237                 mth.hashes = append(mth.hashes, mth.hashes[0])
 238                 if mth.events != nil {
 239                         mth.events <- MTHEvent{MTHEventAppend, 0, 1, mth.hashes[1][:]}
 240                 }
 241         }
 242         mth.hasher = blake3.New(MTHSize, MTHNodeKey[:])
 243         level := 1
 244         for len(mth.hashes) != 1 {
 245                 hashesUp := make([][MTHSize]byte, 0, 1+len(mth.hashes)/2)
 246                 pairs := (len(mth.hashes) / 2) * 2
 247                 for i := 0; i < pairs; i += 2 {
 248                         if _, err := mth.hasher.Write(mth.hashes[i][:]); err != nil {
 249                                 panic(err)
 250                         }
 251                         if _, err := mth.hasher.Write(mth.hashes[i+1][:]); err != nil {
 252                                 panic(err)
 253                         }
 254                         h := new([MTHSize]byte)
 255                         mth.hasher.Sum(h[:0])
 256                         mth.hasher.Reset()
 257                         hashesUp = append(hashesUp, *h)
 258                         if mth.events != nil {
 259                                 mth.events <- MTHEvent{
 260                                         MTHEventFold,
 261                                         level, len(hashesUp) - 1,
 262                                         hashesUp[len(hashesUp)-1][:],
 263                                 }
 264                         }
 265                 }
 266                 if len(mth.hashes)%2 == 1 {
 267                         hashesUp = append(hashesUp, mth.hashes[len(mth.hashes)-1])
 268                         if mth.events != nil {
 269                                 mth.events <- MTHEvent{
 270                                         MTHEventAppend,
 271                                         level, len(hashesUp) - 1,
 272                                         hashesUp[len(hashesUp)-1][:],
 273                                 }
 274                         }
 275                 }
 276                 mth.hashes = hashesUp
 277                 level++
 278         }
 279         mth.finished = true
 280         if mth.events != nil {
 281                 close(mth.events)
 282         }
 283         return append(b, mth.hashes[0][:]...)
 284 }
 285
 286 type MTHSeqEnt struct {
 287         l int
 288         h [MTHSize]byte
 289 }
 290
 291 type MTHSeq struct {
 292         hasherLeaf *blake3.Hasher
 293         hasherNode *blake3.Hasher
 294         hashes     []MTHSeqEnt
 295         buf        *bytes.Buffer
 296         events     chan MTHEvent
 297         ctrs       []int
 298         finished   bool
 299 }
 300
 301 func MTHSeqNew() *MTHSeq {
 302         mth := MTHSeq{
 303                 hasherLeaf: blake3.New(MTHSize, MTHLeafKey[:]),
 304                 hasherNode: blake3.New(MTHSize, MTHNodeKey[:]),
 305                 buf:        bytes.NewBuffer(make([]byte, 0, 2*MTHBlockSize)),
 306                 ctrs:       make([]int, 1, 2),
 307         }
 308         return &mth
 309 }
 310
 311 func (mth *MTHSeq) Reset() { panic("not implemented") }
 312
 313 func (mth *MTHSeq) Size() int { return MTHSize }
 314
 315 func (mth *MTHSeq) BlockSize() int { return MTHBlockSize }
 316
 317 func (mth *MTHSeq) PrependFrom(r io.Reader) (int, error) {
 318         panic("must not reach that code")
 319 }
 320
 321 func (mth *MTHSeq) Events() chan MTHEvent {
 322         mth.events = make(chan MTHEvent)
 323         return mth.events
 324 }
 325
 326 func (mth *MTHSeq) SetPktName(pktName string) {}
 327
 328 func (mth *MTHSeq) PrependSize() int64 { return 0 }
 329
 330 func (mth *MTHSeq) leafAdd() {
 331         ent := MTHSeqEnt{l: 0}
 332         mth.hasherLeaf.Sum(ent.h[:0])
 333         mth.hasherLeaf.Reset()
 334         mth.hashes = append(mth.hashes, ent)
 335         if mth.events != nil {
 336                 mth.events <- MTHEvent{
 337                         MTHEventAppend, 0, mth.ctrs[0],
 338                         mth.hashes[len(mth.hashes)-1].h[:],
 339                 }
 340         }
 341         mth.ctrs[0]++
 342 }
 343
 344 func (mth *MTHSeq) incr(l int) {
 345         if len(mth.ctrs) <= l {
 346                 mth.ctrs = append(mth.ctrs, 0)
 347         } else {
 348                 mth.ctrs[l]++
 349         }
 350 }
 351
 352 func (mth *MTHSeq) fold() {
 353         for len(mth.hashes) >= 2 {
 354                 if mth.hashes[len(mth.hashes)-2].l != mth.hashes[len(mth.hashes)-1].l {
 355                         break
 356                 }
 357                 if _, err := mth.hasherNode.Write(mth.hashes[len(mth.hashes)-2].h[:]); err != nil {
 358                         panic(err)
 359                 }
 360                 if _, err := mth.hasherNode.Write(mth.hashes[len(mth.hashes)-1].h[:]); err != nil {
 361                         panic(err)
 362                 }
 363                 mth.hashes = mth.hashes[:len(mth.hashes)-1]
 364                 end := &mth.hashes[len(mth.hashes)-1]
 365                 end.l++
 366                 mth.incr(end.l)
 367                 mth.hasherNode.Sum(end.h[:0])
 368                 mth.hasherNode.Reset()
 369                 if mth.events != nil {
 370                         mth.events <- MTHEvent{MTHEventFold, end.l, mth.ctrs[end.l], end.h[:]}
 371                 }
 372         }
 373 }
 374
 375 func (mth *MTHSeq) Write(data []byte) (int, error) {
 376         if mth.finished {
 377                 return 0, errors.New("already Sum()ed")
 378         }
 379         n, err := mth.buf.Write(data)
 380         if err != nil {
 381                 return n, err
 382         }
 383         for mth.buf.Len() >= MTHBlockSize {
 384                 if _, err = mth.hasherLeaf.Write(mth.buf.Next(MTHBlockSize)); err != nil {
 385                         return n, err
 386                 }
 387                 mth.leafAdd()
 388                 mth.fold()
 389         }
 390         return n, err
 391 }
 392
 393 func (mth *MTHSeq) Sum(b []byte) []byte {
 394         if mth.finished {
 395                 return append(b, mth.hashes[0].h[:]...)
 396         }
 397         if mth.buf.Len() > 0 {
 398                 if _, err := mth.hasherLeaf.Write(mth.buf.Next(MTHBlockSize)); err != nil {
 399                         panic(err)
 400                 }
 401                 mth.leafAdd()
 402                 mth.fold()
 403         }
 404         switch mth.ctrs[0] {
 405         case 0:
 406                 if _, err := mth.hasherLeaf.Write(nil); err != nil {
 407                         panic(err)
 408                 }
 409                 mth.leafAdd()
 410                 fallthrough
 411         case 1:
 412                 mth.hashes = append(mth.hashes, mth.hashes[0])
 413                 mth.ctrs[0]++
 414                 if mth.events != nil {
 415                         mth.events <- MTHEvent{
 416                                 MTHEventAppend, 0, mth.ctrs[0],
 417                                 mth.hashes[len(mth.hashes)-1].h[:],
 418                         }
 419                 }
 420                 mth.fold()
 421         }
 422         for len(mth.hashes) >= 2 {
 423                 l := mth.hashes[len(mth.hashes)-2].l
 424                 mth.incr(l)
 425                 mth.hashes[len(mth.hashes)-1].l = l
 426                 if mth.events != nil {
 427                         mth.events <- MTHEvent{
 428                                 MTHEventAppend, l, mth.ctrs[l],
 429                                 mth.hashes[len(mth.hashes)-1].h[:],
 430                         }
 431                 }
 432                 mth.fold()
 433         }
 434         mth.finished = true
 435         if mth.events != nil {
 436                 close(mth.events)
 437         }
 438         return append(b, mth.hashes[0].h[:]...)
 439 }
 440
 441 func MTHNew(size, offset int64) MTH {
 442         if offset == 0 {
 443                 return MTHSeqNew()
 444         }
 445         return MTHFatNew(size, offset)
 446 }