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