decoder_options.go raw

   1  // Copyright 2019+ Klaus Post. All rights reserved.
   2  // License information can be found in the LICENSE file.
   3  // Based on work by Yann Collet, released under BSD License.
   4  
   5  package zstd
   6  
   7  import (
   8  	"errors"
   9  	"fmt"
  10  	"math/bits"
  11  	"runtime"
  12  )
  13  
  14  // DOption is an option for creating a decoder.
  15  type DOption func(*decoderOptions) error
  16  
  17  // options retains accumulated state of multiple options.
  18  type decoderOptions struct {
  19  	lowMem          bool
  20  	concurrent      int
  21  	maxDecodedSize  uint64
  22  	maxWindowSize   uint64
  23  	dicts           []*dict
  24  	ignoreChecksum  bool
  25  	limitToCap      bool
  26  	decodeBufsBelow int
  27  }
  28  
  29  func (o *decoderOptions) setDefault() {
  30  	*o = decoderOptions{
  31  		// use less ram: true for now, but may change.
  32  		lowMem:          true,
  33  		concurrent:      runtime.GOMAXPROCS(0),
  34  		maxWindowSize:   MaxWindowSize,
  35  		decodeBufsBelow: 128 << 10,
  36  	}
  37  	if o.concurrent > 4 {
  38  		o.concurrent = 4
  39  	}
  40  	o.maxDecodedSize = 64 << 30
  41  }
  42  
  43  // WithDecoderLowmem will set whether to use a lower amount of memory,
  44  // but possibly have to allocate more while running.
  45  func WithDecoderLowmem(b bool) DOption {
  46  	return func(o *decoderOptions) error { o.lowMem = b; return nil }
  47  }
  48  
  49  // WithDecoderConcurrency sets the number of created decoders.
  50  // When decoding block with DecodeAll, this will limit the number
  51  // of possible concurrently running decodes.
  52  // When decoding streams, this will limit the number of
  53  // inflight blocks.
  54  // When decoding streams and setting maximum to 1,
  55  // no async decoding will be done.
  56  // When a value of 0 is provided GOMAXPROCS will be used.
  57  // By default this will be set to 4 or GOMAXPROCS, whatever is lower.
  58  func WithDecoderConcurrency(n int) DOption {
  59  	return func(o *decoderOptions) error {
  60  		if n < 0 {
  61  			return errors.New("concurrency must be at least 1")
  62  		}
  63  		if n == 0 {
  64  			o.concurrent = runtime.GOMAXPROCS(0)
  65  		} else {
  66  			o.concurrent = n
  67  		}
  68  		return nil
  69  	}
  70  }
  71  
  72  // WithDecoderMaxMemory allows to set a maximum decoded size for in-memory
  73  // non-streaming operations or maximum window size for streaming operations.
  74  // This can be used to control memory usage of potentially hostile content.
  75  // Maximum is 1 << 63 bytes. Default is 64GiB.
  76  func WithDecoderMaxMemory(n uint64) DOption {
  77  	return func(o *decoderOptions) error {
  78  		if n == 0 {
  79  			return errors.New("WithDecoderMaxMemory must be at least 1")
  80  		}
  81  		if n > 1<<63 {
  82  			return errors.New("WithDecoderMaxmemory must be less than 1 << 63")
  83  		}
  84  		o.maxDecodedSize = n
  85  		return nil
  86  	}
  87  }
  88  
  89  // WithDecoderDicts allows to register one or more dictionaries for the decoder.
  90  //
  91  // Each slice in dict must be in the [dictionary format] produced by
  92  // "zstd --train" from the Zstandard reference implementation.
  93  //
  94  // If several dictionaries with the same ID are provided, the last one will be used.
  95  //
  96  // [dictionary format]: https://github.com/facebook/zstd/blob/dev/doc/zstd_compression_format.md#dictionary-format
  97  func WithDecoderDicts(dicts ...[]byte) DOption {
  98  	return func(o *decoderOptions) error {
  99  		for _, b := range dicts {
 100  			d, err := loadDict(b)
 101  			if err != nil {
 102  				return err
 103  			}
 104  			o.dicts = append(o.dicts, d)
 105  		}
 106  		return nil
 107  	}
 108  }
 109  
 110  // WithDecoderDictRaw registers a dictionary that may be used by the decoder.
 111  // The slice content can be arbitrary data.
 112  func WithDecoderDictRaw(id uint32, content []byte) DOption {
 113  	return func(o *decoderOptions) error {
 114  		if bits.UintSize > 32 && uint(len(content)) > dictMaxLength {
 115  			return fmt.Errorf("dictionary of size %d > 2GiB too large", len(content))
 116  		}
 117  		o.dicts = append(o.dicts, &dict{id: id, content: content, offsets: [3]int{1, 4, 8}})
 118  		return nil
 119  	}
 120  }
 121  
 122  // WithDecoderMaxWindow allows to set a maximum window size for decodes.
 123  // This allows rejecting packets that will cause big memory usage.
 124  // The Decoder will likely allocate more memory based on the WithDecoderLowmem setting.
 125  // If WithDecoderMaxMemory is set to a lower value, that will be used.
 126  // Default is 512MB, Maximum is ~3.75 TB as per zstandard spec.
 127  func WithDecoderMaxWindow(size uint64) DOption {
 128  	return func(o *decoderOptions) error {
 129  		if size < MinWindowSize {
 130  			return errors.New("WithMaxWindowSize must be at least 1KB, 1024 bytes")
 131  		}
 132  		if size > (1<<41)+7*(1<<38) {
 133  			return errors.New("WithMaxWindowSize must be less than (1<<41) + 7*(1<<38) ~ 3.75TB")
 134  		}
 135  		o.maxWindowSize = size
 136  		return nil
 137  	}
 138  }
 139  
 140  // WithDecodeAllCapLimit will limit DecodeAll to decoding cap(dst)-len(dst) bytes,
 141  // or any size set in WithDecoderMaxMemory.
 142  // This can be used to limit decoding to a specific maximum output size.
 143  // Disabled by default.
 144  func WithDecodeAllCapLimit(b bool) DOption {
 145  	return func(o *decoderOptions) error {
 146  		o.limitToCap = b
 147  		return nil
 148  	}
 149  }
 150  
 151  // WithDecodeBuffersBelow will fully decode readers that have a
 152  // `Bytes() []byte` and `Len() int` interface similar to bytes.Buffer.
 153  // This typically uses less allocations but will have the full decompressed object in memory.
 154  // Note that DecodeAllCapLimit will disable this, as well as giving a size of 0 or less.
 155  // Default is 128KiB.
 156  func WithDecodeBuffersBelow(size int) DOption {
 157  	return func(o *decoderOptions) error {
 158  		o.decodeBufsBelow = size
 159  		return nil
 160  	}
 161  }
 162  
 163  // IgnoreChecksum allows to forcibly ignore checksum checking.
 164  func IgnoreChecksum(b bool) DOption {
 165  	return func(o *decoderOptions) error {
 166  		o.ignoreChecksum = b
 167  		return nil
 168  	}
 169  }
 170