github.com/cloud-foundations/dominator@v0.0.0-20221004181915-6e4fee580046/lib/fsutil/checksummer.go (about)

     1  package fsutil
     2  
     3  import (
     4  	"bufio"
     5  	"bytes"
     6  	"crypto/sha512"
     7  	"fmt"
     8  	"io"
     9  )
    10  
    11  func newChecksumReader(reader io.Reader) *ChecksumReader {
    12  	r := new(ChecksumReader)
    13  	r.checksummer = sha512.New()
    14  	if _, ok := reader.(io.ByteReader); !ok {
    15  		r.reader = bufio.NewReader(reader)
    16  	} else {
    17  		r.reader = reader
    18  	}
    19  	return r
    20  }
    21  
    22  func (r *ChecksumReader) getChecksum() []byte {
    23  	return r.checksummer.Sum(nil)
    24  }
    25  
    26  func (r *ChecksumReader) read(p []byte) (int, error) {
    27  	if nRead, err := r.reader.Read(p); err != nil && err != io.EOF {
    28  		return nRead, err
    29  	} else if _, e := r.checksummer.Write(p[:nRead]); e != nil {
    30  		return nRead, e
    31  	} else {
    32  		return nRead, err
    33  	}
    34  }
    35  
    36  func (r *ChecksumReader) readByte() (byte, error) {
    37  	buf := make([]byte, 1)
    38  	_, err := r.read(buf)
    39  	return buf[0], err
    40  }
    41  
    42  func (r *ChecksumReader) verifyChecksum() error {
    43  	buf := make([]byte, r.checksummer.Size())
    44  	nRead, err := io.ReadAtLeast(r.reader, buf, len(buf))
    45  	if err != nil {
    46  		return err
    47  	}
    48  	if nRead != r.checksummer.Size() {
    49  		return fmt.Errorf(
    50  			"ChecksumReader.Checksum(): expected: %d got: %d bytes",
    51  			r.checksummer.Size(), nRead)
    52  	}
    53  	if !bytes.Equal(buf, r.getChecksum()) {
    54  		return ErrorChecksumMismatch
    55  	}
    56  	return nil
    57  }
    58  
    59  func newChecksumWriter(writer io.Writer) *ChecksumWriter {
    60  	w := new(ChecksumWriter)
    61  	w.checksummer = sha512.New()
    62  	w.writer = writer
    63  	return w
    64  }
    65  
    66  func (w *ChecksumWriter) write(p []byte) (int, error) {
    67  	if _, err := w.checksummer.Write(p); err != nil {
    68  		return 0, err
    69  	}
    70  	return w.writer.Write(p)
    71  }
    72  
    73  func (w *ChecksumWriter) writeChecksum() error {
    74  	_, err := w.writer.Write(w.checksummer.Sum(nil))
    75  	return err
    76  }