MCorange/wings
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Streaming Transfers (#153)

Browse Source
This commit is contained in:
Matthew Penner
2022-11-14 18:25:01 -07:00
committed by GitHub
parent 4781eeaedc
commit 57e7eb714c
21 changed files with 1015 additions and 612 deletions
Download Patch File Download Diff File Expand all files Collapse all files

118
server/filesystem/archive.go
View File

@@ -2,13 +2,13 @@ package filesystem
import (
"archive/tar"
"context"
"io"
"io/fs"
"os"
"path/filepath"
"strings"
"sync"
"sync/atomic"
"emperror.dev/errors"
"github.com/apex/log"
@@ -18,7 +18,7 @@ import (
ignore "github.com/sabhiram/go-gitignore"
"github.com/pterodactyl/wings/config"
"github.com/pterodactyl/wings/system"
"github.com/pterodactyl/wings/internal/progress"
)
const memory = 4 * 1024
@@ -33,13 +33,13 @@ var pool = sync.Pool{
// TarProgress .
type TarProgress struct {
*tar.Writer
p *Progress
p *progress.Progress
}
// NewTarProgress .
func NewTarProgress(w *tar.Writer, p *Progress) *TarProgress {
func NewTarProgress(w *tar.Writer, p *progress.Progress) *TarProgress {
if p != nil {
p.w = w
p.Writer = w
}
return &TarProgress{
Writer: w,
@@ -47,6 +47,7 @@ func NewTarProgress(w *tar.Writer, p *Progress) *TarProgress {
}
}
// Write .
func (p *TarProgress) Write(v []byte) (int, error) {
if p.p == nil {
return p.Writer.Write(v)
@@ -54,84 +55,6 @@ func (p *TarProgress) Write(v []byte) (int, error) {
return p.p.Write(v)
}
// Progress is used to track the progress of any I/O operation that are being
// performed.
type Progress struct {
// written is the total size of the files that have been written to the writer.
written int64
// Total is the total size of the archive in bytes.
total int64
// w .
w io.Writer
}
// NewProgress .
func NewProgress(total int64) *Progress {
return &Progress{total: total}
}
// SetWriter sets the writer progress will forward writes to.
// NOTE: This function is not thread safe.
func (p *Progress) SetWriter(w io.Writer) {
p.w = w
}
// Written returns the total number of bytes written.
// This function should be used when the progress is tracking data being written.
func (p *Progress) Written() int64 {
return atomic.LoadInt64(&p.written)
}
// Total returns the total size in bytes.
func (p *Progress) Total() int64 {
return atomic.LoadInt64(&p.total)
}
// Write totals the number of bytes that have been written to the writer.
func (p *Progress) Write(v []byte) (int, error) {
n := len(v)
atomic.AddInt64(&p.written, int64(n))
if p.w != nil {
return p.w.Write(v)
}
return n, nil
}
// Progress returns a formatted progress string for the current progress.
func (p *Progress) Progress(width int) string {
// current = 100 (Progress, dynamic)
// total = 1000 (Content-Length, dynamic)
// width = 25 (Number of ticks to display, static)
// widthPercentage = 100 / width (What percentage does each tick represent, static)
//
// percentageDecimal = current / total = 0.1
// percentage = percentageDecimal * 100 = 10%
// ticks = percentage / widthPercentage = 2.5
//
// ticks is a float64, so we cast it to an int which rounds it down to 2.
// Values are cast to floats to prevent integer division.
current := p.Written()
total := p.Total()
// width := is passed as a parameter
widthPercentage := float64(100) / float64(width)
percentageDecimal := float64(current) / float64(total)
percentage := percentageDecimal * 100
ticks := int(percentage / widthPercentage)
// Ensure that we never get a negative number of ticks, this will prevent strings#Repeat
// from panicking. A negative number of ticks is likely to happen when the total size is
// inaccurate, such as when we are going off of rough disk usage calculation.
if ticks < 0 {
ticks = 0
} else if ticks > width {
ticks = width
}
bar := strings.Repeat("=", ticks) + strings.Repeat(" ", width-ticks)
return "[" + bar + "] " + system.FormatBytes(current) + " / " + system.FormatBytes(total)
}
type Archive struct {
// BasePath is the absolute path to create the archive from where Files and Ignore are
// relative to.
@@ -146,12 +69,12 @@ type Archive struct {
Files []string
// Progress wraps the writer of the archive to pass through the progress tracker.
Progress *Progress
Progress *progress.Progress
}
// Create creates an archive at dst with all the files defined in the
// included Files array.
func (a *Archive) Create(dst string) error {
func (a *Archive) Create(ctx context.Context, dst string) error {
f, err := os.OpenFile(dst, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0o600)
if err != nil {
return err
@@ -169,6 +92,11 @@ func (a *Archive) Create(dst string) error {
writer = f
}
return a.Stream(ctx, writer)
}
// Stream .
func (a *Archive) Stream(ctx context.Context, w io.Writer) error {
// Choose which compression level to use based on the compression_level configuration option
var compressionLevel int
switch config.Get().System.Backups.CompressionLevel {
@@ -183,7 +111,7 @@ func (a *Archive) Create(dst string) error {
}
// Create a new gzip writer around the file.
gw, _ := pgzip.NewWriterLevel(writer, compressionLevel)
gw, _ := pgzip.NewWriterLevel(w, compressionLevel)
_ = gw.SetConcurrency(1<<20, 1)
defer gw.Close()
@@ -197,16 +125,16 @@ func (a *Archive) Create(dst string) error {
options := &godirwalk.Options{
FollowSymbolicLinks: false,
Unsorted: true,
Callback: a.callback(pw),
}
// If we're specifically looking for only certain files, or have requested
// that certain files be ignored we'll update the callback function to reflect
// that request.
var callback godirwalk.WalkFunc
if len(a.Files) == 0 && len(a.Ignore) > 0 {
i := ignore.CompileIgnoreLines(strings.Split(a.Ignore, "\n")...)
options.Callback = a.callback(pw, func(_ string, rp string) error {
callback = a.callback(pw, func(_ string, rp string) error {
if i.MatchesPath(rp) {
return godirwalk.SkipThis
}
@@ -214,7 +142,19 @@ func (a *Archive) Create(dst string) error {
return nil
})
} else if len(a.Files) > 0 {
options.Callback = a.withFilesCallback(pw)
callback = a.withFilesCallback(pw)
} else {
callback = a.callback(pw)
}
// Set the callback function, wrapped with support for context cancellation.
options.Callback = func(path string, de *godirwalk.Dirent) error {
select {
case <-ctx.Done():
return ctx.Err()
default:
return callback(path, de)
}
}
// Recursively walk the path we are archiving.

48
server/filesystem/archive_test.go
View File

@@ -1,48 +0,0 @@
package filesystem
import (
"bytes"
"testing"
. "github.com/franela/goblin"
)
func TestProgress(t *testing.T) {
g := Goblin(t)
g.Describe("Progress", func() {
g.It("properly initializes", func() {
total := int64(1000)
p := NewProgress(total)
g.Assert(p).IsNotNil()
g.Assert(p.Total()).Equal(total)
g.Assert(p.Written()).Equal(int64(0))
})
g.It("increments written when Write is called", func() {
v := []byte("hello")
p := NewProgress(1000)
_, err := p.Write(v)
g.Assert(err).IsNil()
g.Assert(p.Written()).Equal(int64(len(v)))
})
g.It("renders a progress bar", func() {
v := bytes.Repeat([]byte{' '}, 100)
p := NewProgress(1000)
_, err := p.Write(v)
g.Assert(err).IsNil()
g.Assert(p.Written()).Equal(int64(len(v)))
g.Assert(p.Progress(25)).Equal("[== ] 100 B / 1000 B")
})
g.It("renders a progress bar when written exceeds total", func() {
v := bytes.Repeat([]byte{' '}, 1001)
p := NewProgress(1000)
_, err := p.Write(v)
g.Assert(err).IsNil()
g.Assert(p.Written()).Equal(int64(len(v)))
g.Assert(p.Progress(25)).Equal("[=========================] 1001 B / 1000 B")
})
})
}

57
server/filesystem/compress.go
View File

@@ -6,6 +6,7 @@ import (
"compress/gzip"
"context"
"fmt"
"io"
iofs "io/fs"
"os"
"path"
@@ -21,7 +22,7 @@ import (
"github.com/mholt/archiver/v4"
)
// CompressFiles compresses all of the files matching the given paths in the
// 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
@@ -36,7 +37,7 @@ func (fs *Filesystem) CompressFiles(dir string, paths []string) (os.FileInfo, er
return nil, err
}
// Take all of the paths passed in and merge them together with the root directory we've gotten.
// Take all the paths passed in and merge them together with the root directory we've gotten.
for i, p := range paths {
paths[i] = filepath.Join(cleanedRootDir, p)
}
@@ -52,7 +53,7 @@ func (fs *Filesystem) CompressFiles(dir string, paths []string) (os.FileInfo, er
fmt.Sprintf("archive-%s.tar.gz", strings.ReplaceAll(time.Now().Format(time.RFC3339), ":", "")),
)
if err := a.Create(d); err != nil {
if err := a.Create(context.Background(), d); err != nil {
return nil, err
}
@@ -147,6 +148,7 @@ func (fs *Filesystem) DecompressFileUnsafe(ctx context.Context, dir string, file
if err != nil {
return err
}
// TODO: defer file close?
// Identify the type of archive we are dealing with.
format, input, err := archiver.Identify(filepath.Base(file), f)
@@ -157,13 +159,49 @@ func (fs *Filesystem) DecompressFileUnsafe(ctx context.Context, dir string, file
return err
}
return fs.extractStream(ctx, extractStreamOptions{
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 {
// Decompress and extract archive
if ex, ok := format.(archiver.Extractor); ok {
return ex.Extract(ctx, input, nil, func(ctx context.Context, f archiver.File) error {
if ex, ok := opts.Format.(archiver.Extractor); ok {
return ex.Extract(ctx, opts.Reader, nil, func(ctx context.Context, f archiver.File) error {
if f.IsDir() {
return nil
}
p := filepath.Join(dir, ExtractNameFromArchive(f))
p := filepath.Join(opts.Directory, ExtractNameFromArchive(f))
// 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
@@ -174,20 +212,19 @@ func (fs *Filesystem) DecompressFileUnsafe(ctx context.Context, dir string, file
}
defer r.Close()
if err := fs.Writefile(p, r); err != nil {
return wrapError(err, file)
return wrapError(err, opts.FileName)
}
// Update the file permissions to the one set in the archive.
if err := fs.Chmod(p, f.Mode()); err != nil {
return wrapError(err, file)
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, file)
return wrapError(err, opts.FileName)
}
return nil
})
}
return nil
}