Warp/server/internal/storage/media_test.go

package storage_test

import (
	"bytes"
	"crypto/sha256"
	"database/sql"
	"encoding/hex"
	"errors"
	"io"
	"os"
	"strings"
	"testing"

	"github.com/dariuszjaniak/warp/server/internal/storage"
	"github.com/dariuszjaniak/warp/server/internal/storage/filestore"
)

// migratedMediaDB opens a driver with the media tables empty.
func migratedMediaDB(t *testing.T, driver string) storage.Database {
	t.Helper()
	db := openTestDB(t, driver)
	if err := db.Migrate(t.Context()); err != nil {
		t.Fatalf("Migrate: %v", err)
	}
	if driver == storage.DriverPostgres {
		pool, err := sql.Open("pgx", os.Getenv("WARP_TEST_POSTGRES_DSN"))
		if err != nil {
			t.Fatalf("open postgres to reset the media tables: %v", err)
		}
		defer func() { _ = pool.Close() }()
		if _, err := pool.ExecContext(t.Context(), `TRUNCATE media_repository, media_thumbnails`); err != nil {
			t.Fatalf("reset the media tables: %v", err)
		}
	}
	return db
}

func TestMediaUploadAndRead(t *testing.T) {
	for _, driver := range testDrivers {
		t.Run(driver, func(t *testing.T) {
			ctx := t.Context()
			media := migratedMediaDB(t, driver).Media()

			m := storage.Media{
				Origin: "warp.example", MediaID: "abc", ContentType: "image/png",
				UploadName: "cat.png", Size: 42, SHA256: strings.Repeat("a", 64),
				UserID: "@alice:warp.example", CreatedAtMS: 1000,
			}
			if err := media.CreateMedia(ctx, m); err != nil {
				t.Fatalf("CreateMedia: %v", err)
			}
			if err := media.CreateMedia(ctx, m); !errors.Is(err, storage.ErrMediaExists) {
				t.Fatalf("CreateMedia twice = %v, want ErrMediaExists", err)
			}
			got, err := media.Media(ctx, "warp.example", "abc")
			if err != nil {
				t.Fatalf("Media: %v", err)
			}
			if got != m {
				t.Errorf("Media = %+v, want %+v", got, m)
			}
			if _, err := media.Media(ctx, "other.example", "abc"); !errors.Is(err, storage.ErrNoSuchMedia) {
				t.Errorf("Media of another origin = %v, want ErrNoSuchMedia", err)
			}

			// Content, once stored, does not change.
			err = media.CompleteUpload(ctx, "warp.example", "abc", storage.MediaContent{SHA256: strings.Repeat("b", 64)})
			if !errors.Is(err, storage.ErrMediaUploaded) {
				t.Errorf("CompleteUpload of uploaded media = %v, want ErrMediaUploaded", err)
			}
			err = media.CompleteUpload(ctx, "warp.example", "nope", storage.MediaContent{})
			if !errors.Is(err, storage.ErrNoSuchMedia) {
				t.Errorf("CompleteUpload of unknown media = %v, want ErrNoSuchMedia", err)
			}
		})
	}
}

func TestMediaPendingUploads(t *testing.T) {
	for _, driver := range testDrivers {
		t.Run(driver, func(t *testing.T) {
			ctx := t.Context()
			media := migratedMediaDB(t, driver).Media()

			const alice = "@alice:warp.example"
			for _, p := range []struct {
				id      string
				expires int64
			}{{"p1", 5000}, {"p2", 5000}, {"old", 1500}} {
				if err := media.CreateMedia(ctx, storage.Media{
					Origin: "warp.example", MediaID: p.id, UserID: alice,
					CreatedAtMS: 1000, Pending: true, ExpiresAtMS: p.expires,
				}); err != nil {
					t.Fatalf("CreateMedia(%s): %v", p.id, err)
				}
			}

			n, err := media.PendingUploads(ctx, alice, 2000)
			if err != nil {
				t.Fatalf("PendingUploads: %v", err)
			}
			if n != 2 {
				t.Errorf("PendingUploads = %d, want 2 (one has expired)", n)
			}

			content := storage.MediaContent{
				ContentType: "text/plain", UploadName: "a.txt", Size: 3, SHA256: strings.Repeat("c", 64),
			}
			if err := media.CompleteUpload(ctx, "warp.example", "p1", content); err != nil {
				t.Fatalf("CompleteUpload: %v", err)
			}
			got, err := media.Media(ctx, "warp.example", "p1")
			if err != nil {
				t.Fatalf("Media: %v", err)
			}
			if got.Pending || got.ExpiresAtMS != 0 || got.SHA256 != content.SHA256 || got.Size != 3 {
				t.Errorf("completed media = %+v", got)
			}
			if err := media.CompleteUpload(ctx, "warp.example", "p1", content); !errors.Is(err, storage.ErrMediaUploaded) {
				t.Errorf("second CompleteUpload = %v, want ErrMediaUploaded", err)
			}

			deleted, err := media.DeleteExpiredUploads(ctx, 2000)
			if err != nil {
				t.Fatalf("DeleteExpiredUploads: %v", err)
			}
			if deleted != 1 {
				t.Errorf("DeleteExpiredUploads = %d, want 1", deleted)
			}
			if _, err := media.Media(ctx, "warp.example", "old"); !errors.Is(err, storage.ErrNoSuchMedia) {
				t.Errorf("expired upload still there: %v", err)
			}
			if _, err := media.Media(ctx, "warp.example", "p1"); err != nil {
				t.Errorf("a completed upload was swept: %v", err)
			}
		})
	}
}

func TestMediaThumbnails(t *testing.T) {
	for _, driver := range testDrivers {
		t.Run(driver, func(t *testing.T) {
			ctx := t.Context()
			media := migratedMediaDB(t, driver).Media()

			th := storage.Thumbnail{
				Origin: "remote.example", MediaID: "xyz", Width: 96, Height: 96, Method: "crop",
				ContentType: "image/jpeg", Size: 10, SHA256: strings.Repeat("d", 64), CreatedAtMS: 7,
			}
			if _, err := media.Thumbnail(ctx, "remote.example", "xyz", 96, 96, "crop"); !errors.Is(err, storage.ErrNoSuchMedia) {
				t.Fatalf("Thumbnail before storing = %v, want ErrNoSuchMedia", err)
			}
			if err := media.StoreThumbnail(ctx, th); err != nil {
				t.Fatalf("StoreThumbnail: %v", err)
			}
			th.SHA256, th.Size = strings.Repeat("e", 64), 11
			if err := media.StoreThumbnail(ctx, th); err != nil {
				t.Fatalf("StoreThumbnail again: %v", err)
			}
			got, err := media.Thumbnail(ctx, "remote.example", "xyz", 96, 96, "crop")
			if err != nil {
				t.Fatalf("Thumbnail: %v", err)
			}
			if got != th {
				t.Errorf("Thumbnail = %+v, want %+v", got, th)
			}
			if _, err := media.Thumbnail(ctx, "remote.example", "xyz", 96, 96, "scale"); !errors.Is(err, storage.ErrNoSuchMedia) {
				t.Errorf("Thumbnail of another method = %v, want ErrNoSuchMedia", err)
			}
		})
	}
}

func TestFileStore(t *testing.T) {
	ctx := t.Context()
	store, err := filestore.New(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}

	content := []byte("hello, media")
	sum := sha256.Sum256(content)
	info, err := store.Put(ctx, bytes.NewReader(content), 100)
	if err != nil {
		t.Fatalf("Put: %v", err)
	}
	if info.SHA256 != hex.EncodeToString(sum[:]) || info.Size != int64(len(content)) {
		t.Fatalf("Put = %+v", info)
	}
	// The same content again is the same blob.
	again, err := store.Put(ctx, bytes.NewReader(content), 100)
	if err != nil || again != info {
		t.Fatalf("Put twice = %+v, %v", again, err)
	}

	f, err := store.Open(ctx, info.SHA256)
	if err != nil {
		t.Fatalf("Open: %v", err)
	}
	got, _ := io.ReadAll(f)
	_ = f.Close()
	if !bytes.Equal(got, content) {
		t.Errorf("Open read %q", got)
	}

	if _, err := store.Put(ctx, bytes.NewReader(content), 5); !errors.Is(err, storage.ErrBlobTooLarge) {
		t.Errorf("Put past the limit = %v, want ErrBlobTooLarge", err)
	}
	// A hash that is not one cannot name a file, inside the store or out.
	for _, bad := range []string{"../../etc/passwd", strings.Repeat("G", 64), ""} {
		if _, err := store.Open(ctx, bad); !errors.Is(err, storage.ErrNoSuchBlob) {
			t.Errorf("Open(%q) = %v, want ErrNoSuchBlob", bad, err)
		}
	}

	if err := store.Delete(ctx, info.SHA256); err != nil {
		t.Fatalf("Delete: %v", err)
	}
	if _, err := store.Open(ctx, info.SHA256); !errors.Is(err, storage.ErrNoSuchBlob) {
		t.Errorf("Open after Delete = %v, want ErrNoSuchBlob", err)
	}
	if err := store.Delete(ctx, info.SHA256); err != nil {
		t.Errorf("Delete twice = %v", err)
	}
}