From 0b9ad3de05b0ff41b6923603b9374ea0de876377 Mon Sep 17 00:00:00 2001
From: Dane Everitt <dane@daneeveritt.com>
Date: Sat, 26 Sep 2020 16:02:41 -0700
Subject: [PATCH] More file size confirmations before performing some actions;
 closes pterodactyl/panel#2420

---
 router/router_server_files.go | 41 +++++++++++++++++++++++++++++------
 server/filesystem.go          | 36 +++++++++++++++++++-----------
 2 files changed, 57 insertions(+), 20 deletions(-)

diff --git a/router/router_server_files.go b/router/router_server_files.go
index 8633fbd..bd9dd03 100644
--- a/router/router_server_files.go
+++ b/router/router_server_files.go
@@ -231,6 +231,12 @@ func postServerWriteFile(c *gin.Context) {
 	}
 	f = "/" + strings.TrimLeft(f, "/")
 
+	// Check if there is enough space available to perform this action.
+	if err := s.Filesystem.HasSpaceFor(c.Request.ContentLength); err != nil {
+		TrackedServerError(err, s).AbortFilesystemError(c)
+		return
+	}
+
 	if err := s.Filesystem.Writefile(f, c.Request.Body); err != nil {
 		if errors.Is(err, server.ErrIsDirectory) {
 			c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{
@@ -307,12 +313,30 @@ func postServerCompressFiles(c *gin.Context) {
 		return
 	}
 
+	d, err := s.Filesystem.SafePath(data.RootPath)
+	if err != nil {
+		TrackedServerError(err, s).AbortWithServerError(c)
+		return
+	}
+
 	f, err := s.Filesystem.CompressFiles(data.RootPath, data.Files)
 	if err != nil {
 		TrackedServerError(err, s).AbortFilesystemError(c)
 		return
 	}
 
+	if err := s.Filesystem.HasSpaceFor(f.Size()); err != nil {
+		if errors.Is(err, server.ErrNotEnoughDiskSpace) {
+			// Exceeding space limits, delete archive and return error. Kind of a waste of resources
+			// I suppose, but oh well.
+			_ = os.Remove(filepath.Join(d, f.Name()))
+		}
+
+		TrackedServerError(err, s).AbortFilesystemError(c)
+		return
+	}
+
+
 	c.JSON(http.StatusOK, &server.Stat{
 		Info:     f,
 		Mimetype: "application/tar+gzip",
@@ -396,13 +420,6 @@ func postServerUploadFiles(c *gin.Context) {
 		return
 	}
 
-	if !s.Filesystem.HasSpaceAvailable(true) {
-		c.AbortWithStatusJSON(http.StatusConflict, gin.H{
-			"error": "This server does not have enough available disk space to accept any file uploads.",
-		})
-		return
-	}
-
 	form, err := c.MultipartForm()
 	if err != nil {
 		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{
@@ -421,6 +438,16 @@ func postServerUploadFiles(c *gin.Context) {
 
 	directory := c.Query("directory")
 
+	var totalSize int64
+	for _, header := range headers {
+		totalSize += header.Size
+	}
+
+	if err := s.Filesystem.HasSpaceFor(totalSize); err != nil {
+		TrackedServerError(err, s).AbortFilesystemError(c)
+		return
+	}
+
 	for _, header := range headers {
 		p, err := s.Filesystem.SafePath(filepath.Join(directory, header.Filename))
 		if err != nil {
diff --git a/server/filesystem.go b/server/filesystem.go
index b0a843c..769401b 100644
--- a/server/filesystem.go
+++ b/server/filesystem.go
@@ -215,6 +215,26 @@ type SpaceCheckingOpts struct {
 	AllowStaleResponse bool
 }
 
+// Helper function to determine if a server has space available for a file of a given size.
+// If space is available, no error will be returned, otherwise an ErrNotEnoughSpace error
+// will be raised.
+func (fs *Filesystem) HasSpaceFor(size int64) error {
+	if fs.Server.DiskSpace() <= 0 {
+		return nil
+	}
+
+	s, err := fs.DiskUsage(true)
+	if err != nil {
+		return err
+	}
+
+	if (s + size) > fs.Server.DiskSpace() {
+		return ErrNotEnoughDiskSpace
+	}
+
+	return nil
+}
+
 // Determines if the directory a file is trying to be added to has enough space available
 // for the file to be written to.
 //
@@ -587,7 +607,7 @@ func (fs *Filesystem) Copy(p string) error {
 		return errors.WithStack(err)
 	}
 
-	s, err := os.Stat(cleaned);
+	s, err := os.Stat(cleaned)
 	if err != nil {
 		return errors.WithStack(err)
 	} else if s.IsDir() || !s.Mode().IsRegular() {
@@ -597,9 +617,8 @@ func (fs *Filesystem) Copy(p string) error {
 	}
 
 	// Check that copying this file wouldn't put the server over its limit.
-	curSize := atomic.LoadInt64(&fs.disk)
-	if (curSize + s.Size()) > fs.Server.DiskSpace() {
-		return ErrNotEnoughDiskSpace
+	if err := fs.HasSpaceFor(s.Size()); err != nil {
+		return err
 	}
 
 	base := filepath.Base(cleaned)
@@ -942,15 +961,6 @@ func (fs *Filesystem) CompressFiles(dir string, paths []string) (os.FileInfo, er
 		return nil, errors.WithStack(err)
 	}
 
-	curSize := atomic.LoadInt64(&fs.disk)
-	if (curSize + f.Size()) > fs.Server.DiskSpace() {
-		// Exceeding space limits, delete archive and return error. Kind of a waste of resources
-		// I suppose, but oh well.
-		_ = os.Remove(d)
-
-		return nil, ErrNotEnoughDiskSpace
-	}
-
 	atomic.AddInt64(&fs.disk, f.Size())
 
 	return f, nil