forked from Sandertv/gophertunnel
-
Notifications
You must be signed in to change notification settings - Fork 0
/
compression.go
151 lines (130 loc) · 4.55 KB
/
compression.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
package packet
import (
"bytes"
"fmt"
"github.com/Adrian8115/gophertunnel-Amethyst-Protocol/minecraft/internal"
"github.com/golang/snappy"
"github.com/klauspost/compress/flate"
"io"
"sync"
)
// Compression represents a compression algorithm that can compress and decompress data.
type Compression interface {
// EncodeCompression encodes the compression algorithm into a uint16 ID.
EncodeCompression() uint16
// Compress compresses the given data and returns the compressed data.
Compress(decompressed []byte) ([]byte, error)
// Decompress decompresses the given data and returns the decompressed data.
Decompress(compressed []byte) ([]byte, error)
}
var (
// FlateCompression is the implementation of the Flate compression
// algorithm. This was used by default until v1.19.30.
FlateCompression flateCompression
// SnappyCompression is the implementation of the Snappy compression
// algorithm. This is used by default.
SnappyCompression snappyCompression
DefaultCompression Compression = FlateCompression
)
type (
// flateCompression is the implementation of the Flate compression algorithm. This was used by default until v1.19.30.
flateCompression struct{}
// snappyCompression is the implementation of the Snappy compression algorithm. This is used by default.
snappyCompression struct{}
)
// flateDecompressPool is a sync.Pool for io.ReadCloser flate readers. These are
// pooled for connections.
var (
flateDecompressPool = sync.Pool{
New: func() any { return flate.NewReader(bytes.NewReader(nil)) },
}
flateCompressPool = sync.Pool{
New: func() any {
w, _ := flate.NewWriter(io.Discard, 6)
return w
},
}
)
// EncodeCompression ...
func (flateCompression) EncodeCompression() uint16 {
return CompressionAlgorithmFlate
}
// Compress ...
func (flateCompression) Compress(decompressed []byte) ([]byte, error) {
compressed := internal.BufferPool.Get().(*bytes.Buffer)
w := flateCompressPool.Get().(*flate.Writer)
defer func() {
// Reset the buffer, so we can return it to the buffer pool safely.
compressed.Reset()
internal.BufferPool.Put(compressed)
flateCompressPool.Put(w)
}()
w.Reset(compressed)
_, err := w.Write(decompressed)
if err != nil {
return nil, fmt.Errorf("compress flate: %w", err)
}
err = w.Close()
if err != nil {
return nil, fmt.Errorf("close flate writer: %w", err)
}
return append([]byte(nil), compressed.Bytes()...), nil
}
// Decompress ...
func (flateCompression) Decompress(compressed []byte) ([]byte, error) {
buf := bytes.NewReader(compressed)
c := flateDecompressPool.Get().(io.ReadCloser)
defer flateDecompressPool.Put(c)
if err := c.(flate.Resetter).Reset(buf, nil); err != nil {
return nil, fmt.Errorf("reset flate: %w", err)
}
_ = c.Close()
// Guess an uncompressed size of 2*len(compressed).
decompressed := bytes.NewBuffer(make([]byte, 0, len(compressed)*2))
if _, err := io.Copy(decompressed, c); err != nil {
return nil, fmt.Errorf("decompress flate: %v", err)
}
return decompressed.Bytes(), nil
}
// EncodeCompression ...
func (snappyCompression) EncodeCompression() uint16 {
return CompressionAlgorithmSnappy
}
// Compress ...
func (snappyCompression) Compress(decompressed []byte) ([]byte, error) {
// Because Snappy allocates a slice only once, it is less important to have
// a dst slice pre-allocated. With flateCompression this is more important,
// because flate does a lot of smaller allocations which causes a
// considerable slowdown.
return snappy.Encode(nil, decompressed), nil
}
// Decompress ...
func (snappyCompression) Decompress(compressed []byte) ([]byte, error) {
// Snappy writes a decoded data length prefix, so it can allocate the
// perfect size right away and only needs to allocate once. No need to pool
// byte slices here either.
decompressed, err := snappy.Decode(nil, compressed)
if err != nil {
return nil, fmt.Errorf("decompress snappy: %w", err)
}
return decompressed, nil
}
// init registers all valid compressions with the protocol.
func init() {
RegisterCompression(flateCompression{})
RegisterCompression(snappyCompression{})
}
var compressions = map[uint16]Compression{}
// RegisterCompression registers a compression so that it can be used by the protocol.
func RegisterCompression(compression Compression) {
compressions[compression.EncodeCompression()] = compression
}
// CompressionByID attempts to return a compression by the ID it was registered with. If found, the compression found
// is returned and the bool is true.
func CompressionByID(id uint16) (Compression, bool) {
c, ok := compressions[id]
if !ok {
c = DefaultCompression
}
return c, ok
}