Update base logic for handling server file downloads

router/downloader/downloader.go

@@ -0,0 +1,133 @@
+package downloader
+
+import (
+	"context"
+	"emperror.dev/errors"
+	"github.com/google/uuid"
+	"github.com/pterodactyl/wings/server"
+	"net/http"
+	"net/url"
+	"path/filepath"
+	"strings"
+	"sync"
+	"time"
+)
+
+type Downloader struct {
+	mu            sync.RWMutex
+	downloadCache map[string]Download
+	serverCache   map[string][]string
+}
+
+type DownloadRequest struct {
+	URL       *url.URL
+	Directory string
+}
+
+type Download struct {
+	Identifier string
+	req        DownloadRequest
+	server     *server.Server
+	cancelFunc *context.CancelFunc
+}
+
+var client = &http.Client{Timeout: time.Hour * 12}
+var instance = &Downloader{
+	// Tracks all of the active downloads.
+	downloadCache: make(map[string]Download),
+	// Tracks all of the downloads active for a given server instance. This is
+	// primarily used to make things quicker and keep the code a little more
+	// legible throughout here.
+	serverCache: make(map[string][]string),
+}
+
+// Starts a new tracked download which allows for cancelation later on by calling
+// the Downloader.Cancel function.
+func New(s *server.Server, r DownloadRequest) *Download {
+	dl := Download{
+		Identifier: uuid.Must(uuid.NewRandom()).String(),
+		req:        r,
+		server:     s,
+	}
+	instance.track(dl)
+	return &dl
+}
+
+// Executes a given download for the server and begins writing the file to the disk. Once
+// completed the download will be removed from the cache.
+func (dl *Download) Execute() error {
+	ctx, cancel := context.WithTimeout(context.Background(), time.Hour*12)
+	dl.cancelFunc = &cancel
+	defer dl.Cancel()
+
+	fnameparts := strings.Split(dl.req.URL.Path, "/")
+	req, _ := http.NewRequestWithContext(ctx, http.MethodGet, dl.req.URL.String(), nil)
+	res, err := client.Do(req)
+	if err != nil {
+		return errors.New("downloader: failed opening request to download file")
+	}
+	defer res.Body.Close()
+	if res.StatusCode >= 300 || res.StatusCode < 200 {
+		return errors.New("downloader: got bad response status from endpoint: " + res.Status)
+	}
+	p := filepath.Join(dl.req.Directory, fnameparts[len(fnameparts)-1])
+	dl.server.Log().WithField("path", p).Debug("writing remote file to disk")
+	if err := dl.server.Filesystem().Writefile(p, res.Body); err != nil {
+		return errors.WrapIf(err, "downloader: failed to write file to server directory")
+	}
+	return nil
+}
+
+// Cancels a running download and frees up the associated resources. If a file is being
+// written a partial file will remain present on the disk.
+func (dl *Download) Cancel() {
+	if dl.cancelFunc != nil {
+		(*dl.cancelFunc)()
+	}
+	instance.remove(dl.Identifier)
+}
+
+// Tracks a download in the internal cache for this instance.
+func (d *Downloader) track(dl Download) {
+	d.mu.Lock()
+	defer d.mu.Unlock()
+	sid := dl.server.Id()
+	if _, ok := d.downloadCache[dl.Identifier]; !ok {
+		d.downloadCache[dl.Identifier] = dl
+		if _, ok := d.serverCache[sid]; !ok {
+			d.serverCache[sid] = []string{}
+		}
+		d.serverCache[sid] = append(d.serverCache[sid], dl.Identifier)
+	}
+}
+
+// Finds a given download entry using the provided ID and returns it.
+func (d *Downloader) find(dlid string) (Download, bool) {
+	d.mu.RLock()
+	defer d.mu.RUnlock()
+	if entry, ok := d.downloadCache[dlid]; ok {
+		return entry, true
+	}
+	return Download{}, false
+}
+
+// Remove the given download reference from the cache storing them. This also updates
+// the slice of active downloads for a given server to not include this download.
+func (d *Downloader) remove(dlid string) {
+	d.mu.Lock()
+	defer d.mu.Unlock()
+	if _, ok := d.downloadCache[dlid]; !ok {
+		return
+	}
+	sid := d.downloadCache[dlid].server.Id()
+	delete(d.downloadCache, dlid)
+	if tracked, ok := d.serverCache[sid]; ok {
+		var out []string
+		for _, k := range tracked {
+			if k != dlid {
+				out = append(out, k)
+			}
+		}
+		d.serverCache[sid] = out
+	}
+}

router/router.go

@@ -82,7 +82,7 @@ func Configure() *gin.Engine {
 			files.PUT("/rename", putServerRenameFiles)
 			files.POST("/copy", postServerCopyFile)
 			files.POST("/write", postServerWriteFile)
-			files.POST("/writeUrl", postServerDownloadRemoteFile)
+			files.POST("/pull", postServerPullRemoteFile)
 			files.POST("/create-directory", postServerCreateDirectory)
 			files.POST("/delete", postServerDeleteFiles)
 			files.POST("/compress", postServerCompressFiles)

router/router_server_files.go

@@ -5,6 +5,7 @@ import (
 	"emperror.dev/errors"
 	"github.com/apex/log"
 	"github.com/gin-gonic/gin"
+	"github.com/pterodactyl/wings/router/downloader"
 	"github.com/pterodactyl/wings/router/tokens"
 	"github.com/pterodactyl/wings/server"
 	"github.com/pterodactyl/wings/server/filesystem"
@@ -233,11 +234,11 @@ func postServerWriteFile(c *gin.Context) {
 }
 
 // Writes the contents of the remote URL to a file on a server.
-func postServerDownloadRemoteFile(c *gin.Context) {
+func postServerPullRemoteFile(c *gin.Context) {
 	s := ExtractServer(c)
 	var data struct {
-		URL      string `binding:"required" json:"url"`
-		BasePath string `json:"path"`
+		URL       string `binding:"required" json:"url"`
+		Directory string `binding:"required,omitempty" json:"directory"`
 	}
 	if err := c.BindJSON(&data); err != nil {
 		return
@@ -251,28 +252,21 @@ func postServerDownloadRemoteFile(c *gin.Context) {
 			})
 			return
 		}
-		NewServerError(err, s).Abort(c)
+		WithError(c, err)
 		return
 	}
 
-	resp, err := http.Get(u.String())
-	if err != nil {
-		NewServerError(err, s).Abort(c)
-		return
-	}
-	defer resp.Body.Close()
+	dl := downloader.New(s, downloader.DownloadRequest{
+		URL:       u,
+		Directory: data.Directory,
+	})
 
-	filename := strings.Split(u.Path, "/")
-	if err := s.Filesystem().Writefile(filepath.Join(data.BasePath, filename[len(filename)-1]), resp.Body); err != nil {
-		if filesystem.IsErrorCode(err, filesystem.ErrCodeIsDirectory) {
-			c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{
-				"error": "Cannot write file, name conflicts with an existing directory by the same name.",
-			})
-			return
-		}
-		NewServerError(err, s).AbortFilesystemError(c)
+	s.Log().WithField("download_id", dl.Identifier).WithField("url", u.String()).Info("starting pull of remote file to disk")
+	if err := dl.Execute(); err != nil {
+		WithError(c, err)
 		return
 	}
+	s.Log().WithField("download_id", dl.Identifier).Info("completed pull of remote file")
 
 	c.Status(http.StatusNoContent)
 }