cubefs

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
	"runtime"
10
)
11

12
// DOption is an option for creating a decoder.
13
type DOption func(*decoderOptions) error
14

15
// options retains accumulated state of multiple options.
16
type decoderOptions struct {
17
	lowMem         bool
18
	concurrent     int
19
	maxDecodedSize uint64
20
	maxWindowSize  uint64
21
	dicts          []dict
22
}
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
	}
31
	if o.concurrent > 4 {
32
		o.concurrent = 4
33
	}
34
	o.maxDecodedSize = 1 << 63
35
}
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

43
// WithDecoderConcurrency sets the number of created decoders.
44
// When decoding block with DecodeAll, this will limit the number
45
// of possible concurrently running decodes.
46
// When decoding streams, this will limit the number of
47
// inflight blocks.
48
// When decoding streams and setting maximum to 1,
49
// no async decoding will be done.
50
// When a value of 0 is provided GOMAXPROCS will be used.
51
// By default this will be set to 4 or GOMAXPROCS, whatever is lower.
52
func WithDecoderConcurrency(n int) DOption {
53
	return func(o *decoderOptions) error {
54
		if n < 0 {
55
			return errors.New("concurrency must be at least 1")
56
		}
57
		if n == 0 {
58
			o.concurrent = runtime.GOMAXPROCS(0)
59
		} else {
60
			o.concurrent = n
61
		}
62
		return nil
63
	}
64
}
65

66
// WithDecoderMaxMemory allows to set a maximum decoded size for in-memory
67
// non-streaming operations or maximum window size for streaming operations.
68
// This can be used to control memory usage of potentially hostile content.
69
// Maximum and default is 1 << 63 bytes.
70
func WithDecoderMaxMemory(n uint64) DOption {
71
	return func(o *decoderOptions) error {
72
		if n == 0 {
73
			return errors.New("WithDecoderMaxMemory must be at least 1")
74
		}
75
		if n > 1<<63 {
76
			return errors.New("WithDecoderMaxmemory must be less than 1 << 63")
77
		}
78
		o.maxDecodedSize = n
79
		return nil
80
	}
81
}
82

83
// WithDecoderDicts allows to register one or more dictionaries for the decoder.
84
// If several dictionaries with the same ID is provided the last one will be used.
85
func WithDecoderDicts(dicts ...[]byte) DOption {
86
	return func(o *decoderOptions) error {
87
		for _, b := range dicts {
88
			d, err := loadDict(b)
89
			if err != nil {
90
				return err
91
			}
92
			o.dicts = append(o.dicts, *d)
93
		}
94
		return nil
95
	}
96
}
97

98
// WithDecoderMaxWindow allows to set a maximum window size for decodes.
99
// This allows rejecting packets that will cause big memory usage.
100
// The Decoder will likely allocate more memory based on the WithDecoderLowmem setting.
101
// If WithDecoderMaxMemory is set to a lower value, that will be used.
102
// Default is 512MB, Maximum is ~3.75 TB as per zstandard spec.
103
func WithDecoderMaxWindow(size uint64) DOption {
104
	return func(o *decoderOptions) error {
105
		if size < MinWindowSize {
106
			return errors.New("WithMaxWindowSize must be at least 1KB, 1024 bytes")
107
		}
108
		if size > (1<<41)+7*(1<<38) {
109
			return errors.New("WithMaxWindowSize must be less than (1<<41) + 7*(1<<38) ~ 3.75TB")
110
		}
111
		o.maxWindowSize = size
112
		return nil
113
	}
114
}
115