decoder_options.go 5.2 KB

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