github.com/NVIDIA/aistore@v1.3.23-0.20240517131212-7df6609be51d/cmn/jsp/io.go (about)

     1  // Package jsp (JSON persistence) provides utilities to store and load arbitrary
     2  // JSON-encoded structures with optional checksumming and compression.
     3  /*
     4   * Copyright (c) 2018-2024, NVIDIA CORPORATION. All rights reserved.
     5   */
     6  package jsp
     7  
     8  import (
     9  	"encoding/binary"
    10  	"encoding/hex"
    11  	"hash"
    12  	"io"
    13  
    14  	"github.com/NVIDIA/aistore/cmn/cos"
    15  	"github.com/NVIDIA/aistore/cmn/debug"
    16  	"github.com/NVIDIA/aistore/cmn/nlog"
    17  	"github.com/OneOfOne/xxhash"
    18  	jsoniter "github.com/json-iterator/go"
    19  	"github.com/pierrec/lz4/v3"
    20  )
    21  
    22  const (
    23  	sizeXXHash64  = cos.SizeofI64
    24  	lz4BufferSize = 64 << 10
    25  )
    26  
    27  func Encode(ws cos.WriterAt, v any, opts Options) (err error) {
    28  	var (
    29  		h       hash.Hash
    30  		w       io.Writer = ws
    31  		encoder *jsoniter.Encoder
    32  		off     int
    33  	)
    34  	if opts.Signature {
    35  		var (
    36  			prefix [prefLen]byte
    37  			flags  uint32
    38  		)
    39  		copy(prefix[:], signature) // [ 0 - 63 ]
    40  		l := len(signature)
    41  		debug.Assert(l < cos.SizeofI64)
    42  		prefix[l] = Metaver // current jsp version
    43  		off += cos.SizeofI64
    44  
    45  		binary.BigEndian.PutUint32(prefix[off:], opts.Metaver) // [ 64 - 95 ]
    46  		off += cos.SizeofI32
    47  
    48  		if opts.Compress { // [ 96 - 127 ]
    49  			flags |= 1 << 0
    50  		}
    51  		if opts.Checksum {
    52  			flags |= 1 << 1
    53  		}
    54  		binary.BigEndian.PutUint32(prefix[off:], flags)
    55  		off += cos.SizeofI32
    56  
    57  		w.Write(prefix[:])
    58  		debug.Assert(off == prefLen)
    59  	}
    60  	if opts.Checksum {
    61  		var cksum [sizeXXHash64]byte
    62  		w.Write(cksum[:]) // reserve for checksum
    63  	}
    64  	if opts.Compress {
    65  		zw := lz4.NewWriter(w)
    66  		zw.BlockMaxSize = lz4BufferSize
    67  		w = zw
    68  		defer zw.Close()
    69  	}
    70  	if opts.Checksum {
    71  		h = xxhash.New64()
    72  		cos.Assert(h.Size() == sizeXXHash64)
    73  		w = io.MultiWriter(h, w)
    74  	}
    75  
    76  	encoder = cos.JSON.NewEncoder(w)
    77  	if opts.Indent {
    78  		encoder.SetIndent("", "  ")
    79  	}
    80  	if err = encoder.Encode(v); err != nil {
    81  		return
    82  	}
    83  	if opts.Checksum {
    84  		if _, err := ws.WriteAt(h.Sum(nil), int64(off)); err != nil {
    85  			return err
    86  		}
    87  	}
    88  	return
    89  }
    90  
    91  func Decode(reader io.ReadCloser, v any, opts Options, tag string) (checksum *cos.Cksum, err error) {
    92  	var (
    93  		r             io.Reader = reader
    94  		expectedCksum uint64
    95  		h             hash.Hash
    96  		jspVer        byte
    97  	)
    98  	defer cos.Close(reader)
    99  	if opts.Signature {
   100  		var (
   101  			prefix  [prefLen]byte
   102  			metaVer uint32
   103  		)
   104  		if _, err = r.Read(prefix[:]); err != nil {
   105  			return
   106  		}
   107  		l := len(signature)
   108  		debug.Assert(l < cos.SizeofI64)
   109  		if signature != string(prefix[:l]) {
   110  			err = &ErrBadSignature{tag, string(prefix[:l]), signature}
   111  			return
   112  		}
   113  		jspVer = prefix[l]
   114  		if jspVer != Metaver {
   115  			err = newErrVersion("jsp", uint32(jspVer), Metaver)
   116  			return
   117  		}
   118  		metaVer = binary.BigEndian.Uint32(prefix[cos.SizeofI64:])
   119  		if metaVer != opts.Metaver {
   120  			if opts.OldMetaverOk == 0 || metaVer > opts.Metaver || metaVer < opts.OldMetaverOk {
   121  				// _not_ backward compatible
   122  				err = newErrVersion(tag, metaVer, opts.Metaver)
   123  				return
   124  			}
   125  			// backward compatible
   126  			erw := newErrVersion(tag, metaVer, opts.Metaver, opts.OldMetaverOk)
   127  			nlog.Warningln(erw)
   128  		}
   129  		flags := binary.BigEndian.Uint32(prefix[cos.SizeofI64+cos.SizeofI32:])
   130  		opts.Compress = flags&(1<<0) != 0
   131  		opts.Checksum = flags&(1<<1) != 0
   132  	}
   133  	if opts.Checksum {
   134  		var cksum [sizeXXHash64]byte
   135  		if _, err = r.Read(cksum[:]); err != nil {
   136  			return
   137  		}
   138  		expectedCksum = binary.BigEndian.Uint64(cksum[:])
   139  	}
   140  	if opts.Compress {
   141  		zr := lz4.NewReader(r)
   142  		zr.BlockMaxSize = lz4BufferSize
   143  		r = zr
   144  	}
   145  	if opts.Checksum {
   146  		h = xxhash.New64()
   147  		r = io.TeeReader(r, h)
   148  	}
   149  	if err = cos.JSON.NewDecoder(r).Decode(v); err != nil {
   150  		return
   151  	}
   152  	if opts.Checksum {
   153  		// We have already parsed `v` but there is still the possibility that `\n` remains
   154  		// not read. Therefore, we read it to include it into the final checksum.
   155  		var b []byte
   156  		if b, err = io.ReadAll(r); err != nil {
   157  			return
   158  		}
   159  		// To be sure that this is exactly the case...
   160  		debug.Assert(len(b) == 0 || (len(b) == 1 && b[0] == '\n'), b)
   161  
   162  		actual := h.Sum(nil)
   163  		actualCksum := binary.BigEndian.Uint64(actual)
   164  		if expectedCksum != actualCksum {
   165  			err = cos.NewErrMetaCksum(expectedCksum, actualCksum, tag)
   166  			return
   167  		}
   168  		checksum = cos.NewCksum(cos.ChecksumXXHash, hex.EncodeToString(actual))
   169  	}
   170  	return
   171  }