pterodactyl-wings/router/downloader/downloader_pull_test.go

package downloader

import (
	"bytes"
	"compress/gzip"
	"context"
	"encoding/json"
	"net"
	"net/http"
	"net/http/httptest"
	"net/url"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	"github.com/pterodactyl/wings/config"
	"github.com/pterodactyl/wings/remote"
	"github.com/pterodactyl/wings/server"
)

const pullTestServerUuid = "ffffffff-ffff-4fff-8fff-ffffffffffff"

// newPullTestServer returns a server whose files are stored in a temporary data
// directory. Loopback is allowed so that the tests can use a local origin.
func newPullTestServer(t *testing.T) *server.Server {
	t.Helper()
	previous := config.Get()
	t.Cleanup(func() { config.Set(previous) })
	next := *previous
	next.System.Data = t.TempDir()
	next.Api.RemoteDownloadAllowlist = []string{"127.0.0.1"}
	config.Set(&next)

	settings, err := json.Marshal(map[string]string{"uuid": pullTestServerUuid})
	require.NoError(t, err)
	s, err := server.NewEmptyManager(nil).InitServer(pullTestServerUuid, remote.ServerConfigurationResponse{
		Settings:             settings,
		ProcessConfiguration: &remote.ProcessConfiguration{},
	})
	require.NoError(t, err)
	require.NoError(t, os.MkdirAll(s.Filesystem().Path(), 0o755))
	return s
}

// A file from a server that compresses responses for clients that accept gzip
// can be pulled.
func TestDownloadAcceptsGzipEncodedResponses(t *testing.T) {
	s := newPullTestServer(t)

	payload := []byte(strings.Repeat("pull test payload\n", 4096))
	var compressed bytes.Buffer
	zw := gzip.NewWriter(&compressed)
	_, err := zw.Write(payload)
	require.NoError(t, err)
	require.NoError(t, zw.Close())

	origin := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		body := payload
		if strings.Contains(r.Header.Get("Accept-Encoding"), "gzip") {
			w.Header().Set("Content-Encoding", "gzip")
			body = compressed.Bytes()
		}
		w.Header().Set("Content-Type", "application/octet-stream")
		w.Header().Set("Content-Length", strconv.Itoa(len(body)))
		_, _ = w.Write(body)
	}))
	t.Cleanup(origin.Close)

	u, err := url.Parse(origin.URL + "/payload.bin")
	require.NoError(t, err)
	require.NoError(t, New(s, DownloadRequest{Directory: "/", URL: u, FileName: "payload.bin"}).Execute())

	got, err := os.ReadFile(filepath.Join(s.Filesystem().Path(), "payload.bin"))
	require.NoError(t, err)
	require.Equal(t, payload, got)
}

// Only the name of a file from Content-Disposition is used, so the file is saved
// in the directory that was asked for.
func TestDownloadFilenameFromHeaderIsSavedInRequestedDirectory(t *testing.T) {
	tests := []struct {
		directory string
		want      string
	}{
		{directory: "/chosen/dir", want: "chosen/dir/named.txt"},
		{directory: "chosen/dir", want: "chosen/dir/named.txt"},
		{directory: "", want: "named.txt"},
	}
	for _, tt := range tests {
		t.Run(tt.directory, func(t *testing.T) {
			s := newPullTestServer(t)

			origin := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				w.Header().Set("Content-Disposition", `attachment; filename="../../named.txt"`)
				w.Header().Set("Content-Length", "5")
				_, _ = w.Write([]byte("hello"))
			}))
			t.Cleanup(origin.Close)

			u, err := url.Parse(origin.URL + "/file")
			require.NoError(t, err)
			require.NoError(t, New(s, DownloadRequest{Directory: tt.directory, URL: u, FileName: "file", UseHeader: true}).Execute())

			root := s.Filesystem().Path()
			require.FileExists(t, filepath.Join(root, tt.want))
			require.NoFileExists(t, filepath.Join(filepath.Dir(root), "named.txt"))
		})
	}
}

// A download that fails to connect is reported as a failed download.
func TestDownloadConnectionFailureIsDownloadError(t *testing.T) {
	s := newPullTestServer(t)

	l, err := net.Listen("tcp", "127.0.0.1:0")
	require.NoError(t, err)
	addr := l.Addr().String()
	require.NoError(t, l.Close())

	u, err := url.Parse("http://" + addr + "/file")
	require.NoError(t, err)
	err = New(s, DownloadRequest{Directory: "/", URL: u, FileName: "file"}).Execute()
	require.Error(t, err)
	require.True(t, IsDownloadError(err), "got %v", err)
	require.ErrorIs(t, err, ErrDownloadFailed)
}

// newStalledOrigin returns the address of a file whose body is never sent, so a
// download of it only ends when it is cancelled. The number of requests made to
// it is counted.
func newStalledOrigin(t *testing.T) (*url.URL, *atomic.Int32) {
	t.Helper()
	var requests atomic.Int32
	release := make(chan struct{})
	origin := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		requests.Add(1)
		w.Header().Set("Content-Length", "1048576")
		w.WriteHeader(http.StatusOK)
		w.(http.Flusher).Flush()
		select {
		case <-r.Context().Done():
		case <-release:
		}
	}))
	t.Cleanup(origin.Close)
	t.Cleanup(func() { close(release) })

	u, err := url.Parse(origin.URL + "/stalled.bin")
	require.NoError(t, err)
	return u, &requests
}

// execute runs the download in the background and returns its outcome.
func execute(dl *Download) <-chan error {
	done := make(chan error, 1)
	go func() { done <- dl.Execute() }()
	return done
}

// waitForDownload fails the test if the download does not end shortly.
func waitForDownload(t *testing.T, done <-chan error) error {
	t.Helper()
	select {
	case err := <-done:
		return err
	case <-time.After(5 * time.Second):
		t.Fatal("download is still running after it was cancelled")
		return nil
	}
}

// A download cancelled before it starts never makes the request.
func TestDownloadCancelledBeforeExecuteIsNotStarted(t *testing.T) {
	s := newPullTestServer(t)
	u, requests := newStalledOrigin(t)

	dl := New(s, DownloadRequest{Directory: "/", URL: u, FileName: "stalled.bin"})
	dl.Cancel()
	require.Nil(t, ByID(dl.Identifier))

	err := waitForDownload(t, execute(dl))
	require.ErrorIs(t, err, context.Canceled)
	require.Zero(t, requests.Load())
	require.NoFileExists(t, filepath.Join(s.Filesystem().Path(), "stalled.bin"))
}

// A download cancelled while it is starting is stopped, wherever the
// cancellation lands.
func TestDownloadCancelledWhileExecuteStartsIsStopped(t *testing.T) {
	s := newPullTestServer(t)
	u, _ := newStalledOrigin(t)

	for range 20 {
		dl := New(s, DownloadRequest{Directory: "/", URL: u, FileName: "stalled.bin"})
		done := execute(dl)
		dl.Cancel()
		waitForDownload(t, done)
		require.Nil(t, ByID(dl.Identifier))
	}
}