Update base logic for handling server file downloads

This commit is contained in:
Dane Everitt 2020-12-20 10:59:07 -08:00
parent 3495fb1c76
commit 17daa2071f
No known key found for this signature in database
GPG Key ID: EEA66103B3D71F53
3 changed files with 147 additions and 20 deletions

View File

@ -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
}
}

View File

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

View File

@ -5,6 +5,7 @@ import (
"emperror.dev/errors" "emperror.dev/errors"
"github.com/apex/log" "github.com/apex/log"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"github.com/pterodactyl/wings/router/downloader"
"github.com/pterodactyl/wings/router/tokens" "github.com/pterodactyl/wings/router/tokens"
"github.com/pterodactyl/wings/server" "github.com/pterodactyl/wings/server"
"github.com/pterodactyl/wings/server/filesystem" "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. // 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) s := ExtractServer(c)
var data struct { var data struct {
URL string `binding:"required" json:"url"` URL string `binding:"required" json:"url"`
BasePath string `json:"path"` Directory string `binding:"required,omitempty" json:"directory"`
} }
if err := c.BindJSON(&data); err != nil { if err := c.BindJSON(&data); err != nil {
return return
@ -251,28 +252,21 @@ func postServerDownloadRemoteFile(c *gin.Context) {
}) })
return return
} }
NewServerError(err, s).Abort(c) WithError(c, err)
return return
} }
resp, err := http.Get(u.String()) dl := downloader.New(s, downloader.DownloadRequest{
if err != nil { URL: u,
NewServerError(err, s).Abort(c) Directory: data.Directory,
return })
}
defer resp.Body.Close()
filename := strings.Split(u.Path, "/") s.Log().WithField("download_id", dl.Identifier).WithField("url", u.String()).Info("starting pull of remote file to disk")
if err := s.Filesystem().Writefile(filepath.Join(data.BasePath, filename[len(filename)-1]), resp.Body); err != nil { if err := dl.Execute(); err != nil {
if filesystem.IsErrorCode(err, filesystem.ErrCodeIsDirectory) { WithError(c, err)
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)
return return
} }
s.Log().WithField("download_id", dl.Identifier).Info("completed pull of remote file")
c.Status(http.StatusNoContent) c.Status(http.StatusNoContent)
} }