server(filesystem): fix inaccurate archive progress (#145)

server/backup.go

@@ -24,7 +24,6 @@ func (s *Server) notifyPanelOfBackup(uuid string, ad *backup.ArchiveDetails, suc
 				"backup": uuid,
 				"error":  err,
 			}).Error("failed to notify panel of backup status due to wings error")
-
 			return err
 		}
 
@@ -127,7 +126,7 @@ func (s *Server) RestoreBackup(b backup.BackupInterface, reader io.ReadCloser) (
 	defer func() {
 		s.Config().SetSuspended(false)
 		if reader != nil {
-			reader.Close()
+			_ = reader.Close()
 		}
 	}()
 	// Send an API call to the Panel as soon as this function is done running so that
@@ -142,7 +141,7 @@ func (s *Server) RestoreBackup(b backup.BackupInterface, reader io.ReadCloser) (
 	// instance, otherwise you'll likely hit all types of write errors due to the
 	// server being suspended.
 	if s.Environment.State() != environment.ProcessOfflineState {
-		if err = s.Environment.WaitForStop(s.Context(), time.Minute*2, false); err != nil {
+		if err = s.Environment.WaitForStop(s.Context(), 2*time.Minute, false); err != nil {
 			if !client.IsErrNotFound(err) {
 				return errors.WrapIf(err, "server/backup: restore: failed to wait for container stop")
 			}
@@ -152,14 +151,19 @@ func (s *Server) RestoreBackup(b backup.BackupInterface, reader io.ReadCloser) (
 	// Attempt to restore the backup to the server by running through each entry
 	// in the file one at a time and writing them to the disk.
 	s.Log().Debug("starting file writing process for backup restoration")
-	err = b.Restore(s.Context(), reader, func(file string, r io.Reader, mode fs.FileMode, atime, mtime time.Time) error {
+	err = b.Restore(s.Context(), reader, func(file string, info fs.FileInfo, r io.ReadCloser) error {
+		defer r.Close()
 		s.Events().Publish(DaemonMessageEvent, "(restoring): "+file)
+
 		if err := s.Filesystem().Writefile(file, r); err != nil {
 			return err
 		}
-		if err := s.Filesystem().Chmod(file, mode); err != nil {
+		if err := s.Filesystem().Chmod(file, info.Mode()); err != nil {
 			return err
 		}
+
+		atime := info.ModTime()
+		mtime := atime
 		return s.Filesystem().Chtimes(file, atime, mtime)
 	})
 

server/backup/backup.go

@@ -8,16 +8,21 @@ import (
 	"io/fs"
 	"os"
 	"path"
-	"time"
 
 	"emperror.dev/errors"
 	"github.com/apex/log"
+	"github.com/mholt/archiver/v4"
 	"golang.org/x/sync/errgroup"
 
 	"github.com/pterodactyl/wings/config"
 	"github.com/pterodactyl/wings/remote"
 )
 
+var format = archiver.CompressedArchive{
+	Compression: archiver.Gz{},
+	Archival:    archiver.Tar{},
+}
+
 type AdapterType string
 
 const (
@@ -27,7 +32,7 @@ const (
 
 // RestoreCallback is a generic restoration callback that exists for both local
 // and remote backups allowing the files to be restored.
-type RestoreCallback func(file string, r io.Reader, mode fs.FileMode, atime, mtime time.Time) error
+type RestoreCallback func(file string, info fs.FileInfo, r io.ReadCloser) error
 
 // noinspection GoNameStartsWithPackageName
 type BackupInterface interface {

server/backup/backup_local.go

@@ -6,8 +6,10 @@ import (
 	"os"
 
 	"emperror.dev/errors"
-	"github.com/mholt/archiver/v3"
+	"github.com/juju/ratelimit"
+	"github.com/mholt/archiver/v4"
 
+	"github.com/pterodactyl/wings/config"
 	"github.com/pterodactyl/wings/remote"
 	"github.com/pterodactyl/wings/server/filesystem"
 )
@@ -79,16 +81,27 @@ func (b *LocalBackup) Generate(ctx context.Context, basePath, ignore string) (*A
 // Restore will walk over the archive and call the callback function for each
 // file encountered.
 func (b *LocalBackup) Restore(ctx context.Context, _ io.Reader, callback RestoreCallback) error {
-	return archiver.Walk(b.Path(), func(f archiver.File) error {
-		select {
-		case <-ctx.Done():
-			// Stop walking if the context is canceled.
-			return archiver.ErrStopWalk
-		default:
-			if f.IsDir() {
-				return nil
-			}
-			return callback(filesystem.ExtractNameFromArchive(f), f, f.Mode(), f.ModTime(), f.ModTime())
+	f, err := os.Open(b.Path())
+	if err != nil {
+		return err
+	}
+
+	var reader io.Reader = f
+	// Steal the logic we use for making backups which will be applied when restoring
+	// this specific backup. This allows us to prevent overloading the disk unintentionally.
+	if writeLimit := int64(config.Get().System.Backups.WriteLimit * 1024 * 1024); writeLimit > 0 {
+		reader = ratelimit.Reader(f, ratelimit.NewBucketWithRate(float64(writeLimit), writeLimit))
+	}
+	if err := format.Extract(ctx, reader, nil, func(ctx context.Context, f archiver.File) error {
+		r, err := f.Open()
+		if err != nil {
+			return err
 		}
-	})
+		defer r.Close()
+
+		return callback(filesystem.ExtractNameFromArchive(f), f.FileInfo, r)
+	}); err != nil {
+		return err
+	}
+	return nil
 }

server/backup/backup_s3.go

@@ -1,8 +1,6 @@
 package backup
 
 import (
-	"archive/tar"
-	"compress/gzip"
 	"context"
 	"fmt"
 	"io"
@@ -13,13 +11,12 @@ import (
 
 	"emperror.dev/errors"
 	"github.com/cenkalti/backoff/v4"
-
-	"github.com/pterodactyl/wings/server/filesystem"
-
 	"github.com/juju/ratelimit"
+	"github.com/mholt/archiver/v4"
 
 	"github.com/pterodactyl/wings/config"
 	"github.com/pterodactyl/wings/remote"
+	"github.com/pterodactyl/wings/server/filesystem"
 )
 
 type S3Backup struct {
@@ -96,31 +93,16 @@ func (s *S3Backup) Restore(ctx context.Context, r io.Reader, callback RestoreCal
 	if writeLimit := int64(config.Get().System.Backups.WriteLimit * 1024 * 1024); writeLimit > 0 {
 		reader = ratelimit.Reader(r, ratelimit.NewBucketWithRate(float64(writeLimit), writeLimit))
 	}
-	gr, err := gzip.NewReader(reader)
-	if err != nil {
-		return err
-	}
-	defer gr.Close()
-	tr := tar.NewReader(gr)
-	for {
-		select {
-		case <-ctx.Done():
-			return nil
-		default:
-			// Do nothing, fall through to the next block of code in this loop.
-		}
-		header, err := tr.Next()
+	if err := format.Extract(ctx, reader, nil, func(ctx context.Context, f archiver.File) error {
+		r, err := f.Open()
 		if err != nil {
-			if err == io.EOF {
-				break
-			}
 			return err
 		}
-		if header.Typeflag == tar.TypeReg {
-			if err := callback(header.Name, tr, header.FileInfo().Mode(), header.AccessTime, header.ModTime); err != nil {
-				return err
-			}
-		}
+		defer r.Close()
+
+		return callback(filesystem.ExtractNameFromArchive(f), f.FileInfo, r)
+	}); err != nil {
+		return err
 	}
 	return nil
 }

server/filesystem/archive.go

@@ -30,6 +30,30 @@ var pool = sync.Pool{
 	},
 }
 
+// TarProgress .
+type TarProgress struct {
+	*tar.Writer
+	p *Progress
+}
+
+// NewTarProgress .
+func NewTarProgress(w *tar.Writer, p *Progress) *TarProgress {
+	if p != nil {
+		p.w = w
+	}
+	return &TarProgress{
+		Writer: w,
+		p:      p,
+	}
+}
+
+func (p *TarProgress) Write(v []byte) (int, error) {
+	if p.p == nil {
+		return p.Writer.Write(v)
+	}
+	return p.p.Write(v)
+}
+
 // Progress is used to track the progress of any I/O operation that are being
 // performed.
 type Progress struct {
@@ -46,6 +70,12 @@ 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 {
@@ -157,23 +187,17 @@ func (a *Archive) Create(dst string) error {
 	_ = gw.SetConcurrency(1<<20, 1)
 	defer gw.Close()
 
-	var pw io.Writer
-	if a.Progress != nil {
-		a.Progress.w = gw
-		pw = a.Progress
-	} else {
-		pw = gw
-	}
-
 	// Create a new tar writer around the gzip writer.
-	tw := tar.NewWriter(pw)
+	tw := tar.NewWriter(gw)
 	defer tw.Close()
 
+	pw := NewTarProgress(tw, a.Progress)
+
 	// Configure godirwalk.
 	options := &godirwalk.Options{
 		FollowSymbolicLinks: false,
 		Unsorted:            true,
-		Callback:            a.callback(tw),
+		Callback:            a.callback(pw),
 	}
 
 	// If we're specifically looking for only certain files, or have requested
@@ -182,7 +206,7 @@ func (a *Archive) Create(dst string) error {
 	if len(a.Files) == 0 && len(a.Ignore) > 0 {
 		i := ignore.CompileIgnoreLines(strings.Split(a.Ignore, "\n")...)
 
-		options.Callback = a.callback(tw, func(_ string, rp string) error {
+		options.Callback = a.callback(pw, func(_ string, rp string) error {
 			if i.MatchesPath(rp) {
 				return godirwalk.SkipThis
 			}
@@ -190,7 +214,7 @@ func (a *Archive) Create(dst string) error {
 			return nil
 		})
 	} else if len(a.Files) > 0 {
-		options.Callback = a.withFilesCallback(tw)
+		options.Callback = a.withFilesCallback(pw)
 	}
 
 	// Recursively walk the path we are archiving.
@@ -199,7 +223,7 @@ func (a *Archive) Create(dst string) error {
 
 // Callback function used to determine if a given file should be included in the archive
 // being generated.
-func (a *Archive) callback(tw *tar.Writer, opts ...func(path string, relative string) error) func(path string, de *godirwalk.Dirent) error {
+func (a *Archive) callback(tw *TarProgress, opts ...func(path string, relative string) error) func(path string, de *godirwalk.Dirent) error {
 	return func(path string, de *godirwalk.Dirent) error {
 		// Skip directories because we are walking them recursively.
 		if de.IsDir() {
@@ -223,7 +247,7 @@ func (a *Archive) callback(tw *tar.Writer, opts ...func(path string, relative st
 }
 
 // Pushes only files defined in the Files key to the final archive.
-func (a *Archive) withFilesCallback(tw *tar.Writer) func(path string, de *godirwalk.Dirent) error {
+func (a *Archive) withFilesCallback(tw *TarProgress) func(path string, de *godirwalk.Dirent) error {
 	return a.callback(tw, func(p string, rp string) error {
 		for _, f := range a.Files {
 			// If the given doesn't match, or doesn't have the same prefix continue
@@ -244,7 +268,7 @@ func (a *Archive) withFilesCallback(tw *tar.Writer) func(path string, de *godirw
 }
 
 // Adds a given file path to the final archive being created.
-func (a *Archive) addToArchive(p string, rp string, w *tar.Writer) error {
+func (a *Archive) addToArchive(p string, rp string, w *TarProgress) error {
 	// Lstat the file, this will give us the same information as Stat except that it will not
 	// follow a symlink to its target automatically. This is important to avoid including
 	// files that exist outside the server root unintentionally in the backup.

server/filesystem/compress.go

@@ -4,7 +4,9 @@ import (
 	"archive/tar"
 	"archive/zip"
 	"compress/gzip"
+	"context"
 	"fmt"
+	iofs "io/fs"
 	"os"
 	"path"
 	"path/filepath"
@@ -13,11 +15,10 @@ import (
 	"sync/atomic"
 	"time"
 
+	"emperror.dev/errors"
 	gzip2 "github.com/klauspost/compress/gzip"
 	zip2 "github.com/klauspost/compress/zip"
-
-	"emperror.dev/errors"
-	"github.com/mholt/archiver/v3"
+	"github.com/mholt/archiver/v4"
 )
 
 // CompressFiles compresses all of the files matching the given paths in the
@@ -73,7 +74,7 @@ func (fs *Filesystem) CompressFiles(dir string, paths []string) (os.FileInfo, er
 
 // 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(dir string, file string) error {
+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 {
@@ -89,69 +90,104 @@ func (fs *Filesystem) SpaceAvailableForDecompression(dir string, file string) er
 	// waiting an unnecessary amount of time on this call.
 	dirSize, err := fs.DiskUsage(false)
 
-	var size int64
-	// Walk over the archive and figure out just how large the final output would be from unarchiving it.
-	err = archiver.Walk(source, func(f archiver.File) error {
-		if atomic.AddInt64(&size, f.Size())+dirSize > fs.MaxDisk() {
-			return newFilesystemError(ErrCodeDiskSpace, nil)
-		}
-		return nil
-	})
+	fsys, err := archiver.FileSystem(source)
 	if err != nil {
-		if IsUnknownArchiveFormatError(err) {
+		if errors.Is(err, archiver.ErrNoMatch) {
 			return newFilesystemError(ErrCodeUnknownArchive, err)
 		}
 		return err
 	}
-	return err
+
+	var size 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 atomic.AddInt64(&size, info.Size())+dirSize > fs.MaxDisk() {
+				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 of the files within the given archive and ensure that there is not a
+// 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(dir string, file string) error {
+func (fs *Filesystem) DecompressFile(ctx context.Context, dir string, file string) error {
 	source, err := fs.SafePath(filepath.Join(dir, file))
 	if err != nil {
 		return err
 	}
-	// Ensure that the source archive actually exists on the system.
-	if _, err := os.Stat(source); err != nil {
+	return fs.DecompressFileUnsafe(ctx, dir, source)
+}
+
+// DecompressFileUnsafe will decompress any file on the local disk without checking
+// if it is owned by the server.  The file will be SAFELY decompressed and extracted
+// into the server's directory.
+func (fs *Filesystem) DecompressFileUnsafe(ctx context.Context, dir string, file string) error {
+	// Ensure that the archive actually exists on the system.
+	if _, err := os.Stat(file); err != nil {
 		return errors.WithStack(err)
 	}
 
-	// Walk all of the files in the archiver file and write them to the disk. If any
-	// directory is encountered it will be skipped since we handle creating any missing
-	// directories automatically when writing files.
-	err = archiver.Walk(source, func(f archiver.File) error {
-		if f.IsDir() {
-			return nil
-		}
-		p := filepath.Join(dir, 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
-		}
-		if err := fs.Writefile(p, f); err != nil {
-			return wrapError(err, source)
-		}
-		// Update the file permissions to the one set in the archive.
-		if err := fs.Chmod(p, f.Mode()); err != nil {
-			return wrapError(err, source)
-		}
-		// 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, source)
-		}
-		return nil
-	})
+	f, err := os.Open(file)
 	if err != nil {
-		if IsUnknownArchiveFormatError(err) {
+		return err
+	}
+
+	// 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
 	}
+
+	// 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 f.IsDir() {
+				return nil
+			}
+			p := filepath.Join(dir, 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
+			}
+			r, err := f.Open()
+			if err != nil {
+				return err
+			}
+			defer r.Close()
+			if err := fs.Writefile(p, r); err != nil {
+				return wrapError(err, file)
+			}
+			// Update the file permissions to the one set in the archive.
+			if err := fs.Chmod(p, f.Mode()); err != nil {
+				return wrapError(err, file)
+			}
+			// 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 nil
+		})
+	}
+
 	return nil
 }
 

server/filesystem/compress_test.go

@@ -1,6 +1,7 @@
 package filesystem
 
 import (
+	"context"
 	"os"
 	"sync/atomic"
 	"testing"
@@ -28,7 +29,7 @@ func TestFilesystem_DecompressFile(t *testing.T) {
 				g.Assert(err).IsNil()
 
 				// decompress
-				err = fs.DecompressFile("/", "test."+ext)
+				err = fs.DecompressFile(context.Background(), "/", "test."+ext)
 				g.Assert(err).IsNil()
 
 				// make sure everything is where it is supposed to be

server/filesystem/errors.go

@@ -4,7 +4,6 @@ import (
 	"fmt"
 	"os"
 	"path/filepath"
-	"strings"
 
 	"emperror.dev/errors"
 	"github.com/apex/log"
@@ -122,15 +121,6 @@ func IsErrorCode(err error, code ErrorCode) bool {
 	return false
 }
 
-// IsUnknownArchiveFormatError checks if the error is due to the archive being
-// in an unexpected file format.
-func IsUnknownArchiveFormatError(err error) bool {
-	if err != nil && strings.HasPrefix(err.Error(), "format ") {
-		return true
-	}
-	return false
-}
-
 // NewBadPathResolution returns a new BadPathResolution error.
 func NewBadPathResolution(path string, resolved string) error {
 	return errors.WithStackDepth(&Error{code: ErrCodePathResolution, path: path, resolved: resolved}, 1)