-
-
Notifications
You must be signed in to change notification settings - Fork 288
/
compress.go
285 lines (252 loc) · 7.55 KB
/
compress.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
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
package filesystem
import (
"context"
"fmt"
"io"
iofs "io/fs"
"path"
"path/filepath"
"strings"
"sync/atomic"
"time"
"emperror.dev/errors"
"github.com/klauspost/compress/zip"
"github.com/mholt/archiver/v4"
"github.com/pterodactyl/wings/internal/ufs"
"github.com/pterodactyl/wings/server/filesystem/archiverext"
)
// CompressFiles compresses all the files matching the given paths in the
// specified directory. This function also supports passing nested paths to only
// compress certain files and folders when working in a larger directory. This
// effectively creates a local backup, but rather than ignoring specific files
// and folders, it takes an allow-list of files and folders.
//
// All paths are relative to the dir that is passed in as the first argument,
// and the compressed file will be placed at that location named
// `archive-{date}.tar.gz`.
func (fs *Filesystem) CompressFiles(dir string, paths []string) (ufs.FileInfo, error) {
a := &Archive{Filesystem: fs, BaseDirectory: dir, Files: paths}
d := path.Join(
dir,
fmt.Sprintf("archive-%s.tar.gz", strings.ReplaceAll(time.Now().Format(time.RFC3339), ":", "")),
)
f, err := fs.unixFS.OpenFile(d, ufs.O_WRONLY|ufs.O_CREATE, 0o644)
if err != nil {
return nil, err
}
defer f.Close()
cw := ufs.NewCountedWriter(f)
if err := a.Stream(context.Background(), cw); err != nil {
return nil, err
}
if !fs.unixFS.CanFit(cw.BytesWritten()) {
_ = fs.unixFS.Remove(d)
return nil, newFilesystemError(ErrCodeDiskSpace, nil)
}
fs.unixFS.Add(cw.BytesWritten())
return f.Stat()
}
func (fs *Filesystem) archiverFileSystem(ctx context.Context, p string) (iofs.FS, error) {
f, err := fs.unixFS.Open(p)
if err != nil {
return nil, err
}
// Do not use defer to close `f`, it will likely be used later.
format, _, err := archiver.Identify(filepath.Base(p), f)
if err != nil && !errors.Is(err, archiver.ErrNoMatch) {
_ = f.Close()
return nil, err
}
// Reset the file reader.
if _, err := f.Seek(0, io.SeekStart); err != nil {
_ = f.Close()
return nil, err
}
info, err := f.Stat()
if err != nil {
_ = f.Close()
return nil, err
}
if format != nil {
switch ff := format.(type) {
case archiver.Zip:
// zip.Reader is more performant than ArchiveFS, because zip.Reader caches content information
// and zip.Reader can open several content files concurrently because of io.ReaderAt requirement
// while ArchiveFS can't.
// zip.Reader doesn't suffer from issue #330 and #310 according to local test (but they should be fixed anyway)
return zip.NewReader(f, info.Size())
case archiver.Archival:
return archiver.ArchiveFS{Stream: io.NewSectionReader(f, 0, info.Size()), Format: ff, Context: ctx}, nil
case archiver.Compression:
return archiverext.FileFS{File: f, Compression: ff}, nil
}
}
_ = f.Close()
return nil, archiver.ErrNoMatch
}
// SpaceAvailableForDecompression looks through a given archive and determines
// if decompressing it would put the server over its allocated disk space limit.
func (fs *Filesystem) SpaceAvailableForDecompression(ctx context.Context, dir string, file string) error {
// Don't waste time trying to determine this if we know the server will have the space for
// it since there is no limit.
if fs.MaxDisk() <= 0 {
return nil
}
fsys, err := fs.archiverFileSystem(ctx, filepath.Join(dir, file))
if err != nil {
if errors.Is(err, archiver.ErrNoMatch) {
return newFilesystemError(ErrCodeUnknownArchive, err)
}
return err
}
var size atomic.Int64
return iofs.WalkDir(fsys, ".", func(path string, d iofs.DirEntry, err error) error {
if err != nil {
return err
}
select {
case <-ctx.Done():
// Stop walking if the context is canceled.
return ctx.Err()
default:
info, err := d.Info()
if err != nil {
return err
}
if !fs.unixFS.CanFit(size.Add(info.Size())) {
return newFilesystemError(ErrCodeDiskSpace, nil)
}
return nil
}
})
}
// DecompressFile will decompress a file in a given directory by using the
// archiver tool to infer the file type and go from there. This will walk over
// all the files within the given archive and ensure that there is not a
// zip-slip attack being attempted by validating that the final path is within
// the server data directory.
func (fs *Filesystem) DecompressFile(ctx context.Context, dir string, file string) error {
f, err := fs.unixFS.Open(filepath.Join(dir, file))
if err != nil {
return err
}
defer f.Close()
// Identify the type of archive we are dealing with.
format, input, err := archiver.Identify(filepath.Base(file), f)
if err != nil {
if errors.Is(err, archiver.ErrNoMatch) {
return newFilesystemError(ErrCodeUnknownArchive, err)
}
return err
}
return fs.extractStream(ctx, extractStreamOptions{
FileName: file,
Directory: dir,
Format: format,
Reader: input,
})
}
// ExtractStreamUnsafe .
func (fs *Filesystem) ExtractStreamUnsafe(ctx context.Context, dir string, r io.Reader) error {
format, input, err := archiver.Identify("archive.tar.gz", r)
if err != nil {
if errors.Is(err, archiver.ErrNoMatch) {
return newFilesystemError(ErrCodeUnknownArchive, err)
}
return err
}
return fs.extractStream(ctx, extractStreamOptions{
Directory: dir,
Format: format,
Reader: input,
})
}
type extractStreamOptions struct {
// The directory to extract the archive to.
Directory string
// File name of the archive.
FileName string
// Format of the archive.
Format archiver.Format
// Reader for the archive.
Reader io.Reader
}
func (fs *Filesystem) extractStream(ctx context.Context, opts extractStreamOptions) error {
// See if it's a compressed archive, such as TAR or a ZIP
ex, ok := opts.Format.(archiver.Extractor)
if !ok {
// If not, check if it's a single-file compression, such as
// .log.gz, .sql.gz, and so on
de, ok := opts.Format.(archiver.Decompressor)
if !ok {
return nil
}
// Strip the compression suffix
p := filepath.Join(opts.Directory, strings.TrimSuffix(opts.FileName, opts.Format.Name()))
// Make sure it's not ignored
if err := fs.IsIgnored(p); err != nil {
return nil
}
reader, err := de.OpenReader(opts.Reader)
if err != nil {
return err
}
defer reader.Close()
// Open the file for creation/writing
f, err := fs.unixFS.OpenFile(p, ufs.O_WRONLY|ufs.O_CREATE, 0o644)
if err != nil {
return err
}
defer f.Close()
// Read in 4 KB chunks
buf := make([]byte, 4096)
for {
n, err := reader.Read(buf)
if n > 0 {
// Check quota before writing the chunk
if quotaErr := fs.HasSpaceFor(int64(n)); quotaErr != nil {
return quotaErr
}
// Write the chunk
if _, writeErr := f.Write(buf[:n]); writeErr != nil {
return writeErr
}
// Add to quota
fs.addDisk(int64(n))
}
if err != nil {
// EOF are expected
if err == io.EOF {
break
}
// Return any other
return err
}
}
return nil
}
// Decompress and extract archive
return ex.Extract(ctx, opts.Reader, nil, func(ctx context.Context, f archiver.File) error {
if f.IsDir() {
return nil
}
p := filepath.Join(opts.Directory, f.NameInArchive)
// If it is ignored, just don't do anything with the file and skip over it.
if err := fs.IsIgnored(p); err != nil {
return nil
}
r, err := f.Open()
if err != nil {
return err
}
defer r.Close()
if err := fs.Write(p, r, f.Size(), f.Mode()); err != nil {
return wrapError(err, opts.FileName)
}
// Update the file modification time to the one set in the archive.
if err := fs.Chtimes(p, f.ModTime(), f.ModTime()); err != nil {
return wrapError(err, opts.FileName)
}
return nil
})
}