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 int64
  50         Ctr   int64
  51         Hsh   []byte
  52 }
  53
  54 type MTH interface {
  55         hash.Hash
  56         PrependFrom(r io.Reader) (int64, 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 := offset / MTHBlockSize
  84         skip := MTHBlockSize - (offset - prepends*MTHBlockSize)
  85         if skip == MTHBlockSize {
  86                 skip = 0
  87         } else if skip > 0 {
  88                 prepends++
  89         }
  90         prependSize := 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, int64(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) (int64, 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 n int
 157         var i, read int64
 158         fullsize := mth.prependSize
 159         les := LEs{{"Pkt", mth.pktName}, {"FullSize", fullsize}, {"Size", 0}}
 160         for mth.prependSize >= MTHBlockSize {
 161                 n, err = io.ReadFull(r, buf)
 162                 read += int64(n)
 163                 mth.prependSize -= MTHBlockSize
 164                 if err != nil {
 165                         return read, err
 166                 }
 167                 if _, err = mth.hasher.Write(buf); err != nil {
 168                         panic(err)
 169                 }
 170                 mth.hasher.Sum(mth.hashes[i][:0])
 171                 mth.hasher.Reset()
 172                 if mth.events != nil {
 173                         mth.events <- MTHEvent{MTHEventPrepend, 0, i, mth.hashes[i][:]}
 174                 }
 175                 if mth.pktName != "" {
 176                         les[len(les)-1].V = read
 177                         Progress("check", les)
 178                 }
 179                 i++
 180         }
 181         if mth.prependSize > 0 {
 182                 n, err = io.ReadFull(r, buf[:mth.prependSize])
 183                 read += int64(n)
 184                 if err != nil {
 185                         return read, err
 186                 }
 187                 if _, err = mth.hasher.Write(buf[:mth.prependSize]); err != nil {
 188                         panic(err)
 189                 }
 190                 mth.hasher.Sum(mth.hashes[i][:0])
 191                 mth.hasher.Reset()
 192                 if mth.events != nil {
 193                         mth.events <- MTHEvent{MTHEventPrepend, 0, i, mth.hashes[i][:]}
 194                 }
 195                 if mth.pktName != "" {
 196                         les[len(les)-1].V = fullsize
 197                         Progress("check", les)
 198                 }
 199         }
 200         return read, nil
 201 }
 202
 203 func (mth *MTHFat) Sum(b []byte) []byte {
 204         if mth.finished {
 205                 return append(b, mth.hashes[0][:]...)
 206         }
 207         if mth.buf.Len() > 0 {
 208                 b := mth.buf.Next(MTHBlockSize)
 209                 if _, err := mth.hasher.Write(b); err != nil {
 210                         panic(err)
 211                 }
 212                 h := new([MTHSize]byte)
 213                 mth.hasher.Sum(h[:0])
 214                 mth.hasher.Reset()
 215                 mth.hashes = append(mth.hashes, *h)
 216                 if mth.events != nil {
 217                         mth.events <- MTHEvent{
 218                                 MTHEventAppend,
 219                                 0, int64(len(mth.hashes) - 1),
 220                                 mth.hashes[len(mth.hashes)-1][:],
 221                         }
 222                 }
 223         }
 224         switch len(mth.hashes) {
 225         case 0:
 226                 h := new([MTHSize]byte)
 227                 if _, err := mth.hasher.Write(nil); err != nil {
 228                         panic(err)
 229                 }
 230                 mth.hasher.Sum(h[:0])
 231                 mth.hasher.Reset()
 232                 mth.hashes = append(mth.hashes, *h)
 233                 if mth.events != nil {
 234                         mth.events <- MTHEvent{MTHEventAppend, 0, 0, mth.hashes[0][:]}
 235                 }
 236                 fallthrough
 237         case 1:
 238                 mth.hashes = append(mth.hashes, mth.hashes[0])
 239                 if mth.events != nil {
 240                         mth.events <- MTHEvent{MTHEventAppend, 0, 1, mth.hashes[1][:]}
 241                 }
 242         }
 243         mth.hasher = blake3.New(MTHSize, MTHNodeKey[:])
 244         level := int64(1)
 245         for len(mth.hashes) != 1 {
 246                 hashesUp := make([][MTHSize]byte, 0, 1+len(mth.hashes)/2)
 247                 pairs := (len(mth.hashes) / 2) * 2
 248                 for i := 0; i < pairs; i += 2 {
 249                         if _, err := mth.hasher.Write(mth.hashes[i][:]); err != nil {
 250                                 panic(err)
 251                         }
 252                         if _, err := mth.hasher.Write(mth.hashes[i+1][:]); err != nil {
 253                                 panic(err)
 254                         }
 255                         h := new([MTHSize]byte)
 256                         mth.hasher.Sum(h[:0])
 257                         mth.hasher.Reset()
 258                         hashesUp = append(hashesUp, *h)
 259                         if mth.events != nil {
 260                                 mth.events <- MTHEvent{
 261                                         MTHEventFold,
 262                                         level, int64(len(hashesUp) - 1),
 263                                         hashesUp[len(hashesUp)-1][:],
 264                                 }
 265                         }
 266                 }
 267                 if len(mth.hashes)%2 == 1 {
 268                         hashesUp = append(hashesUp, mth.hashes[len(mth.hashes)-1])
 269                         if mth.events != nil {
 270                                 mth.events <- MTHEvent{
 271                                         MTHEventAppend,
 272                                         level, int64(len(hashesUp) - 1),
 273                                         hashesUp[len(hashesUp)-1][:],
 274                                 }
 275                         }
 276                 }
 277                 mth.hashes = hashesUp
 278                 level++
 279         }
 280         mth.finished = true
 281         if mth.events != nil {
 282                 close(mth.events)
 283         }
 284         return append(b, mth.hashes[0][:]...)
 285 }
 286
 287 type MTHSeqEnt struct {
 288         l int64
 289         h [MTHSize]byte
 290 }
 291
 292 type MTHSeq struct {
 293         hasherLeaf *blake3.Hasher
 294         hasherNode *blake3.Hasher
 295         hashes     []MTHSeqEnt
 296         buf        *bytes.Buffer
 297         events     chan MTHEvent
 298         ctrs       []int64
 299         finished   bool
 300 }
 301
 302 func MTHSeqNew() *MTHSeq {
 303         mth := MTHSeq{
 304                 hasherLeaf: blake3.New(MTHSize, MTHLeafKey[:]),
 305                 hasherNode: blake3.New(MTHSize, MTHNodeKey[:]),
 306                 buf:        bytes.NewBuffer(make([]byte, 0, 2*MTHBlockSize)),
 307                 ctrs:       make([]int64, 1, 2),
 308         }
 309         return &mth
 310 }
 311
 312 func (mth *MTHSeq) Reset() { panic("not implemented") }
 313
 314 func (mth *MTHSeq) Size() int { return MTHSize }
 315
 316 func (mth *MTHSeq) BlockSize() int { return MTHBlockSize }
 317
 318 func (mth *MTHSeq) PrependFrom(r io.Reader) (int64, error) {
 319         panic("must not reach that code")
 320 }
 321
 322 func (mth *MTHSeq) Events() chan MTHEvent {
 323         mth.events = make(chan MTHEvent)
 324         return mth.events
 325 }
 326
 327 func (mth *MTHSeq) SetPktName(pktName string) {}
 328
 329 func (mth *MTHSeq) PrependSize() int64 { return 0 }
 330
 331 func (mth *MTHSeq) leafAdd() {
 332         ent := MTHSeqEnt{l: 0}
 333         mth.hasherLeaf.Sum(ent.h[:0])
 334         mth.hasherLeaf.Reset()
 335         mth.hashes = append(mth.hashes, ent)
 336         if mth.events != nil {
 337                 mth.events <- MTHEvent{
 338                         MTHEventAppend, 0, mth.ctrs[0],
 339                         mth.hashes[len(mth.hashes)-1].h[:],
 340                 }
 341         }
 342         mth.ctrs[0]++
 343 }
 344
 345 func (mth *MTHSeq) incr(l int64) {
 346         if int64(len(mth.ctrs)) <= l {
 347                 mth.ctrs = append(mth.ctrs, 0)
 348         } else {
 349                 mth.ctrs[l]++
 350         }
 351 }
 352
 353 func (mth *MTHSeq) fold() {
 354         for len(mth.hashes) >= 2 {
 355                 if mth.hashes[len(mth.hashes)-2].l != mth.hashes[len(mth.hashes)-1].l {
 356                         break
 357                 }
 358                 if _, err := mth.hasherNode.Write(mth.hashes[len(mth.hashes)-2].h[:]); err != nil {
 359                         panic(err)
 360                 }
 361                 if _, err := mth.hasherNode.Write(mth.hashes[len(mth.hashes)-1].h[:]); err != nil {
 362                         panic(err)
 363                 }
 364                 mth.hashes = mth.hashes[:len(mth.hashes)-1]
 365                 end := &mth.hashes[len(mth.hashes)-1]
 366                 end.l++
 367                 mth.incr(end.l)
 368                 mth.hasherNode.Sum(end.h[:0])
 369                 mth.hasherNode.Reset()
 370                 if mth.events != nil {
 371                         mth.events <- MTHEvent{MTHEventFold, end.l, mth.ctrs[end.l], end.h[:]}
 372                 }
 373         }
 374 }
 375
 376 func (mth *MTHSeq) Write(data []byte) (int, error) {
 377         if mth.finished {
 378                 return 0, errors.New("already Sum()ed")
 379         }
 380         n, err := mth.buf.Write(data)
 381         if err != nil {
 382                 return n, err
 383         }
 384         for mth.buf.Len() >= MTHBlockSize {
 385                 if _, err = mth.hasherLeaf.Write(mth.buf.Next(MTHBlockSize)); err != nil {
 386                         return n, err
 387                 }
 388                 mth.leafAdd()
 389                 mth.fold()
 390         }
 391         return n, err
 392 }
 393
 394 func (mth *MTHSeq) Sum(b []byte) []byte {
 395         if mth.finished {
 396                 return append(b, mth.hashes[0].h[:]...)
 397         }
 398         if mth.buf.Len() > 0 {
 399                 if _, err := mth.hasherLeaf.Write(mth.buf.Next(MTHBlockSize)); err != nil {
 400                         panic(err)
 401                 }
 402                 mth.leafAdd()
 403                 mth.fold()
 404         }
 405         switch mth.ctrs[0] {
 406         case 0:
 407                 if _, err := mth.hasherLeaf.Write(nil); err != nil {
 408                         panic(err)
 409                 }
 410                 mth.leafAdd()
 411                 fallthrough
 412         case 1:
 413                 mth.hashes = append(mth.hashes, mth.hashes[0])
 414                 mth.ctrs[0]++
 415                 if mth.events != nil {
 416                         mth.events <- MTHEvent{
 417                                 MTHEventAppend, 0, mth.ctrs[0],
 418                                 mth.hashes[len(mth.hashes)-1].h[:],
 419                         }
 420                 }
 421                 mth.fold()
 422         }
 423         for len(mth.hashes) >= 2 {
 424                 l := mth.hashes[len(mth.hashes)-2].l
 425                 mth.incr(l)
 426                 mth.hashes[len(mth.hashes)-1].l = l
 427                 if mth.events != nil {
 428                         mth.events <- MTHEvent{
 429                                 MTHEventAppend, l, mth.ctrs[l],
 430                                 mth.hashes[len(mth.hashes)-1].h[:],
 431                         }
 432                 }
 433                 mth.fold()
 434         }
 435         mth.finished = true
 436         if mth.events != nil {
 437                 close(mth.events)
 438         }
 439         return append(b, mth.hashes[0].h[:]...)
 440 }
 441
 442 func MTHNew(size, offset int64) MTH {
 443         if offset == 0 {
 444                 return MTHSeqNew()
 445         }
 446         return MTHFatNew(size, offset)
 447 }