github.com/m3db/m3@v1.5.0/src/dbnode/persist/fs/msgpack/stream_with_digest.go (about)

     1  // Copyright (c) 2020 Uber Technologies, Inc.
     2  //
     3  // Permission is hereby granted, free of charge, to any person obtaining a copy
     4  // of this software and associated documentation files (the "Software"), to deal
     5  // in the Software without restriction, including without limitation the rights
     6  // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
     7  // copies of the Software, and to permit persons to whom the Software is
     8  // furnished to do so, subject to the following conditions:
     9  //
    10  // The above copyright notice and this permission notice shall be included in
    11  // all copies or substantial portions of the Software.
    12  //
    13  // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    14  // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    15  // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    16  // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    17  // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    18  // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    19  // THE SOFTWARE.
    20  
    21  package msgpack
    22  
    23  import (
    24  	"errors"
    25  	"hash"
    26  	"hash/adler32"
    27  )
    28  
    29  var (
    30  	// errChecksumMismatch returned when the calculated checksum doesn't match the stored checksum
    31  	errChecksumMismatch = errors.New("calculated checksum doesn't match stored checksum")
    32  )
    33  
    34  var _ DecoderStream = &decoderStreamWithDigest{}
    35  
    36  // decoderStreamWithDigest calculates the digest as it processes a decoder stream.
    37  type decoderStreamWithDigest struct {
    38  	reader        DecoderStream
    39  	readerDigest  hash.Hash32
    40  	unreadByte    bool
    41  	enabled       bool
    42  	singleByteBuf []byte
    43  }
    44  
    45  // newDecoderStreamWithDigest returns a new decoderStreamWithDigest
    46  func newDecoderStreamWithDigest(reader DecoderStream) *decoderStreamWithDigest {
    47  	return &decoderStreamWithDigest{
    48  		reader:        reader,
    49  		readerDigest:  adler32.New(),
    50  		singleByteBuf: make([]byte, 1),
    51  	}
    52  }
    53  
    54  func (d *decoderStreamWithDigest) Read(p []byte) (n int, err error) {
    55  	n, err = d.reader.Read(p)
    56  	if err != nil {
    57  		return n, err
    58  	}
    59  	if n <= 0 {
    60  		return n, nil
    61  	}
    62  
    63  	start := 0
    64  	if d.unreadByte {
    65  		d.unreadByte = false
    66  		start++
    67  	}
    68  	if d.enabled {
    69  		if _, err := d.readerDigest.Write(p[start:n]); err != nil {
    70  			return 0, err
    71  		}
    72  	}
    73  	return n, err
    74  }
    75  
    76  func (d *decoderStreamWithDigest) ReadByte() (byte, error) {
    77  	b, err := d.reader.ReadByte()
    78  	if err != nil {
    79  		return 0, err
    80  	}
    81  
    82  	if d.unreadByte {
    83  		d.unreadByte = false
    84  	} else if d.enabled {
    85  		d.singleByteBuf[0] = b
    86  		if _, err := d.readerDigest.Write(d.singleByteBuf); err != nil {
    87  			return b, err
    88  		}
    89  	}
    90  	return b, err
    91  }
    92  
    93  func (d *decoderStreamWithDigest) UnreadByte() error {
    94  	err := d.reader.UnreadByte()
    95  	if err == nil {
    96  		d.unreadByte = true
    97  	}
    98  	return err
    99  }
   100  
   101  // reset resets the reader for use with a new reader.
   102  func (d *decoderStreamWithDigest) reset(stream DecoderStream) {
   103  	d.reader = stream
   104  	d.readerDigest.Reset()
   105  }
   106  
   107  // digest returns the digest
   108  func (d *decoderStreamWithDigest) digest() hash.Hash32 {
   109  	return d.readerDigest
   110  }
   111  
   112  // validate compares the current digest against the expected digest and returns
   113  // an error if they don't match.
   114  func (d *decoderStreamWithDigest) validate(expectedDigest uint32) error {
   115  	if d.readerDigest.Sum32() != expectedDigest {
   116  		return errChecksumMismatch
   117  	}
   118  	return nil
   119  }
   120  
   121  // capture provides a mechanism for manually capturing bytes to add to digest when reader is manipulated
   122  // through atypical means (e.g. reading directly from the backing byte slice of a ByteReader)
   123  func (d *decoderStreamWithDigest) capture(bytes []byte) error {
   124  	// No-op if not actually capturing at the moment
   125  	if d.enabled {
   126  		if _, err := d.readerDigest.Write(bytes); err != nil {
   127  			return err
   128  		}
   129  	}
   130  	return nil
   131  }
   132  
   133  // setDigestReaderEnabled enables calculating of digest for bytes read. If this is false, behaves like a regular reader.
   134  func (d *decoderStreamWithDigest) setDigestReaderEnabled(enabled bool) {
   135  	if !d.enabled && enabled {
   136  		d.enabled = true
   137  		d.readerDigest.Reset()
   138  	} else if d.enabled && !enabled {
   139  		d.enabled = false
   140  	}
   141  }
   142  
   143  // Returns the decoder stream wrapped by this object
   144  func (d *decoderStreamWithDigest) wrappedStream() DecoderStream {
   145  	return d.reader
   146  }